# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "unsloth",
#     "trl>=0.12.0",
#     "transformers>=4.45.0",
#     "datasets>=3.0.0",
#     "peft>=0.13.0",
#     "accelerate>=0.34.0",
#     "trackio",
#     "huggingface_hub>=0.25.0",
# ]
# ///
"""H1 cloud SFT — Qwen3-32B-Instruct-2507 + Unsloth + LoRA r=16 on ATC parser dataset.

Submit via:
    hf jobs uv run --flavor a100-large --timeout 3h --secrets HF_TOKEN \\
      "https://huggingface.co/<USER>/atc-parser-train-script/resolve/main/train_h1_cloud.py"

Goal: break M3-local 65% intent_em ceiling. Target 85%+.

Inputs (env, HF Jobs auto-injects HF_TOKEN):
    HF_DATASET_REPO    default kinglyai/atc-parser-spike-v0
    HF_OUTPUT_REPO     default kinglyai/qwen3-32b-atc-parser-v1
    HF_TRACKIO_PROJECT default atc-parser
    RUN_NAME           default qwen3-32b-r16-h1-v1
"""
from __future__ import annotations
import os

from datasets import load_dataset
from peft import LoraConfig
from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig
import trackio


BASE_MODEL = os.environ.get("HF_BASE_MODEL", "Qwen/Qwen3-32B-Instruct-2507")
DATASET_REPO = os.environ.get("HF_DATASET_REPO", "kinglyai/atc-parser-spike-v0")
OUTPUT_REPO = os.environ.get("HF_OUTPUT_REPO", "kinglyai/qwen3-32b-atc-parser-v1")
TRACKIO_PROJECT = os.environ.get("HF_TRACKIO_PROJECT", "atc-parser")
RUN_NAME = os.environ.get("RUN_NAME", "qwen3-32b-r16-h1-v1")


def main() -> None:
    trackio.init(project=TRACKIO_PROJECT, name=RUN_NAME)

    # Load with Unsloth (60% less VRAM, 2× faster)
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=BASE_MODEL,
        max_seq_length=1024,
        dtype=None,  # auto bf16
        load_in_4bit=False,  # use 8-bit or full bf16 on a100-80GB; LoRA only adapter
    )

    # LoRA r=16 — 2× capacity vs M3 spike's r=8
    model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        bias="none",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        use_gradient_checkpointing="unsloth",
        random_state=42,
    )

    # Load dataset (3 JSONL files in HF dataset repo)
    ds = load_dataset(DATASET_REPO, data_files={
        "train": "train.jsonl",
        "valid": "valid.jsonl",
        "test": "test.jsonl",
    })

    cfg = SFTConfig(
        output_dir="/tmp/sft_out",
        num_train_epochs=2,
        per_device_train_batch_size=4,
        per_device_eval_batch_size=4,
        gradient_accumulation_steps=4,
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        warmup_ratio=0.05,
        bf16=True,
        eval_strategy="steps",
        eval_steps=100,
        save_strategy="steps",
        save_steps=200,
        save_total_limit=3,
        logging_steps=10,
        report_to="trackio",
        run_name=RUN_NAME,
        push_to_hub=True,
        hub_model_id=OUTPUT_REPO,
        hub_strategy="every_save",
        hub_private_repo=True,
        save_safetensors=True,
        max_length=1024,
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=ds["train"],
        eval_dataset=ds["valid"],
        args=cfg,
    )

    trainer.train()
    trainer.save_model()
    trainer.push_to_hub()

    trackio.finish()


if __name__ == "__main__":
    main()