# /// script
# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate", "torch"]
# ///
"""SFT on compile-verified C++ curriculum. Designed for Hugging Face Jobs (uv)."""

import os

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

DATASET_ID = os.environ.get("DATASET_ID", "gonzalolinares/cpp-compiler-curriculum")
MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "gonzalolinares/qwen25-1.5b-cpp-sft")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "qwen25-1.5b-cpp-sft")


def main() -> None:
    ds = load_dataset(DATASET_ID, split="train")
    # Keep messages format for SFTTrainer chat templates
    if "messages" not in ds.column_names:
        raise SystemExit(f"Dataset must have 'messages'; got {ds.column_names}")

    split = ds.train_test_split(test_size=0.1, seed=42)

    trainer = SFTTrainer(
        model=MODEL_ID,
        train_dataset=split["train"],
        eval_dataset=split["test"],
        peft_config=LoraConfig(
            r=16,
            lora_alpha=32,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        ),
        args=SFTConfig(
            output_dir=OUTPUT_DIR,
            num_train_epochs=4,
            per_device_train_batch_size=2,
            per_device_eval_batch_size=2,
            gradient_accumulation_steps=8,
            learning_rate=2e-4,
            logging_steps=10,
            eval_strategy="steps",
            eval_steps=40,
            save_strategy="epoch",
            save_total_limit=1,
            max_length=1024,
            bf16=True,
            push_to_hub=False,  # push once at end (avoid Trackio parquet bug mid-save)
            hub_model_id=HUB_MODEL_ID,
            report_to="none",
        ),
    )
    trainer.train()
    trainer.model.push_to_hub(HUB_MODEL_ID, private=False)
    trainer.processing_class.push_to_hub(HUB_MODEL_ID, private=False)
    print(f"Pushed to {HUB_MODEL_ID}")


if __name__ == "__main__":
    main()