Upload folder using huggingface_hub

Dockerfile

@@ -1,4 +1,10 @@
-FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel
+FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-devel
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1
+ENV HF_HOME=/app/.cache/huggingface
+ENV TRANSFORMERS_CACHE=/app/.cache/huggingface/transformers
+ENV DEBIAN_FRONTEND=noninteractive
 
 WORKDIR /app
 
@@ -6,18 +12,25 @@ WORKDIR /app
 RUN apt-get update && apt-get install -y \
     git \
     curl \
+    wget \
+    build-essential \
     && rm -rf /var/lib/apt/lists/*
 
+# Download orbital CLI from HuggingFace (has native validation, no kflow needed)
+RUN wget -q https://huggingface.co/orbital-ai/orbital-cli/resolve/main/orbital-linux-x86_64 -O /usr/local/bin/orbital \
+    && chmod +x /usr/local/bin/orbital \
+    && orbital --version || echo "Orbital CLI installed"
+
 # Copy requirements and install Python dependencies
 COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
 
 # Copy training code
 COPY . .
 
-# Set environment variables
-ENV PYTHONUNBUFFERED=1
-ENV HF_HOME=/app/.cache/huggingface
+# Create cache directories
+RUN mkdir -p /app/.cache/huggingface
 
-# Default command
-CMD ["python", "train.py", "--config", "config.yaml"]
+# Default command - uses HuggingFace config
+CMD ["python", "train.py", "--config", "config-huggingface.yaml"]

README.md

@@ -1,37 +1,47 @@
+---
+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
 
-### 1. Install Dependencies
+### Local Training
 
 ```bash
+# Install dependencies
 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
+# 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
+```
 
-# Resume from checkpoint
-python train.py --config config.yaml --resume_from_checkpoint ./orbgen-1.5b/checkpoint-500
+### HuggingFace Cloud Training
+
+```bash
+# Use the HuggingFace-optimized config
+python train.py --config config-huggingface.yaml
 ```
 
-### 4. Evaluate
+## 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
@@ -41,7 +51,7 @@ python evaluate.py --checkpoint ./orbgen-1.5b/final
 python evaluate.py --checkpoint ./orbgen-1.5b/final --use_validator
 ```
 
-### 5. Generate
+## Generate
 
 ```bash
 # Single generation
@@ -61,28 +71,41 @@ python generate.py --prompt "..." --output schema.orb --validate
 | `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 |
+| `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
 
-1. Push this directory to `orbital-ai/orbgen-training`:
+### Prerequisites
+
+1. Upload dataset to HuggingFace:
    ```bash
-   hf upload orbital-ai/orbgen-training . --repo-type space
+   python scripts/upload_dataset.py
    ```
 
-2. Configure the Space with A10G GPU
+2. Set Space secrets:
+   - `HUGGINGFACE_TOKEN` - HF token with write access
+   - `WANDB_API_KEY` - Weights & Biases API key
 
-3. Training will start automatically
+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 | Time |
-|-------|-----|------|------|
-| Training (3 epochs) | A10G | 24GB | ~6 hours |
-| Evaluation | T4 | 16GB | ~30 min |
-| Inference | T4 | 16GB | ~1 sec/gen |
+| 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
 
@@ -92,7 +115,4 @@ After training, the model is saved to `./orbgen-1.5b/final/`:
 - `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
-```
+Model is automatically pushed to `orbital-ai/orbgen-1.5b` when using HuggingFace config.

config-huggingface.yaml

@@ -0,0 +1,67 @@
+# OrbGen Training Configuration
+# Optimized for HuggingFace Spaces with A10G GPU (24GB VRAM)
+
+model:
+  base_model: "Qwen/Qwen2.5-Coder-1.5B"
+  output_dir: "./orbgen-1.5b"
+  max_seq_length: 4096  # Full context for schema generation
+
+data:
+  # Load from HuggingFace Hub (upload dataset first)
+  dataset: "orbital-ai/orbital-schemas"
+  train_split: "train"
+  eval_split: "validation"
+
+training:
+  # SFT Configuration - optimized for A10G (24GB VRAM)
+  num_epochs: 3
+  per_device_train_batch_size: 4   # Can use larger batches
+  per_device_eval_batch_size: 4
+  gradient_accumulation_steps: 4   # Effective batch size = 16
+  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          # Full LoRA rank
+  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"
+
+# No quantization needed - enough VRAM for bf16
+quantization:
+  enabled: false
+
+generation:
+  max_new_tokens: 4096
+  temperature: 0.7
+  top_p: 0.95
+  do_sample: true
+
+wandb:
+  project: "orbgen-training"
+  entity: null  # Will use default
+  run_name: "orbgen-1.5b-sft-hf"
+
+# HuggingFace Hub settings
+hub:
+  push_to_hub: true
+  hub_model_id: "orbital-ai/orbgen-1.5b"
+  hub_strategy: "checkpoint"

config.yaml

@@ -1,21 +1,22 @@
 # OrbGen Training Configuration
+# Optimized for 6GB VRAM (Quadro RTX 3000)
 
 model:
   base_model: "Qwen/Qwen2.5-Coder-1.5B"
   output_dir: "./orbgen-1.5b"
-  max_seq_length: 8192
+  max_seq_length: 2048  # Reduced from 8192 for low VRAM
 
 data:
-  dataset: "orbital-ai/orbital-schemas"
-  train_split: "train"
-  eval_split: "validation"
+  # Use local files instead of HuggingFace (schema mismatch issue)
+  train_file: "../training-data/combined-train.jsonl"
+  eval_file: "../training-data/combined-validation.jsonl"
 
 training:
-  # SFT Configuration
+  # SFT Configuration - optimized for 6GB VRAM
   num_epochs: 3
-  per_device_train_batch_size: 2
-  per_device_eval_batch_size: 2
-  gradient_accumulation_steps: 8
+  per_device_train_batch_size: 1  # Reduced from 2
+  per_device_eval_batch_size: 1   # Reduced from 2
+  gradient_accumulation_steps: 16  # Increased to compensate for smaller batch
   learning_rate: 2.0e-5
   warmup_ratio: 0.1
   weight_decay: 0.01
@@ -25,12 +26,12 @@ training:
   logging_steps: 10
   eval_steps: 50
   save_steps: 100
-  save_total_limit: 3
+  save_total_limit: 2  # Reduced to save disk space
 
 lora:
   enabled: true
-  r: 64
-  lora_alpha: 128
+  r: 32  # Reduced from 64 to save memory
+  lora_alpha: 64  # Reduced proportionally
   lora_dropout: 0.05
   target_modules:
     - "q_proj"
@@ -43,8 +44,16 @@ lora:
   bias: "none"
   task_type: "CAUSAL_LM"
 
+# 4-bit quantization for low VRAM
+quantization:
+  enabled: true
+  load_in_4bit: true
+  bnb_4bit_compute_dtype: "bfloat16"
+  bnb_4bit_quant_type: "nf4"
+  bnb_4bit_use_double_quant: true
+
 generation:
-  max_new_tokens: 4096
+  max_new_tokens: 2048  # Reduced from 4096
   temperature: 0.7
   top_p: 0.95
   do_sample: true
@@ -52,4 +61,4 @@ generation:
 wandb:
   project: "orbgen-training"
   entity: "orbital-ai"
-  run_name: "orbgen-1.5b-sft"
+  run_name: "orbgen-1.5b-sft-qlora"

evaluate.py

@@ -36,12 +36,18 @@ def validate_schema(schema_json: str) -> tuple[bool, list[str]]:
         temp_path = f.name
 
     try:
+        # Find orbital binary - check multiple locations
+        orbital_cmd = 'orbital'
+        for path in ['/usr/local/bin/orbital', os.path.expanduser('~/kflow.ai.builder/orbital-rust/target/release/orbital')]:
+            if os.path.exists(path):
+                orbital_cmd = path
+                break
+
         result = subprocess.run(
-            ['orbital', 'validate', temp_path],
+            [orbital_cmd, '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:
@@ -52,7 +58,7 @@ def validate_schema(schema_json: str) -> tuple[bool, list[str]]:
     except subprocess.TimeoutExpired:
         return False, ["Validation timeout"]
     except FileNotFoundError:
-        return False, ["Orbital CLI not found"]
+        return False, ["Orbital CLI not found - install it or use --use_validator=False"]
     except Exception as e:
         return False, [f"Validation error: {e}"]
     finally:

requirements.txt

@@ -1,15 +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
+transformers>=4.40.0
+datasets>=2.18.0
+peft>=0.10.0
+trl>=0.8.0
+accelerate>=0.28.0
+bitsandbytes>=0.43.0
 wandb>=0.16.0
-huggingface_hub>=0.20.0
+huggingface_hub>=0.22.0
 safetensors>=0.4.0
-sentencepiece>=0.1.99
+sentencepiece>=0.2.0
 protobuf>=4.25.0
 pyyaml>=6.0
 fire>=0.5.0

scripts/bundle_kflow.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+# Bundle kflow CLI into a single file for distribution
+# Run from builder/ directory
+
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+BUILDER_DIR="$SCRIPT_DIR/../../../../builder"
+OUTPUT_DIR="$SCRIPT_DIR/../bin"
+
+echo "Bundling kflow CLI..."
+echo "Builder dir: $BUILDER_DIR"
+
+cd "$BUILDER_DIR"
+
+# Ensure dependencies are installed
+npm install --workspace=@kflow-builder/compiler --workspace=@kflow-builder/shared
+
+# Build shared first
+npm run build:shared
+
+# Bundle kflow CLI with esbuild
+mkdir -p "$OUTPUT_DIR"
+
+npx esbuild packages/compiler/src/cli/index.ts \
+  --bundle \
+  --platform=node \
+  --target=node18 \
+  --outfile="$OUTPUT_DIR/kflow-bundle.js" \
+  --external:esbuild \
+  --external:typescript \
+  --external:@swc/core \
+  --format=esm \
+  --sourcemap
+
+# Create wrapper script
+cat > "$OUTPUT_DIR/kflow" << 'EOF'
+#!/usr/bin/env node
+import('./kflow-bundle.js');
+EOF
+
+chmod +x "$OUTPUT_DIR/kflow"
+
+echo "Bundle created at: $OUTPUT_DIR/kflow-bundle.js"
+echo "Wrapper at: $OUTPUT_DIR/kflow"
+
+# Get size
+ls -lh "$OUTPUT_DIR/kflow-bundle.js"

scripts/upload_dataset.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python3
+"""
+Upload OrbGen training dataset to HuggingFace Hub.
+
+Usage:
+    python scripts/upload_dataset.py
+    python scripts/upload_dataset.py --repo orbital-ai/orbital-schemas
+"""
+
+import os
+import json
+import argparse
+from pathlib import Path
+from huggingface_hub import HfApi, create_repo, upload_file
+from datasets import Dataset, DatasetDict
+
+
+def load_jsonl(path: str) -> list:
+    """Load JSONL file into list of dicts."""
+    data = []
+    with open(path, 'r') as f:
+        for line in f:
+            if line.strip():
+                data.append(json.loads(line))
+    return data
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Upload dataset to HuggingFace')
+    parser.add_argument('--repo', default='orbital-ai/orbital-schemas',
+                        help='HuggingFace dataset repository')
+    parser.add_argument('--data-dir', default='../../training-data',
+                        help='Directory containing JSONL files')
+    parser.add_argument('--private', action='store_true',
+                        help='Make dataset private')
+    args = parser.parse_args()
+
+    # Resolve paths
+    script_dir = Path(__file__).parent
+    data_dir = (script_dir / args.data_dir).resolve()
+
+    print(f"Loading data from: {data_dir}")
+
+    # Load training data
+    train_path = data_dir / 'combined-train.jsonl'
+    val_path = data_dir / 'combined-validation.jsonl'
+    test_path = data_dir / 'test.jsonl'
+
+    if not train_path.exists():
+        print(f"Error: {train_path} not found")
+        return 1
+
+    train_data = load_jsonl(str(train_path))
+    val_data = load_jsonl(str(val_path)) if val_path.exists() else []
+    test_data = load_jsonl(str(test_path)) if test_path.exists() else []
+
+    print(f"Loaded {len(train_data)} train, {len(val_data)} validation, {len(test_data)} test examples")
+
+    # Create datasets
+    def process_examples(examples):
+        """Ensure consistent schema."""
+        processed = []
+        for ex in examples:
+            processed.append({
+                'prompt': ex['prompt'],
+                'completion': ex['completion'],
+                'domain': ex.get('metadata', {}).get('domain', 'general'),
+                'complexity': ex.get('metadata', {}).get('complexity', 'medium'),
+                'source': ex.get('metadata', {}).get('source', 'unknown'),
+            })
+        return processed
+
+    train_ds = Dataset.from_list(process_examples(train_data))
+    val_ds = Dataset.from_list(process_examples(val_data)) if val_data else None
+    test_ds = Dataset.from_list(process_examples(test_data)) if test_data else None
+
+    # Create DatasetDict
+    splits = {'train': train_ds}
+    if val_ds:
+        splits['validation'] = val_ds
+    if test_ds:
+        splits['test'] = test_ds
+
+    dataset_dict = DatasetDict(splits)
+
+    print(f"\nDataset structure:")
+    print(dataset_dict)
+
+    # Create repo if needed
+    api = HfApi()
+    try:
+        create_repo(args.repo, repo_type='dataset', private=args.private, exist_ok=True)
+        print(f"\nRepository: https://huggingface.co/datasets/{args.repo}")
+    except Exception as e:
+        print(f"Note: {e}")
+
+    # Push to hub
+    print(f"\nPushing to HuggingFace Hub...")
+    dataset_dict.push_to_hub(
+        args.repo,
+        private=args.private,
+        commit_message="Update training dataset"
+    )
+
+    print(f"\nDataset uploaded successfully!")
+    print(f"View at: https://huggingface.co/datasets/{args.repo}")
+
+    # Create dataset card
+    dataset_card = f"""---
+license: apache-2.0
+task_categories:
+  - text-generation
+language:
+  - en
+tags:
+  - orbital
+  - schema-generation
+  - code
+size_categories:
+  - n<1K
+---
+
+# Orbital Schemas Dataset
+
+Training data for OrbGen - a model that generates valid Orbital schemas (.orb files).
+
+## Dataset Structure
+
+- **train**: {len(train_data)} examples
+- **validation**: {len(val_data)} examples
+- **test**: {len(test_data)} examples
+
+## Features
+
+- `prompt`: Natural language description of the desired schema
+- `completion`: Valid Orbital schema in JSON format
+- `domain`: Application domain (ecommerce, game, productivity, etc.)
+- `complexity`: Schema complexity (simple, medium, complex)
+- `source`: Source of the example (synthetic, pattern, integrator)
+
+## Usage
+
+```python
+from datasets import load_dataset
+
+dataset = load_dataset("{args.repo}")
+print(dataset["train"][0])
+```
+
+## Example
+
+```json
+{{
+  "prompt": "Create a task management app with projects and due dates",
+  "completion": "{{...valid orbital schema...}}",
+  "domain": "productivity",
+  "complexity": "medium",
+  "source": "synthetic"
+}}
+```
+
+## License
+
+Apache 2.0
+"""
+
+    # Save and upload README
+    readme_path = data_dir / 'README.md'
+    with open(readme_path, 'w') as f:
+        f.write(dataset_card)
+
+    upload_file(
+        path_or_fileobj=str(readme_path),
+        path_in_repo='README.md',
+        repo_id=args.repo,
+        repo_type='dataset',
+        commit_message='Add dataset card'
+    )
+
+    print(f"Dataset card uploaded!")
+    return 0
+
+
+if __name__ == '__main__':
+    exit(main())

scripts/upload_orbital_cli.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+"""
+Upload orbital CLI binary to HuggingFace Hub.
+
+Usage:
+    # First build the binary:
+    cd orbital-rust && cargo build --release --bin orbital
+
+    # Then upload:
+    python scripts/upload_orbital_cli.py
+"""
+
+import os
+import argparse
+from pathlib import Path
+from huggingface_hub import HfApi, create_repo, upload_file
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Upload orbital CLI to HuggingFace')
+    parser.add_argument('--repo', default='orbital-ai/orbital-cli',
+                        help='HuggingFace repository')
+    parser.add_argument('--binary', default=None,
+                        help='Path to orbital binary (auto-detected if not specified)')
+    args = parser.parse_args()
+
+    # Find binary
+    if args.binary:
+        binary_path = Path(args.binary)
+    else:
+        # Try common locations
+        candidates = [
+            Path.home() / 'kflow.ai.builder/orbital-rust/target/release/orbital',
+            Path(__file__).parent.parent.parent.parent / 'orbital-rust/target/release/orbital',
+            Path('/home/osamah/kflow.ai.builder/orbital-rust/target/release/orbital'),
+        ]
+        binary_path = None
+        for candidate in candidates:
+            if candidate.exists():
+                binary_path = candidate
+                break
+
+        if not binary_path:
+            print("Error: Could not find orbital binary. Build it first:")
+            print("  cd orbital-rust && cargo build --release --bin orbital")
+            return 1
+
+    print(f"Binary: {binary_path}")
+    print(f"Size: {binary_path.stat().st_size / 1024 / 1024:.1f} MB")
+
+    # Create repo
+    api = HfApi()
+    try:
+        create_repo(args.repo, repo_type='model', exist_ok=True)
+        print(f"Repository: https://huggingface.co/{args.repo}")
+    except Exception as e:
+        print(f"Note: {e}")
+
+    # Upload binary
+    print(f"\nUploading to {args.repo}...")
+    upload_file(
+        path_or_fileobj=str(binary_path),
+        path_in_repo='orbital-linux-x86_64',
+        repo_id=args.repo,
+        commit_message='Upload orbital CLI binary for Linux x86_64'
+    )
+
+    # Create README
+    readme = """---
+license: apache-2.0
+tags:
+  - orbital
+  - cli
+  - schema-validation
+---
+
+# Orbital CLI
+
+Binary releases of the Orbital CLI for use in training pipelines.
+
+## Files
+
+- `orbital-linux-x86_64` - Linux x86_64 binary
+
+## Usage
+
+```bash
+# Download
+wget https://huggingface.co/orbital-ai/orbital-cli/resolve/main/orbital-linux-x86_64 -O orbital
+chmod +x orbital
+
+# Validate a schema
+./orbital validate schema.orb
+```
+
+## Building from Source
+
+```bash
+cd orbital-rust
+cargo build --release --bin orbital
+```
+"""
+
+    readme_path = Path('/tmp/orbital-cli-readme.md')
+    readme_path.write_text(readme)
+
+    upload_file(
+        path_or_fileobj=str(readme_path),
+        path_in_repo='README.md',
+        repo_id=args.repo,
+        commit_message='Add README'
+    )
+
+    print(f"\nDone! Binary available at:")
+    print(f"  https://huggingface.co/{args.repo}/resolve/main/orbital-linux-x86_64")
+    return 0
+
+
+if __name__ == '__main__':
+    exit(main())

train.py

@@ -19,6 +19,7 @@ from transformers import (
     AutoTokenizer,
     TrainingArguments,
     DataCollatorForSeq2Seq,
+    BitsAndBytesConfig,
 )
 from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
 from trl import SFTTrainer, SFTConfig
@@ -93,18 +94,41 @@ def main(
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "right"
 
-    # Load model
+    # Load model with optional quantization
     print("Loading model...")
+
+    model_kwargs = {
+        "trust_remote_code": True,
+        "device_map": "auto",
+    }
+
+    # Check if 4-bit quantization is enabled
+    quant_cfg = cfg.get('quantization', {})
+    if quant_cfg.get('enabled', False) and quant_cfg.get('load_in_4bit', False):
+        print("Using 4-bit quantization (QLoRA)...")
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=getattr(torch, quant_cfg.get('bnb_4bit_compute_dtype', 'bfloat16')),
+            bnb_4bit_quant_type=quant_cfg.get('bnb_4bit_quant_type', 'nf4'),
+            bnb_4bit_use_double_quant=quant_cfg.get('bnb_4bit_use_double_quant', True),
+        )
+        model_kwargs["quantization_config"] = bnb_config
+    else:
+        model_kwargs["torch_dtype"] = torch.bfloat16
+
     model = AutoModelForCausalLM.from_pretrained(
         cfg['model']['base_model'],
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        trust_remote_code=True,
+        **model_kwargs,
     )
 
     # Prepare model for training
     model.config.use_cache = False
-    model.enable_input_require_grads()
+
+    # For quantized models, use prepare_model_for_kbit_training
+    if quant_cfg.get('enabled', False):
+        model = prepare_model_for_kbit_training(model)
+    else:
+        model.enable_input_require_grads()
 
     # Configure LoRA
     if cfg['lora']['enabled']:
@@ -122,10 +146,22 @@ def main(
 
     # 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']]
+    # Support both HuggingFace dataset and local files
+    if 'train_file' in cfg['data']:
+        # Load from local JSONL files
+        data_files = {
+            'train': cfg['data']['train_file'],
+            'validation': cfg['data']['eval_file'],
+        }
+        dataset = load_dataset('json', data_files=data_files)
+        train_dataset = dataset['train']
+        eval_dataset = dataset['validation']
+    else:
+        # Load from HuggingFace Hub
+        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)}")
@@ -175,19 +211,19 @@ def main(
         bf16=True,
         gradient_checkpointing=True,
         gradient_checkpointing_kwargs={"use_reentrant": False},
-        max_seq_length=cfg['model']['max_seq_length'],
+        max_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
+    # Create trainer (TRL v0.27+ API)
     trainer = SFTTrainer(
         model=model,
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
-        tokenizer=tokenizer,
+        processing_class=tokenizer,
     )
 
     # Train
@@ -199,6 +235,14 @@ def main(
     trainer.save_model(f"{cfg['model']['output_dir']}/final")
     tokenizer.save_pretrained(f"{cfg['model']['output_dir']}/final")
 
+    # Push to HuggingFace Hub if configured
+    hub_cfg = cfg.get('hub', {})
+    if hub_cfg.get('push_to_hub', False) and not debug:
+        print("\nPushing model to HuggingFace Hub...")
+        hub_model_id = hub_cfg.get('hub_model_id', 'orbital-ai/orbgen-1.5b')
+        trainer.push_to_hub(commit_message="Final model after SFT training")
+        print(f"Model pushed to: https://huggingface.co/{hub_model_id}")
+
     # Finish wandb
     if not debug:
         wandb.finish()