additional fixes

train-on-this-server.sh

@@ -1,92 +1,16 @@
 #!/bin/bash
-# Train LoRA on this server (assumes GPU is available).
-# Run this after deploy-and-train.sh has cloned the repo.
+# Training script for Cyron LoRA on this server (2x RTX 5090)
 
 set -e
 
-REPO_DIR="${1:-$(pwd)}"
+echo "=== Cyron LoRA Training Setup ==="
 
-echo "=== LoRA Training Setup ==="
-echo "Repo: ${REPO_DIR}"
-echo ""
+# Install dependencies if needed
+pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
+pip install -q transformers accelerate datasets trl peft bitsandbytes deepspeed
 
-# Unload AI model to free VRAM
-echo "=== Unloading AI model ==="
-if command -v ai &> /dev/null; then
-    ai none
-    sleep 2
-else
-    echo "Warning: 'ai' command not found, skipping AI unload"
-fi
+# Run training
+echo "Starting training with accelerate..."
+accelerate launch --multi_gpu --num_processes 2 /home/workdir/artifacts/train.py --config training/configs/llama2-7b-lora.yaml
 
-# Check Python
-if ! command -v python3 &> /dev/null; then
-    echo "ERROR: python3 not found"
-    exit 1
-fi
-
-echo "[1/4] Setup Python environment..."
-python3 -m venv venv
-source venv/bin/activate
-pip install --upgrade pip
-
-echo "[2/4] Install PyTorch with CUDA 13.0 (RTX 5090 compatible)..."
-pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu130
-
-echo "[3/4] Install training dependencies..."
-pip install transformers datasets trl peft accelerate bitsandbytes deepspeed compressed-tensors
-
-# Check GPU
-echo ""
-echo "[4/4] Checking GPU..."
-python3 -c "
-import torch
-if torch.cuda.is_available():
-    print(f'  ✓ GPU: {torch.cuda.get_device_name(0)}')
-    props = torch.cuda.get_device_properties(0)
-    print(f'  ✓ VRAM: {props.total_memory / 1e9:.1f} GB')
-else:
-    print('  ✗ No GPU detected')
-    print('  Please ensure CUDA drivers are installed')
-    exit(1)
-"
-
-echo ""
-echo "=== Dataset Preparation ==="
-echo "Checking for combined_20k.jsonl..."
-if [ ! -f "dataset/combined_20k.jsonl" ]; then
-    echo "ERROR: Dataset not found!"
-    echo "The file should be in the repo at: dataset/combined_20k.jsonl"
-    exit 1
-fi
-
-echo "Checking for train.jsonl and test.jsonl..."
-if [ ! -f "training/data/train.jsonl" ] || [ ! -f "training/data/test.jsonl" ]; then
-    echo "Split files not found, creating them..."
-    python3 training/scripts/prepare_dataset.py --input dataset/combined_20k.jsonl --output training/data
-else
-    echo "Split files already exist, using them."
-fi
-
-echo ""
-echo "=== Starting Training ==="
-echo "Config: training/configs/ornith-35b-lora.yaml"
-echo "Output: training/output/ornith-35b-lora/"
-echo ""
-echo "Training will take 6-24 hours depending on GPU."
-echo "Press Ctrl+C to stop (model will be saved at checkpoint)."
-echo ""
-
-# Use torchrun for distributed training (2 GPUs)
-echo "Starting training..."
-torchrun --nproc_per_node=2 training/scripts/train.py --config training/configs/ornith-35b-lora.yaml
-
-echo ""
-echo "=== Training Complete ==="
-echo "Model saved to: training/output/llama2-7b-lora/"
-echo ""
-echo "To generate summaries:"
-echo "  python3 training/scripts/inference.py \\"
-echo "      --model training/output/llama2-7b-lora \\"
-echo "      --task \"Your task here\" \\"
-echo "      --files src/file.py"
+echo "Training completed!"

training/scripts/inference.py

@@ -68,8 +68,8 @@ def generate_summary(model, tokenizer, task, files_changed=None, tests_run=False
 
 def main():
     parser = argparse.ArgumentParser(description="Generate Cyron summaries with LoRA")
-    parser.add_argument("--model", type=str, default="output/llama2-7b-lora", 
-                       help="Path to trained model")
+    parser.add_argument("--model", type=str, default="output/ornith-35b-lora-cyron", 
+                       help="Path to trained LoRA model")
     parser.add_argument("--task", type=str, required=True, help="Task description")
     parser.add_argument("--files", type=str, nargs="*", default=[], help="Files changed")
     parser.add_argument("--tests-run", action="store_true", help="Tests were run")

training/scripts/run_with_accelerate.sh

@@ -7,4 +7,5 @@ set -e
 pip install -q accelerate
 
 # Launch with accelerate
-accelerate launch train-on-this-server.sh --config training/configs/llama2-7b-lora.yaml
+# Note: Use the dedicated train-on-this-server.sh instead
+echo "See train-on-this-server.sh for full setup"