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

import os

# Import unsloth FIRST before anything else
from unsloth import FastLanguageModel
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from transformers import AutoTokenizer
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-text")
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 Unsloth...")
model, _ = FastLanguageModel.from_pretrained(
    model_name=BASE_MODEL,
    max_seq_length=1024,
    dtype=None,
    load_in_4bit=True,
    trust_remote_code=True,
)

# Load tokenizer directly
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
# Explicitly set special tokens to avoid TokenizersBackend defaults
tokenizer.eos_token = "<|im_end|>"
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids(tokenizer.eos_token)

model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=3407,
)

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

trainer = SFTTrainer(
    model=model,
    processing_class=tokenizer,
    train_dataset=train_ds,
    eval_dataset=eval_ds,
    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=2,
        gradient_accumulation_steps=4,
        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,
        dataset_text_field="text",
        packing=False,
        eos_token="<|im_end|>",
    ),
)

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