Upload folder using huggingface_hub

Dockerfile

@@ -0,0 +1,23 @@
+FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements and install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy training code
+COPY . .
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1
+ENV HF_HOME=/app/.cache/huggingface
+
+# Default command
+CMD ["python", "train.py", "--config", "config.yaml"]

README.md

@@ -1,10 +1,98 @@
----
-title: Orbgen Training
-emoji: 🚀
-colorFrom: yellow
-colorTo: red
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# OrbGen Training
+
+Training code for OrbGen - a model that generates valid Orbital schemas from natural language.
+
+## Quick Start
+
+### 1. Install Dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
+### 2. Configure Training
+
+Edit `config.yaml` to adjust:
+- Base model
+- Training hyperparameters
+- LoRA configuration
+- W&B settings
+
+### 3. Train
+
+```bash
+# Full training
+python train.py --config config.yaml
+
+# Debug mode (1 epoch, no W&B)
+python train.py --config config.yaml --debug
+
+# Resume from checkpoint
+python train.py --config config.yaml --resume_from_checkpoint ./orbgen-1.5b/checkpoint-500
+```
+
+### 4. Evaluate
+
+```bash
+# Basic evaluation
+python evaluate.py --checkpoint ./orbgen-1.5b/final
+
+# With Orbital validator
+python evaluate.py --checkpoint ./orbgen-1.5b/final --use_validator
+```
+
+### 5. 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` | Training configuration |
+| `Dockerfile` | Container for training |
+| `requirements.txt` | Python dependencies |
+
+## Training on HuggingFace Spaces
+
+1. Push this directory to `orbital-ai/orbgen-training`:
+   ```bash
+   hf upload orbital-ai/orbgen-training . --repo-type space
+   ```
+
+2. Configure the Space with A10G GPU
+
+3. Training will start automatically
+
+## Hardware Requirements
+
+| Phase | GPU | VRAM | Time |
+|-------|-----|------|------|
+| Training (3 epochs) | A10G | 24GB | ~6 hours |
+| Evaluation | T4 | 16GB | ~30 min |
+| Inference | T4 | 16GB | ~1 sec/gen |
+
+## 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
+
+To push to HuggingFace:
+```bash
+hf upload orbital-ai/orbgen-1.5b ./orbgen-1.5b/final --repo-type model
+```

config.yaml

@@ -0,0 +1,55 @@
+# OrbGen Training Configuration
+
+model:
+  base_model: "Qwen/Qwen2.5-Coder-1.5B"
+  output_dir: "./orbgen-1.5b"
+  max_seq_length: 8192
+
+data:
+  dataset: "orbital-ai/orbital-schemas"
+  train_split: "train"
+  eval_split: "validation"
+
+training:
+  # SFT Configuration
+  num_epochs: 3
+  per_device_train_batch_size: 2
+  per_device_eval_batch_size: 2
+  gradient_accumulation_steps: 8
+  learning_rate: 2.0e-5
+  warmup_ratio: 0.1
+  weight_decay: 0.01
+  max_grad_norm: 1.0
+
+  # Logging
+  logging_steps: 10
+  eval_steps: 50
+  save_steps: 100
+  save_total_limit: 3
+
+lora:
+  enabled: true
+  r: 64
+  lora_alpha: 128
+  lora_dropout: 0.05
+  target_modules:
+    - "q_proj"
+    - "k_proj"
+    - "v_proj"
+    - "o_proj"
+    - "gate_proj"
+    - "up_proj"
+    - "down_proj"
+  bias: "none"
+  task_type: "CAUSAL_LM"
+
+generation:
+  max_new_tokens: 4096
+  temperature: 0.7
+  top_p: 0.95
+  do_sample: true
+
+wandb:
+  project: "orbgen-training"
+  entity: "orbital-ai"
+  run_name: "orbgen-1.5b-sft"

evaluate.py

@@ -0,0 +1,254 @@
+#!/usr/bin/env python3
+"""
+OrbGen Evaluation Script
+
+Evaluates a trained model on the test set with Orbital validation metrics.
+
+Usage:
+    python evaluate.py --checkpoint ./orbgen-1.5b/final
+    python evaluate.py --checkpoint ./orbgen-1.5b/final --use_validator
+"""
+
+import os
+import json
+import fire
+import torch
+import subprocess
+import tempfile
+from pathlib import Path
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+from tqdm import tqdm
+
+
+def validate_schema(schema_json: str) -> tuple[bool, list[str]]:
+    """Validate schema using orbital CLI."""
+    # Check valid JSON first
+    try:
+        json.loads(schema_json)
+    except json.JSONDecodeError as e:
+        return False, [f"Invalid JSON: {e}"]
+
+    # Write to temp file and validate
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.orb', delete=False) as f:
+        f.write(schema_json)
+        temp_path = f.name
+
+    try:
+        result = subprocess.run(
+            ['orbital', 'validate', temp_path],
+            capture_output=True,
+            text=True,
+            timeout=30,
+            cwd=os.path.expanduser('~/kflow.ai.builder/builder')
+        )
+
+        if result.returncode == 0 or 'Schema is valid' in result.stdout:
+            return True, []
+        else:
+            errors = [line for line in result.stderr.split('\n') if line.strip()]
+            return False, errors[:5]
+    except subprocess.TimeoutExpired:
+        return False, ["Validation timeout"]
+    except FileNotFoundError:
+        return False, ["Orbital CLI not found"]
+    except Exception as e:
+        return False, [f"Validation error: {e}"]
+    finally:
+        Path(temp_path).unlink(missing_ok=True)
+
+
+def extract_completion(generated_text: str) -> str:
+    """Extract the completion from generated text."""
+    # Try to find assistant response
+    if '<|im_start|>assistant' in generated_text:
+        parts = generated_text.split('<|im_start|>assistant')
+        if len(parts) > 1:
+            completion = parts[-1]
+            if '<|im_end|>' in completion:
+                completion = completion.split('<|im_end|>')[0]
+            return completion.strip()
+
+    # Try to find JSON object
+    start = generated_text.find('{')
+    if start != -1:
+        # Find matching closing brace
+        depth = 0
+        for i, char in enumerate(generated_text[start:]):
+            if char == '{':
+                depth += 1
+            elif char == '}':
+                depth -= 1
+                if depth == 0:
+                    return generated_text[start:start + i + 1]
+
+    return generated_text
+
+
+def main(
+    checkpoint: str = "./orbgen-1.5b/final",
+    dataset: str = "orbital-ai/orbital-schemas",
+    split: str = "test",
+    use_validator: bool = False,
+    max_samples: int = -1,
+    output_file: str = "evaluation_results.json",
+):
+    """Evaluate model on test set."""
+
+    print("=" * 60)
+    print("OrbGen Evaluation")
+    print("=" * 60)
+    print(f"Checkpoint: {checkpoint}")
+    print(f"Dataset: {dataset}")
+    print(f"Split: {split}")
+    print(f"Use Validator: {use_validator}")
+    print("=" * 60)
+
+    # Load tokenizer and model
+    print("\nLoading model...")
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        checkpoint,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    model.eval()
+
+    # Load dataset
+    print("Loading dataset...")
+    ds = load_dataset(dataset)
+    test_data = ds[split]
+
+    if max_samples > 0:
+        test_data = test_data.select(range(min(max_samples, len(test_data))))
+
+    print(f"Evaluating on {len(test_data)} examples...")
+
+    # Metrics
+    metrics = {
+        'total': len(test_data),
+        'valid_json': 0,
+        'valid_schema': 0,
+        'generation_errors': 0,
+    }
+    results = []
+
+    system_prompt = """You are OrbGen, a specialized AI that generates valid Orbital schemas (.orb files) from natural language descriptions.
+
+Rules:
+1. Output ONLY valid JSON - no explanations, no markdown code blocks
+2. Every schema must have: name, version, orbitals array
+3. Each orbital must have: name, entity, traits, pages
+4. Each entity must have: name, collection (or runtime/singleton), fields
+5. Each trait must have: name, category (interaction/integration), linkedEntity, stateMachine
+6. State machines must have: states (with one isInitial:true), events, transitions
+7. Use S-expression arrays for effects: ["set", "field", "value"], ["emit", "EVENT", {}], ["render-ui", "slot", {...}]
+8. Pages must have: name, path, entity, traits"""
+
+    for i, example in enumerate(tqdm(test_data)):
+        prompt = example['prompt']
+        expected = example['completion']
+
+        # Format input
+        input_text = f"""<|im_start|>system
+{system_prompt}
+<|im_end|>
+<|im_start|>user
+{prompt}
+<|im_end|>
+<|im_start|>assistant
+"""
+
+        try:
+            # Generate
+            inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=4096,
+                    temperature=0.7,
+                    top_p=0.95,
+                    do_sample=True,
+                    pad_token_id=tokenizer.eos_token_id,
+                )
+
+            generated = tokenizer.decode(outputs[0], skip_special_tokens=False)
+            completion = extract_completion(generated)
+
+            # Check valid JSON
+            is_valid_json = False
+            is_valid_schema = False
+            errors = []
+
+            try:
+                json.loads(completion)
+                is_valid_json = True
+                metrics['valid_json'] += 1
+
+                # Check valid schema
+                if use_validator:
+                    is_valid_schema, errors = validate_schema(completion)
+                    if is_valid_schema:
+                        metrics['valid_schema'] += 1
+                else:
+                    # Basic structural check
+                    parsed = json.loads(completion)
+                    if 'name' in parsed and 'orbitals' in parsed:
+                        is_valid_schema = True
+                        metrics['valid_schema'] += 1
+
+            except json.JSONDecodeError as e:
+                errors = [f"JSON error: {e}"]
+
+            results.append({
+                'prompt': prompt,
+                'expected': expected[:500] + '...' if len(expected) > 500 else expected,
+                'generated': completion[:500] + '...' if len(completion) > 500 else completion,
+                'valid_json': is_valid_json,
+                'valid_schema': is_valid_schema,
+                'errors': errors,
+            })
+
+        except Exception as e:
+            metrics['generation_errors'] += 1
+            results.append({
+                'prompt': prompt,
+                'error': str(e),
+                'valid_json': False,
+                'valid_schema': False,
+            })
+
+    # Calculate percentages
+    metrics['valid_json_pct'] = metrics['valid_json'] / metrics['total'] * 100
+    metrics['valid_schema_pct'] = metrics['valid_schema'] / metrics['total'] * 100
+
+    # Print results
+    print("\n" + "=" * 60)
+    print("Results")
+    print("=" * 60)
+    print(f"Total examples: {metrics['total']}")
+    print(f"Valid JSON: {metrics['valid_json']} ({metrics['valid_json_pct']:.1f}%)")
+    print(f"Valid Schema: {metrics['valid_schema']} ({metrics['valid_schema_pct']:.1f}%)")
+    print(f"Generation errors: {metrics['generation_errors']}")
+
+    # Save results
+    output = {
+        'metrics': metrics,
+        'results': results,
+    }
+
+    with open(output_file, 'w') as f:
+        json.dump(output, f, indent=2)
+
+    print(f"\nResults saved to: {output_file}")
+
+    return metrics
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

generate.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python3
+"""
+OrbGen Inference Script
+
+Generate Orbital schemas from natural language prompts.
+
+Usage:
+    python generate.py --prompt "Create a task management app"
+    python generate.py --prompt "..." --checkpoint ./orbgen-1.5b/final
+    python generate.py --interactive
+"""
+
+import os
+import json
+import fire
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+SYSTEM_PROMPT = """You are OrbGen, a specialized AI that generates valid Orbital schemas (.orb files) from natural language descriptions.
+
+Rules:
+1. Output ONLY valid JSON - no explanations, no markdown code blocks
+2. Every schema must have: name, version, orbitals array
+3. Each orbital must have: name, entity, traits, pages
+4. Each entity must have: name, collection (or runtime/singleton), fields
+5. Each trait must have: name, category (interaction/integration), linkedEntity, stateMachine
+6. State machines must have: states (with one isInitial:true), events, transitions
+7. Use S-expression arrays for effects: ["set", "field", "value"], ["emit", "EVENT", {}], ["render-ui", "slot", {...}]
+8. Pages must have: name, path, entity, traits"""
+
+
+class OrbGen:
+    """OrbGen schema generator."""
+
+    def __init__(
+        self,
+        checkpoint: str = "orbital-ai/orbgen-1.5b",
+        device: str = "auto",
+    ):
+        """Initialize the generator."""
+        print(f"Loading model from {checkpoint}...")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            checkpoint,
+            trust_remote_code=True,
+        )
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            checkpoint,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+            trust_remote_code=True,
+        )
+        self.model.eval()
+
+        print("Model loaded!")
+
+    def generate(
+        self,
+        prompt: str,
+        temperature: float = 0.7,
+        top_p: float = 0.95,
+        max_new_tokens: int = 4096,
+    ) -> str:
+        """Generate a schema from a prompt."""
+
+        input_text = f"""<|im_start|>system
+{SYSTEM_PROMPT}
+<|im_end|>
+<|im_start|>user
+{prompt}
+<|im_end|>
+<|im_start|>assistant
+"""
+
+        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id,
+            )
+
+        generated = self.tokenizer.decode(outputs[0], skip_special_tokens=False)
+
+        # Extract completion
+        if '<|im_start|>assistant' in generated:
+            parts = generated.split('<|im_start|>assistant')
+            if len(parts) > 1:
+                completion = parts[-1]
+                if '<|im_end|>' in completion:
+                    completion = completion.split('<|im_end|>')[0]
+                return completion.strip()
+
+        # Try to find JSON
+        start = generated.find('{')
+        if start != -1:
+            depth = 0
+            for i, char in enumerate(generated[start:]):
+                if char == '{':
+                    depth += 1
+                elif char == '}':
+                    depth -= 1
+                    if depth == 0:
+                        return generated[start:start + i + 1]
+
+        return generated
+
+
+def main(
+    prompt: str = None,
+    checkpoint: str = "orbital-ai/orbgen-1.5b",
+    output: str = None,
+    temperature: float = 0.7,
+    top_p: float = 0.95,
+    interactive: bool = False,
+    validate: bool = False,
+):
+    """Generate Orbital schemas."""
+
+    generator = OrbGen(checkpoint=checkpoint)
+
+    if interactive:
+        print("\n" + "=" * 60)
+        print("OrbGen Interactive Mode")
+        print("=" * 60)
+        print("Enter prompts to generate schemas. Type 'quit' to exit.\n")
+
+        while True:
+            try:
+                prompt = input("Prompt> ").strip()
+                if prompt.lower() in ['quit', 'exit', 'q']:
+                    break
+                if not prompt:
+                    continue
+
+                print("\nGenerating...")
+                result = generator.generate(prompt, temperature=temperature, top_p=top_p)
+
+                try:
+                    parsed = json.loads(result)
+                    print(json.dumps(parsed, indent=2))
+                except json.JSONDecodeError:
+                    print(result)
+
+                print()
+
+            except KeyboardInterrupt:
+                print("\nExiting...")
+                break
+
+    elif prompt:
+        print(f"\nPrompt: {prompt}\n")
+        print("Generating...")
+
+        result = generator.generate(prompt, temperature=temperature, top_p=top_p)
+
+        try:
+            parsed = json.loads(result)
+            formatted = json.dumps(parsed, indent=2)
+
+            if output:
+                with open(output, 'w') as f:
+                    f.write(formatted)
+                print(f"Schema saved to: {output}")
+            else:
+                print(formatted)
+
+            # Validate if requested
+            if validate:
+                import subprocess
+                import tempfile
+                from pathlib import Path
+
+                with tempfile.NamedTemporaryFile(mode='w', suffix='.orb', delete=False) as f:
+                    f.write(formatted)
+                    temp_path = f.name
+
+                try:
+                    result = subprocess.run(
+                        ['orbital', 'validate', temp_path],
+                        capture_output=True,
+                        text=True,
+                        cwd=os.path.expanduser('~/kflow.ai.builder/builder')
+                    )
+                    print("\nValidation:")
+                    print(result.stdout or result.stderr)
+                finally:
+                    Path(temp_path).unlink(missing_ok=True)
+
+        except json.JSONDecodeError as e:
+            print(f"Warning: Generated invalid JSON: {e}")
+            print(result)
+
+    else:
+        print("Usage: python generate.py --prompt 'Your prompt here'")
+        print("       python generate.py --interactive")
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

requirements.txt

@@ -0,0 +1,15 @@
+# OrbGen Training Dependencies
+torch>=2.1.0
+transformers>=4.36.0
+datasets>=2.16.0
+peft>=0.7.0
+trl>=0.7.0
+accelerate>=0.25.0
+bitsandbytes>=0.41.0
+wandb>=0.16.0
+huggingface_hub>=0.20.0
+safetensors>=0.4.0
+sentencepiece>=0.1.99
+protobuf>=4.25.0
+pyyaml>=6.0
+fire>=0.5.0

train.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python3
+"""
+OrbGen Training Script
+
+Fine-tunes a base model to generate valid Orbital schemas (.orb files).
+
+Usage:
+    python train.py --config config.yaml
+    python train.py --config config.yaml --debug --max_steps 100
+"""
+
+import os
+import yaml
+import fire
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainingArguments,
+    DataCollatorForSeq2Seq,
+)
+from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
+from trl import SFTTrainer, SFTConfig
+import wandb
+
+
+def load_config(config_path: str) -> dict:
+    """Load configuration from YAML file."""
+    with open(config_path, 'r') as f:
+        return yaml.safe_load(f)
+
+
+def format_example(example: dict, tokenizer) -> str:
+    """Format a single training example as a chat conversation."""
+    system_prompt = """You are OrbGen, a specialized AI that generates valid Orbital schemas (.orb files) from natural language descriptions.
+
+Rules:
+1. Output ONLY valid JSON - no explanations, no markdown code blocks
+2. Every schema must have: name, version, orbitals array
+3. Each orbital must have: name, entity, traits, pages
+4. Each entity must have: name, collection (or runtime/singleton), fields
+5. Each trait must have: name, category (interaction/integration), linkedEntity, stateMachine
+6. State machines must have: states (with one isInitial:true), events, transitions
+7. Use S-expression arrays for effects: ["set", "field", "value"], ["emit", "EVENT", {}], ["render-ui", "slot", {...}]
+8. Pages must have: name, path, entity, traits"""
+
+    return f"""<|im_start|>system
+{system_prompt}
+<|im_end|>
+<|im_start|>user
+{example['prompt']}
+<|im_end|>
+<|im_start|>assistant
+{example['completion']}
+<|im_end|>"""
+
+
+def main(
+    config: str = "config.yaml",
+    debug: bool = False,
+    max_steps: int = -1,
+    resume_from_checkpoint: str = None,
+):
+    """Main training function."""
+
+    # Load configuration
+    cfg = load_config(config)
+
+    print("=" * 60)
+    print("OrbGen Training")
+    print("=" * 60)
+    print(f"Base model: {cfg['model']['base_model']}")
+    print(f"Output dir: {cfg['model']['output_dir']}")
+    print(f"Debug mode: {debug}")
+    print("=" * 60)
+
+    # Initialize wandb
+    if not debug:
+        wandb.init(
+            project=cfg['wandb']['project'],
+            entity=cfg['wandb'].get('entity'),
+            name=cfg['wandb']['run_name'],
+            config=cfg,
+        )
+
+    # Load tokenizer
+    print("\nLoading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        cfg['model']['base_model'],
+        trust_remote_code=True,
+    )
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+
+    # Load model
+    print("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        cfg['model']['base_model'],
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+
+    # Prepare model for training
+    model.config.use_cache = False
+    model.enable_input_require_grads()
+
+    # Configure LoRA
+    if cfg['lora']['enabled']:
+        print("Configuring LoRA...")
+        lora_config = LoraConfig(
+            r=cfg['lora']['r'],
+            lora_alpha=cfg['lora']['lora_alpha'],
+            lora_dropout=cfg['lora']['lora_dropout'],
+            target_modules=cfg['lora']['target_modules'],
+            bias=cfg['lora']['bias'],
+            task_type=TaskType.CAUSAL_LM,
+        )
+        model = get_peft_model(model, lora_config)
+        model.print_trainable_parameters()
+
+    # Load dataset
+    print("\nLoading dataset...")
+    dataset = load_dataset(cfg['data']['dataset'])
+
+    train_dataset = dataset[cfg['data']['train_split']]
+    eval_dataset = dataset[cfg['data']['eval_split']]
+
+    print(f"Train examples: {len(train_dataset)}")
+    print(f"Eval examples: {len(eval_dataset)}")
+
+    # Format dataset
+    def format_dataset(examples):
+        texts = []
+        for i in range(len(examples['prompt'])):
+            example = {
+                'prompt': examples['prompt'][i],
+                'completion': examples['completion'][i],
+            }
+            texts.append(format_example(example, tokenizer))
+        return {'text': texts}
+
+    train_dataset = train_dataset.map(
+        format_dataset,
+        batched=True,
+        remove_columns=train_dataset.column_names,
+    )
+
+    eval_dataset = eval_dataset.map(
+        format_dataset,
+        batched=True,
+        remove_columns=eval_dataset.column_names,
+    )
+
+    # Training arguments
+    training_args = SFTConfig(
+        output_dir=cfg['model']['output_dir'],
+        num_train_epochs=cfg['training']['num_epochs'] if not debug else 1,
+        per_device_train_batch_size=cfg['training']['per_device_train_batch_size'],
+        per_device_eval_batch_size=cfg['training']['per_device_eval_batch_size'],
+        gradient_accumulation_steps=cfg['training']['gradient_accumulation_steps'],
+        learning_rate=cfg['training']['learning_rate'],
+        warmup_ratio=cfg['training']['warmup_ratio'],
+        weight_decay=cfg['training']['weight_decay'],
+        max_grad_norm=cfg['training']['max_grad_norm'],
+        logging_steps=cfg['training']['logging_steps'],
+        eval_strategy="steps",
+        eval_steps=cfg['training']['eval_steps'],
+        save_steps=cfg['training']['save_steps'],
+        save_total_limit=cfg['training']['save_total_limit'],
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        greater_is_better=False,
+        bf16=True,
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        max_seq_length=cfg['model']['max_seq_length'],
+        dataset_text_field="text",
+        report_to="wandb" if not debug else "none",
+        max_steps=max_steps if max_steps > 0 else -1,
+    )
+
+    # Create trainer
+    trainer = SFTTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer,
+    )
+
+    # Train
+    print("\nStarting training...")
+    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+
+    # Save final model
+    print("\nSaving model...")
+    trainer.save_model(f"{cfg['model']['output_dir']}/final")
+    tokenizer.save_pretrained(f"{cfg['model']['output_dir']}/final")
+
+    # Finish wandb
+    if not debug:
+        wandb.finish()
+
+    print("\nTraining complete!")
+    print(f"Model saved to: {cfg['model']['output_dir']}/final")
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

utils/__init__.py

@@ -0,0 +1 @@
+# OrbGen Training Utilities