"""Stage 3: full fine-tune of Qwen2.5-0.5B-Instruct into the Shaktimaan router (loss on the route only)."""
import os
from datasets import load_from_disk
from trl import SFTConfig, SFTTrainer

BASE = os.environ.get("BASE", "Qwen/Qwen2.5-0.5B-Instruct")
ds = load_from_disk("data/shaktimaan_routing").map(
    lambda r: {"prompt": r["messages"][:2], "completion": r["messages"][2:]}).select_columns(["prompt", "completion"])

cfg = SFTConfig(
    output_dir="ckpt", num_train_epochs=3, per_device_train_batch_size=8, gradient_accumulation_steps=4,
    learning_rate=2e-5, lr_scheduler_type="cosine", warmup_steps=20, weight_decay=0.0, bf16=True,
    max_length=3072, completion_only_loss=True, logging_steps=10, eval_strategy="epoch", save_strategy="epoch",
    save_total_limit=1, report_to=[], gradient_checkpointing=True, dataloader_num_workers=4,
)
trainer = SFTTrainer(model=BASE, args=cfg, train_dataset=ds["train"], eval_dataset=ds["test"])
trainer.train()
trainer.save_model("shaktimaan-0.5b")
trainer.processing_class.save_pretrained("shaktimaan-0.5b")