""" CodeArena — Fine-Tuning Script Fine-tunes LLaMA / Gemma models on the XCoder-80K dataset using Unsloth. Supported base models (pick one): - unsloth/Llama-3.2-3B-Instruct (recommended for code tasks) - unsloth/gemma-3-4b-it - unsloth/gemma-3-1b-it - unsloth/llava-1.5-7b-hf (multimodal — skip for code-only) Usage: python finetune.py --model llama3 --output ./finetuned_model After training: The model is saved to ./finetuned_model (GGUF + LoRA adapter) Pull into Ollama: ollama create codearena -f ./finetuned_model/Modelfile """ import argparse import os import sys # ─── Check GPU ──────────────────────────────────────────────────────────────── def check_gpu(): try: import torch if not torch.cuda.is_available(): print("⚠ WARNING: No CUDA GPU found. Fine-tuning will be very slow on CPU.") print(" Recommended: Use Google Colab (free T4 GPU) or Kaggle Notebooks.") else: print(f"✓ GPU: {torch.cuda.get_device_name(0)}") except ImportError: print("✗ PyTorch not installed. Run: pip install torch torchvision") sys.exit(1) # ─── Model Registry ─────────────────────────────────────────────────────────── MODELS = { "llama3": "unsloth/Llama-3.2-3B-Instruct", "llama3_8b": "unsloth/Meta-Llama-3.1-8B-Instruct", "gemma4b": "unsloth/gemma-3-4b-it", "gemma1b": "unsloth/gemma-3-1b-it", } # ─── Dataset Formatter ──────────────────────────────────────────────────────── def format_xcoder_example(example: dict) -> dict: """ Convert XCoder-80K format to chat-style instruction tuning. XCoder format: { instruction, input, output, system? } """ instruction = example.get("instruction", "") inp = example.get("input", "") output = example.get("output", "") system = example.get("system", "You are an expert Python debugging assistant.") user_msg = instruction if inp: user_msg += f"\n\n```python\n{inp}\n```" return { "messages": [ {"role": "system", "content": system}, {"role": "user", "content": user_msg}, {"role": "assistant", "content": output}, ] } def load_xcoder_dataset(max_samples: int = 5000): """Load and format the XCoder-80K dataset.""" from datasets import load_dataset print("📦 Loading banksy235/XCoder-80K dataset...") ds = load_dataset("banksy235/XCoder-80K", split="train") # Filter for code-related examples def is_code_task(ex): text = (ex.get("instruction", "") + ex.get("input", "") + ex.get("output", "")).lower() return any(kw in text for kw in ["python", "def ", "function", "bug", "error", "fix", "optimize", "algorithm"]) print(f" Total examples: {len(ds)}") ds = ds.filter(is_code_task) print(f" Code-related: {len(ds)}") if max_samples and len(ds) > max_samples: ds = ds.select(range(max_samples)) print(f" Using {max_samples} samples for training") ds = ds.map(format_xcoder_example, remove_columns=ds.column_names) return ds # ─── Main Fine-Tuning ───────────────────────────────────────────────────────── def run_finetune(model_key: str, output_dir: str, max_samples: int, epochs: int, batch_size: int): check_gpu() try: from unsloth import FastLanguageModel from unsloth.chat_templates import get_chat_template, train_on_responses_only from trl import SFTTrainer from transformers import TrainingArguments, DataCollatorForSeq2Seq except ImportError: print("\n✗ Unsloth not installed. Install it first:") print(" pip install unsloth trl transformers accelerate bitsandbytes datasets") sys.exit(1) model_id = MODELS.get(model_key, MODELS["llama3"]) print(f"\n🚀 Loading model: {model_id}") # Load model with 4-bit quantization (fits in ~6GB VRAM) model, tokenizer = FastLanguageModel.from_pretrained( model_name=model_id, max_seq_length=2048, dtype=None, # Auto-detect (bfloat16 on modern GPUs) load_in_4bit=True, # QLoRA — use less VRAM ) # Apply LoRA adapters (PEFT — only train ~1% of params) model = FastLanguageModel.get_peft_model( model, r=16, # LoRA rank target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_alpha=16, lora_dropout=0, bias="none", use_gradient_checkpointing="unsloth", random_state=42, ) # Apply chat template tokenizer = get_chat_template(tokenizer, chat_template="llama-3") def apply_template(examples): texts = tokenizer.apply_chat_template( examples["messages"], tokenize=False, add_generation_prompt=False, ) return {"text": texts} # Load dataset dataset = load_xcoder_dataset(max_samples) dataset = dataset.map(apply_template, batched=True, remove_columns=["messages"]) print(f"\n📊 Training on {len(dataset)} examples for {epochs} epoch(s)") trainer = SFTTrainer( model=model, tokenizer=tokenizer, train_dataset=dataset, dataset_text_field="text", max_seq_length=2048, data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True), dataset_num_proc=2, packing=False, args=TrainingArguments( per_device_train_batch_size=batch_size, gradient_accumulation_steps=4, warmup_steps=10, num_train_epochs=epochs, learning_rate=2e-4, fp16=False, bf16=True, logging_steps=10, optim="adamw_8bit", weight_decay=0.01, lr_scheduler_type="cosine", seed=42, output_dir=output_dir, save_strategy="epoch", report_to="none", ), ) # Only train on assistant responses, not user prompts trainer = train_on_responses_only( trainer, instruction_part="<|start_header_id|>user<|end_header_id|>\n\n", response_part="<|start_header_id|>assistant<|end_header_id|>\n\n", ) print("\n🔥 Starting training...") trainer_stats = trainer.train() print(f"\n✓ Training complete! Stats: {trainer_stats.metrics}") # Save model print(f"\n💾 Saving LoRA adapter to {output_dir}/lora_model") model.save_pretrained(f"{output_dir}/lora_model") tokenizer.save_pretrained(f"{output_dir}/lora_model") # Export to GGUF for Ollama print("\n📦 Exporting to GGUF (Q4_K_M quantization)...") try: model.save_pretrained_gguf( f"{output_dir}/gguf_model", tokenizer, quantization_method="q4_k_m", ) # Write Modelfile for Ollama modelfile = f"""FROM {output_dir}/gguf_model/model-q4_k_m.gguf SYSTEM You are CodeArena, an expert Python debugging and code optimization agent. You fix bugs, optimize algorithms, and improve code quality. Always return ONLY the fixed code without explanation unless asked. PARAMETER temperature 0.1 PARAMETER num_ctx 2048 """ with open(f"{output_dir}/Modelfile", "w") as f: f.write(modelfile) print(f""" ╔═══════════════════════════════════════════════════════╗ ║ ✓ Fine-tuning complete! ║ ║ ║ ║ To use in CodeArena: ║ ║ 1. Install the model into Ollama: ║ ║ ollama create codearena -f {output_dir}/Modelfile ║ ║ 2. Set model name to "codearena" in the dashboard ║ ╚═══════════════════════════════════════════════════════╝ """) except Exception as e: print(f"⚠ GGUF export failed: {e}") print(" LoRA adapter saved. You can merge it manually later.") # ─── CLI ───────────────────────────────────────────────────────────────────── if __name__ == "__main__": parser = argparse.ArgumentParser(description="Fine-tune a model on XCoder-80K for CodeArena") parser.add_argument("--model", choices=list(MODELS.keys()), default="llama3", help="Base model to fine-tune") parser.add_argument("--output", default="./finetuned_model", help="Output directory for the fine-tuned model") parser.add_argument("--samples", type=int, default=5000, help="Max training samples from XCoder-80K (default: 5000)") parser.add_argument("--epochs", type=int, default=1, help="Number of training epochs (default: 1)") parser.add_argument("--batch-size", type=int, default=2, help="Batch size per device (default: 2)") args = parser.parse_args() run_finetune( model_key=args.model, output_dir=args.output, max_samples=args.samples, epochs=args.epochs, batch_size=args.batch_size, )