init: add LoRA training infrastructure and 20k dataset

training/README.md

@@ -0,0 +1,129 @@
+# Cyron Summary LoRA Training
+
+This directory contains scripts and configurations for training a LoRA adapter on the Cyron summary dataset.
+
+## Directory Structure
+
+```
+training/
+├── configs/          # Training configurations
+│   └── llama2-7b-lora.yaml
+├── data/            # Prepared datasets (train/test splits)
+├── output/          # Trained model outputs
+├── scripts/         # Training and inference scripts
+│   ├── prepare_dataset.py
+│   ├── train.py
+│   └── inference.py
+└── notebooks/       # Jupyter notebooks for analysis
+```
+
+## Prerequisites
+
+- Python 3.10+
+- GPU with 40GB+ VRAM (A100 recommended)
+- 64GB+ system RAM
+- 100GB+ free disk space
+
+## Installation
+
+```bash
+cd scripts/lora_training/training
+
+# Create virtual environment
+python -m venv venv
+source venv/bin/activate
+
+# Install dependencies
+pip install transformers datasets trl peft accelerate bitsandbytes
+
+# Or use conda
+conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
+pip install transformers datasets trl peft accelerate bitsandbytes
+```
+
+## Usage
+
+### 1. Prepare Dataset
+
+```bash
+python scripts/prepare_dataset.py --input ../combined_20k.jsonl --output data
+```
+
+This splits the 20k dataset into train/test sets (95/5).
+
+### 2. Train Model
+
+```bash
+python scripts/train.py --config configs/llama2-7b-lora.yaml
+```
+
+Or with custom parameters:
+
+```bash
+python scripts/train.py --config configs/llama2-7b-lora.yaml --epochs 5 --batch-size 8
+```
+
+Training takes approximately 6-24 hours depending on GPU.
+
+### 3. Generate Summaries
+
+```bash
+python scripts/inference.py \
+    --model output/llama2-7b-lora \
+    --task "Fix parser crash on malformed JSON" \
+    --files src/parser.cpp \
+    --tests-run --test-count 294 \
+    --commit --push
+```
+
+## Configuration
+
+Edit `configs/llama2-7b-lora.yaml` to adjust:
+
+- **Base model**: Change `base_model` to use different foundation models
+- **LoRA rank**: Adjust `lora_r` (higher = more capacity, slower)
+- **Learning rate**: Tune `learning_rate`
+- **Epochs**: Change `num_train_epochs`
+- **Batch size**: Adjust `per_device_train_batch_size`
+
+## Dataset Format
+
+The dataset (`combined_20k.jsonl`) contains 20,000 examples with:
+
+- `task`: Task description
+- `files_changed`: List of changed files
+- `tests_run`: Boolean
+- `commit`: Boolean
+- `push`: Boolean
+- `errors_seen`: Boolean
+- `test_count`: Number of tests
+- `output`: Expected Cyron summary
+
+## Output
+
+Trained model is saved to `output/llama2-7b-lora/`:
+
+- `adapter_model.safetensors` - LoRA weights
+- `config.json` - Model configuration
+- `tokenizer.json` - Tokenizer
+- `peft_config.json` - LoRA configuration
+
+## Troubleshooting
+
+**Out of memory:**
+- Reduce `per_device_train_batch_size`
+- Enable `gradient_checkpointing: true`
+- Use `load_in_4bit: true` (already enabled)
+
+**Training loss not decreasing:**
+- Lower learning rate (try 1e-4)
+- Increase warmup ratio
+- Check dataset quality
+
+**Generation too long/short:**
+- Adjust `max_new_tokens` in inference script
+- Tune `temperature` and `top_p`
+
+## License
+
+This training infrastructure is part of the AgenX project.

training/configs/llama2-7b-lora.yaml

@@ -0,0 +1,55 @@
+# LoRA Training Configuration for Llama-2-7b
+# Dataset: cyron_summary_lora_dataset (20k examples)
+
+base_model: meta-llama/Llama-2-7b-hf
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+
+# Quantization (QLoRA)
+load_in_4bit: true
+bnb_4bit_compute_dtype: bfloat16
+bnb_4bit_quant_type: nf4
+use_nested_quant: false
+
+# LoRA Configuration
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+target_modules:
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+lora_task_type: CAUSAL_LM
+
+# Dataset
+dataset:
+  - path: ../combined_20k.jsonl
+    type: completion
+    text_column: text
+
+# Training Parameters
+train_params:
+  num_train_epochs: 3
+  per_device_train_batch_size: 4
+  gradient_accumulation_steps: 4
+  learning_rate: 2e-4
+  lr_scheduler_type: cosine
+  weight_decay: 0.01
+  warmup_ratio: 0.03
+  max_seq_length: 1024
+  logging_steps: 10
+  save_steps: 100
+  save_total_limit: 3
+  output_dir: ../../output/llama2-7b-lora
+
+# Precision
+mixed_precision: bf16
+
+# Evaluation
+eval_strategy: steps
+eval_steps: 100
+eval_accumulation_steps: 10
+
+# Gradient Checkpointing
+gradient_checkpointing: true

training/scripts/inference.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""
+Inference script for trained LoRA adapter.
+
+Loads the trained model and generates Cyron summaries.
+"""
+
+import argparse
+from pathlib import Path
+
+
+def load_model(model_path):
+    """Load trained LoRA model."""
+    from peft import PeftModel
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    
+    print(f"Loading base model from {model_path}...")
+    base_model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        device_map="auto",
+        torch_dtype="auto",
+    )
+    
+    print(f"Loading LoRA weights from {model_path}...")
+    model = PeftModel.from_pretrained(base_model, model_path)
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    
+    return model, tokenizer
+
+
+def generate_summary(model, tokenizer, task, files_changed=None, tests_run=False, test_count=0, commit=False, push=False, errors_seen=False):
+    """Generate a Cyron summary for a task."""
+    
+    # Build prompt
+    prompt = f"Task: {task}"
+    if files_changed:
+        prompt += f"\nFiles: {', '.join(files_changed[:3])}"
+    
+    if tests_run:
+        prompt += f"\nTests run: {test_count}"
+    if commit:
+        prompt += "\nCommit: yes"
+    if push:
+        prompt += "\nPush: yes"
+    if errors_seen:
+        prompt += "\nErrors seen: yes"
+    
+    prompt += "\n\nGenerate a Cyron summary:"
+    
+    # Tokenize
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # Generate
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=500,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.1,
+    )
+    
+    # Decode
+    generated = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    
+    return generated
+
+
+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("--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")
+    parser.add_argument("--test-count", type=int, default=0, help="Number of tests")
+    parser.add_argument("--commit", action="store_true", help="Commit was made")
+    parser.add_argument("--push", action="store_true", help="Code was pushed")
+    parser.add_argument("--errors", action="store_true", help="Errors were seen")
+    args = parser.parse_args()
+    
+    model, tokenizer = load_model(args.model)
+    
+    summary = generate_summary(
+        model,
+        tokenizer,
+        task=args.task,
+        files_changed=args.files,
+        tests_run=args.tests_run,
+        test_count=args.test_count,
+        commit=args.commit,
+        push=args.push,
+        errors_seen=args.errors,
+    )
+    
+    print("\nGenerated Summary:")
+    print("=" * 70)
+    print(summary)
+    print("=" * 70)
+
+
+if __name__ == "__main__":
+    main()

training/scripts/prepare_dataset.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+"""
+Prepare Cyron summary dataset for LoRA training.
+
+Reads combined_20k.jsonl and formats it for Hugging Face TRL training.
+"""
+
+import json
+import argparse
+from pathlib import Path
+
+
+def prepare_dataset(input_file, output_file, test_size=0.05):
+    """Prepare dataset for training."""
+    
+    with open(input_file) as f:
+        examples = [json.loads(line) for line in f]
+    
+    print(f"Loaded {len(examples)} examples")
+    
+    # Format for training
+    formatted = []
+    for ex in examples:
+        # Create instruction-input-output format
+        instruction = f"Task: {ex['task']}"
+        if ex['files_changed']:
+            instruction += f"\nFiles: {', '.join(ex['files_changed'][:3])}"
+        
+        input_text = ""
+        if ex['tests_run']:
+            input_text += f"Tests run: {ex['test_count']}"
+        if ex['commit']:
+            input_text += f"\nCommit: {ex['git']['commit'] if ex.get('git') else 'yes'}"
+        if ex['push']:
+            input_text += "\nPush: yes"
+        
+        output_text = ex['output']
+        
+        # Create conversation format
+        conversation = {
+            "conversations": [
+                {
+                    "from": "human",
+                    "value": f"Generate a Cyron summary for this task:\n\n{instruction}\n\n{input_text}"
+                },
+                {
+                    "from": "gpt",
+                    "value": output_text
+                }
+            ]
+        }
+        
+        formatted.append(conversation)
+    
+    # Split into train/test
+    import random
+    random.seed(42)
+    random.shuffle(formatted)
+    
+    split_point = int(len(formatted) * (1 - test_size))
+    train_data = formatted[:split_point]
+    test_data = formatted[split_point:]
+    
+    # Save
+    with open(output_file.parent / "train.jsonl", "w") as f:
+        for item in train_data:
+            f.write(json.dumps(item) + "\n")
+    
+    with open(output_file.parent / "test.jsonl", "w") as f:
+        for item in test_data:
+            f.write(json.dumps(item) + "\n")
+    
+    print(f"Train: {len(train_data)} examples")
+    print(f"Test: {len(test_data)} examples")
+    print(f"Saved to {output_file.parent}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare LoRA training dataset")
+    parser.add_argument("--input", type=str, default="../combined_20k.jsonl", 
+                       help="Input combined dataset")
+    parser.add_argument("--output", type=str, default="data", 
+                       help="Output directory")
+    parser.add_argument("--test-size", type=float, default=0.05, 
+                       help="Test set percentage")
+    args = parser.parse_args()
+    
+    output_dir = Path(args.output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    prepare_dataset(args.input, output_dir, args.test_size)
+
+
+if __name__ == "__main__":
+    main()

training/scripts/train.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""
+Train LoRA adapter on Cyron summary dataset.
+
+Uses Hugging Face TRL for SFT training with QLoRA.
+"""
+
+import argparse
+import yaml
+from pathlib import Path
+
+
+def train(config_path):
+    """Train LoRA adapter using TRL."""
+    
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoTokenizer,
+        BitsAndBytesConfig,
+        TrainingArguments,
+    )
+    from peft import (
+        prepare_model_for_kbit_training,
+        LoraConfig,
+        get_peft_model,
+    )
+    from trl import SFTTrainer, SFTConfig
+    
+    print(f"Loading model: {config['base_model']}")
+    
+    # Load model with quantization
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=config.get("load_in_4bit", True),
+        bnb_4bit_compute_dtype=config.get("bnb_4bit_compute_dtype", "bfloat16"),
+        bnb_4bit_quant_type=config.get("bnb_4bit_quant_type", "nf4"),
+        use_nested_quant=config.get("use_nested_quant", False),
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        config["base_model"],
+        quantization_config=bnb_config,
+        device_map="auto",
+    )
+    model = prepare_model_for_kbit_training(model)
+    
+    # Add LoRA
+    lora_config = LoraConfig(
+        r=config["lora_r"],
+        lora_alpha=config["lora_alpha"],
+        lora_dropout=config["lora_dropout"],
+        target_modules=config["target_modules"],
+        task_type=config.get("lora_task_type", "CAUSAL_LM"),
+    )
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(config["base_model"])
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+    
+    # Load dataset
+    from datasets import load_dataset
+    
+    dataset = load_dataset(
+        "json",
+        data_files={
+            "train": config["dataset"][0]["path"].replace("../", ""),
+            "test": config["dataset"][0]["path"].replace("../", "").replace("combined_20k.jsonl", "test.jsonl"),
+        },
+    )
+    
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir=config["train_params"]["output_dir"],
+        num_train_epochs=config["train_params"]["num_train_epochs"],
+        per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
+        gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
+        learning_rate=config["train_params"]["learning_rate"],
+        lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
+        weight_decay=config["train_params"]["weight_decay"],
+        warmup_ratio=config["train_params"]["warmup_ratio"],
+        max_seq_length=config["train_params"]["max_seq_length"],
+        logging_steps=config["train_params"]["logging_steps"],
+        save_steps=config["train_params"]["save_steps"],
+        save_total_limit=config["train_params"]["save_total_limit"],
+        evaluation_strategy=config.get("eval_strategy", "steps"),
+        eval_steps=config.get("eval_steps", 100),
+        mixed_precision=config.get("mixed_precision", "bf16"),
+        gradient_checkpointing=config.get("gradient_checkpointing", True),
+    )
+    
+    # SFT Trainer
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["test"],
+        args=training_args,
+        max_seq_length=config["train_params"]["max_seq_length"],
+    )
+    
+    # Train
+    print("Starting training...")
+    trainer.train()
+    
+    # Save
+    trainer.save_model(config["train_params"]["output_dir"])
+    tokenizer.save_pretrained(config["train_params"]["output_dir"])
+    
+    print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train LoRA adapter")
+    parser.add_argument("--config", type=str, default="configs/llama2-7b-lora.yaml", 
+                       help="Training configuration file")
+    args = parser.parse_args()
+    
+    train(args.config)
+
+
+if __name__ == "__main__":
+    main()