Metadata-Version: 2.4
Name: tinyaicli
Version: 0.1.0
Summary: A CLI tool for training small LLMs and vision models with reproducible research
Home-page: https://github.com/nheinstein/tinyai
Author: Nathan Heinstein
Author-email: nathan.heinstein@dal.ca
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: torchvision>=0.15.0
Requires-Dist: torchaudio>=2.0.0
Requires-Dist: hydra-core>=1.3.0
Requires-Dist: hydra-colorlog>=1.2.0
Requires-Dist: wandb>=0.15.0
Requires-Dist: omegaconf>=2.3.0
Requires-Dist: datasets>=2.14.0
Requires-Dist: transformers>=4.30.0
Requires-Dist: tokenizers>=0.13.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: pandas>=2.0.0
Requires-Dist: matplotlib>=3.7.0
Requires-Dist: seaborn>=0.12.0
Requires-Dist: tqdm>=4.65.0
Requires-Dist: pyyaml>=6.0
Requires-Dist: click>=8.1.0
Requires-Dist: rich>=13.0.0
Requires-Dist: tensorboard>=2.13.0
Requires-Dist: pillow>=10.0.0
Requires-Dist: scikit-learn>=1.3.0
Requires-Dist: accelerate>=0.20.0
Requires-Dist: peft>=0.4.0
Requires-Dist: bitsandbytes>=0.41.0
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: license-file
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# Tiny AI

A powerful CLI tool for training small Language Models (LLMs) and Vision models using custom YAML configurations for reproducible research.

## Features

-  **Multi-Model Support**: Train both LLMs and Vision models
-  **Hydra Configuration**: YAML-based configs for reproducible experiments
-  **Experiment Tracking**: Integrated wandb logging and TensorBoard support
-  **Modular Architecture**: Easy to extend and customize
-  **Rich Logging**: Beautiful CLI output with progress tracking
-  **Reproducible Research**: Deterministic training with seed management


### Installation

```bash
# Clone the repository
git clone https://github.com/nheinstein/TinyAI.git
cd TinyAI

# Install dependencies
pip install -r requirements.txt

# Install the package in development mode
pip install -e .
```

### Basic Usage

```bash
# Train an LLM with default config
python -m tinyai.train model=llm

# Train a vision model with custom config
python -m tinyai.train model=vision config=vision_custom

# Override config parameters
python -m tinyai.train model=llm training.batch_size=32 training.learning_rate=1e-4

# Run with wandb logging
python -m tinyai.train model=llm logging.wandb=true logging.project_name=my_experiment
```

## Configuration

The trainer uses Hydra for configuration management. Configs are stored in `configs/` directory:

```
configs/
├── model/
│   ├── llm.yaml
│   └── vision.yaml
├── training/
│   ├── default.yaml
│   └── custom.yaml
└── logging/
    └── default.yaml
```

### Example Config Structure

```yaml
# configs/model/llm.yaml
defaults:
  - training: default
  - logging: default

model:
  type: "transformer"
  vocab_size: 50257
  hidden_size: 768
  num_layers: 12
  num_heads: 12

training:
  batch_size: 16
  learning_rate: 1e-4
  num_epochs: 10
  warmup_steps: 1000

data:
  train_path: "data/train.txt"
  val_path: "data/val.txt"
  max_length: 512

logging:
  wandb: true
  project_name: "tiny-llm"
  log_interval: 100
```

## Project Structure

```
TinyAI/
├── tinyai/
│   ├── __init__.py
│   ├── train.py              # Main training script
│   ├── models/               # Model implementations
│   │   ├── __init__.py
│   │   ├── llm.py
│   │   └── vision.py
│   ├── data/                 # Data loading utilities
│   │   ├── __init__.py
│   │   ├── datasets.py
│   │   └── tokenizers.py
│   ├── training/             # Training utilities
│   │   ├── __init__.py
│   │   ├── trainer.py
│   │   └── optimizer.py
│   └── utils/                # Utility functions
│       ├── __init__.py
│       ├── logging.py
│       └── metrics.py
├── configs/                  # Hydra configurations
├── scripts/                  # Utility scripts
├── requirements.txt
├── setup.py
└── README.md
```

## Advanced Usage

### Custom Model Implementation

```python
# tinyai/models/custom_model.py
from tinyai.models.base import BaseModel

class CustomModel(BaseModel):
    def __init__(self, config):
        super().__init__(config)
        # Your model implementation

    def forward(self, x):
        # Forward pass
        return output
```

### Custom Training Loop

```python
# Override training behavior
class CustomTrainer(Trainer):
    def training_step(self, batch):
        # Custom training logic
        pass
```

## License

MIT License - see LICENSE file for details.
