"""Stage 3: full fine-tune of Qwen2.5-0.5B-Instruct into Jedi (loss on the one-word verdict 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/jedi_judge").select_columns(["prompt", "completion"])

cfg = SFTConfig(
    output_dir="ckpt", num_train_epochs=2, per_device_train_batch_size=8, gradient_accumulation_steps=4,
    learning_rate=2e-5, lr_scheduler_type="cosine", warmup_steps=30, weight_decay=0.0, bf16=True,
    max_length=2048, completion_only_loss=True, logging_steps=20, 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["val"])
trainer.train()
trainer.save_model("jedi-0.5b")
trainer.processing_class.save_pretrained("jedi-0.5b")