README.md

---
title: OrbGen Training
emoji: 🚀
colorFrom: blue
colorTo: purple
sdk: docker
pinned: false
license: apache-2.0
---

# OrbGen Training

Training code for OrbGen - a model that generates valid Orbital schemas from natural language.

## Quick Start

### Local Training

```bash
# Install dependencies
pip install -r requirements.txt

# Full training (with local GPU)
python train.py --config config.yaml

# Debug mode (1 epoch, no W&B)
python train.py --config config.yaml --debug
```

### HuggingFace Cloud Training

```bash
# Use the HuggingFace-optimized config
python train.py --config config-huggingface.yaml
```

## Configuration Files

| Config | GPU | VRAM | Use Case |
|--------|-----|------|----------|
| `config.yaml` | Local (RTX 3000) | 6GB | Local testing with QLoRA |
| `config-huggingface.yaml` | A10G | 24GB | HuggingFace Spaces training |

## Evaluation

```bash
# Basic evaluation
python evaluate.py --checkpoint ./orbgen-1.5b/final

# With Orbital validator
python evaluate.py --checkpoint ./orbgen-1.5b/final --use_validator
```

## Generate

```bash
# Single generation
python generate.py --prompt "Create a task management app with projects and due dates"

# Interactive mode
python generate.py --interactive

# Save to file
python generate.py --prompt "..." --output schema.orb --validate
```

## Files

| File | Description |
|------|-------------|
| `train.py` | Main training script with SFT |
| `evaluate.py` | Evaluation with Orbital validation |
| `generate.py` | Inference and generation |
| `config.yaml` | Config for local 6GB GPU (QLoRA) |
| `config-huggingface.yaml` | Config for HuggingFace A10G (24GB) |
| `Dockerfile` | Container for HuggingFace Spaces |
| `requirements.txt` | Python dependencies |

## Training on HuggingFace Spaces

### Prerequisites

1. Upload dataset to HuggingFace:
   ```bash
   python scripts/upload_dataset.py
   ```

2. Set Space secrets:
   - `HUGGINGFACE_TOKEN` - HF token with write access
   - `WANDB_API_KEY` - Weights & Biases API key

3. Push training code:
   ```bash
   cd orbgen-training
   huggingface-cli upload orbital-ai/orbgen-training . --repo-type space
   ```

4. Configure Space with A10G GPU in settings

5. Training will start automatically

## Hardware Requirements

| Phase | GPU | VRAM | Notes |
|-------|-----|------|-------|
| Training (local) | RTX 3000 | 6GB | Uses QLoRA (4-bit) |
| Training (cloud) | A10G | 24GB | Full bf16 training |
| Inference | T4 | 16GB | Production inference |

## Model Output

After training, the model is saved to `./orbgen-1.5b/final/`:
- `adapter_config.json` - LoRA configuration
- `adapter_model.safetensors` - LoRA weights
- `tokenizer.json` - Tokenizer
- `config.json` - Model config

Model is automatically pushed to `orbital-ai/orbgen-1.5b` when using HuggingFace config.