"""Qwen3-8B coding agent LoRA — trained on opencode tool-use + coding patterns.

Usage (on RunPod A100/H100):
    python train_qwen3_8b.py
"""

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

# --- Config ---
MODEL = "unsloth/Qwen3-8B-unsloth-bnb-4bit"
MAX_SEQ = 4096
RANK = 16
ALPHA = 16
DATA = "./bt7274_v3.jsonl"
OUT = "./qwen3-8b-lora-v1"
EPOCHS = 3
BATCH = 4        # A100/H100 can handle bigger batches
GRAD_ACCUM = 4   # effective batch = 16
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}/")