"""BT-7274 LoRA v2 — 500 examples with tool calls.

Usage:
    # just train
    cd ~/Projects/lora
    python train_v2.py
"""

from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset

# --- Config ---
MODEL = "unsloth/Qwen2.5-7B-Instruct-bnb-4bit"
MAX_SEQ = 4096  # longer for tool call chains
RANK = 16
ALPHA = 16
DATA = "./bt7274_v2.jsonl"
OUT = "./bt7274-lora-v2"
EPOCHS = 3       # 5x more data, fewer epochs
BATCH = 1        # larger sequences need smaller batch
GRAD_ACCUM = 8   # effective batch = 8
LR = 2e-4

# --- Load model (4-bit quantized) ---
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL,
    max_seq_length=MAX_SEQ,
    load_in_4bit=True,
    dtype=None,
)

# --- LoRA adapter ---
model = FastLanguageModel.get_peft_model(
    model,
    r=RANK,
    lora_alpha=ALPHA,
    lora_dropout=0,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=42,
)

# --- Dataset ---
ds = load_dataset("json", data_files=DATA, split="train")

def to_chatml(ex):
    text = tokenizer.apply_chat_template(
        ex["messages"], tokenize=False, add_generation_prompt=False
    )
    return {"text": text}

ds = ds.map(to_chatml)

steps = (len(ds) * EPOCHS) // (BATCH * GRAD_ACCUM)
print(f"Dataset: {len(ds)} examples")
print(f"Epochs: {EPOCHS}, effective batch: {BATCH * GRAD_ACCUM}")
print(f"Estimated steps: {steps}")
print(f"LoRA: r={RANK}, alpha={ALPHA}")
print(f"Max seq: {MAX_SEQ}")

# --- Train ---
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=ds,
    args=TrainingArguments(
        output_dir=OUT,
        per_device_train_batch_size=BATCH,
        gradient_accumulation_steps=GRAD_ACCUM,
        num_train_epochs=EPOCHS,
        learning_rate=LR,
        bf16=True,
        logging_steps=5,
        save_steps=50,
        save_total_limit=2,
        warmup_steps=10,
        optim="adamw_8bit",
        seed=42,
        report_to="none",
    ),
    dataset_text_field="text",
    max_seq_length=MAX_SEQ,
    packing=False,
)

trainer.train()

# --- Save LoRA adapter ---
model.save_pretrained(OUT)
tokenizer.save_pretrained(OUT)
print(f"\nSaved LoRA adapter to {OUT}/")