# /// script
# dependencies = [
#     "trl>=0.12.0",
#     "peft>=0.7.0",
#     "bitsandbytes",
#     "trackio",
#     "datasets",
#     "transformers>=4.40.0",
#     "torch>=2.0.0",
#     "torchvision",
#     "huggingface-hub",
#     "accelerate",
# ]
# ///

import os
import torch
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import trackio

hf_token = os.getenv("HF_TOKEN")
if hf_token:
    os.environ.setdefault("HF_HUB_TOKEN", hf_token)

DATASET = os.getenv("DATASET", "oExz/reasoning-distill-claude-opus-4-7-max-qwen-format")
BASE_MODEL = os.getenv("BASE_MODEL", "Qwen/Qwen3.5-4B")
OUTPUT_REPO = os.getenv("OUTPUT_REPO", "oExz/qwen3.5-4b-reasoning-lora")

print(f"Loading dataset {DATASET}...")
dataset = load_dataset(DATASET)
train_ds = dataset["train"]
eval_ds = dataset["eval"]
print(f"Train: {len(train_ds)}, Eval: {len(eval_ds)}")

print(f"Loading base model {BASE_MODEL} with 4-bit quantization...")

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    trust_remote_code=True,
    quantization_config=bnb_config,
    device_map="auto",
)

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

trackio.init(
    project="qwen-reasoning-finetune",
    space_id="oExz/trackio",
)

trainer = SFTTrainer(
    model=model,
    train_dataset=train_ds,
    eval_dataset=eval_ds,
    peft_config=peft_config,
    args=SFTConfig(
        output_dir="/tmp/qwen-lora",
        push_to_hub=True,
        hub_model_id=OUTPUT_REPO,
        hub_strategy="every_save",
        num_train_epochs=2,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        warmup_ratio=0.1,
        eval_strategy="steps",
        eval_steps=50,
        save_strategy="steps",
        save_steps=50,
        logging_steps=10,
        report_to="trackio",
        run_name="qwen3.5-4b-reasoning-qlora",
        max_length=1024,
        bf16=True,
        gradient_checkpointing=True,
    ),
)

print("Starting training...")
trainer.train()
print("Training complete, pushing to hub...")
trainer.push_to_hub()
print("Done!")
trackio.finish()