"""
Stage 1 SFT Training Script: Fine-tune Qwen2.5-Coder-7B-Instruct for
Python code generation + tool-calling + RAG-aware generation.

Based on:
  - Qwen2.5-Coder Technical Report (arxiv:2409.12186): coarse-to-fine SFT recipe
  - ToolACE (arxiv:2409.00920): LoRA r=16, alpha=32, LR=1e-4, cosine, 3 epochs
  - Gorilla (arxiv:2305.15334): retriever-aware training pattern
  - TRL v1.2.0 SFTTrainer with ChatML messages format

Reference implementation: TRL SFT docs (https://huggingface.co/docs/trl/sft_trainer)
Dataset format: conversational ChatML with "messages" column

Usage:
  # Launch via hf_jobs (recommended)
  # Or run locally:
  python train_sft.py
"""

import os
import torch
import trackio
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

# ============================================================================
# Configuration — EDIT THESE FOR YOUR USE CASE
# ============================================================================
BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
DATASET_REPO = "your-username/code-toolcall-sft-data"  # From prepare_data.py
OUTPUT_DIR = "./qwen25-coder-7b-code-toolcall"
HUB_MODEL_ID = "your-username/qwen25-coder-7b-code-toolcall"
HF_TOKEN = os.environ.get("HF_TOKEN")

# Training hyperparameters (from ToolACE + Qwen2.5-Coder papers)
LEARNING_RATE = 1e-4          # LoRA LR (10x FFT rate)
NUM_EPOCHS = 2                # 2 epochs (Magicoder recipe)
BATCH_SIZE = 2                # Per-device batch size
GRAD_ACCUM = 8                # Effective batch = 2 * 8 = 16
MAX_SEQ_LENGTH = 8192         # 8K context for tool-calling
WARMUP_RATIO = 0.1
LORA_R = 32                   # Rank (32-64 for code tasks)
LORA_ALPHA = 64               # Alpha = 2 * r
LORA_DROPOUT = 0.05

# ============================================================================
# Trackio Monitoring
# ============================================================================
trackio.init(name="code-toolcall-sft", project="code-llm-finetuning")

# ============================================================================
# Load Dataset
# ============================================================================
print(f"Loading dataset: {DATASET_REPO}")
dataset = load_dataset(DATASET_REPO, split="train")
print(f"Dataset size: {len(dataset)} examples")

split = dataset.train_test_split(test_size=0.02, seed=42)
train_dataset = split["train"]
eval_dataset = split["test"]
print(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")

# ============================================================================
# LoRA Configuration
# ============================================================================
peft_config = LoraConfig(
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    lora_dropout=LORA_DROPOUT,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
)

# ============================================================================
# SFT Training Configuration
# ============================================================================
training_args = SFTConfig(
    output_dir=OUTPUT_DIR,
    hub_model_id=HUB_MODEL_ID,
    push_to_hub=True,
    hub_token=HF_TOKEN,
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=LEARNING_RATE,
    lr_scheduler_type="cosine",
    warmup_ratio=WARMUP_RATIO,
    weight_decay=0.01,
    max_grad_norm=1.0,
    bf16=True,
    gradient_checkpointing=True,
    max_seq_length=MAX_SEQ_LENGTH,
    assistant_only_loss=True,  # Only train on assistant responses
    packing=False,
    logging_strategy="steps",
    logging_steps=10,
    logging_first_step=True,
    disable_tqdm=True,
    eval_strategy="steps",
    eval_steps=200,
    per_device_eval_batch_size=2,
    save_strategy="steps",
    save_steps=500,
    save_total_limit=3,
    load_best_model_at_end=True,
    seed=42,
    report_to="none",
)

# ============================================================================
# Train
# ============================================================================
print(f"Loading model: {BASE_MODEL}")
trainer = SFTTrainer(
    model=BASE_MODEL,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
)

print("Starting training...")
result = trainer.train()
print(f"Training complete! Loss: {result.training_loss:.4f}")

trainer.push_to_hub(commit_message="Final model after SFT training")
print(f"Model pushed to: https://huggingface.co/{HUB_MODEL_ID}")

trackio.log({"final_loss": result.training_loss})
trackio.finish()