"""Modal-deployed SFT training for ATC parser.

Run: modal run poc/llm-finetune/training/train_modal.py

Pulls dataset from HF Hub (kinglyai/atc-parser-spike-v0), trains Qwen3-32B
with LoRA r=32, pushes adapter back to HF Hub (kinglyai/qwen3-32b-atc-parser-v1).

Requires:
  - Modal account + token (modal token new)
  - HF_TOKEN secret created via: modal secret create huggingface HF_TOKEN=<token>

Cost: ~$3-7 on A100-80GB for 1500 iters.
"""
import modal

GPU = "A100-80GB"  # or "H100" if available; A10G-24GB also works for r=16 LoRA on smaller models
TIMEOUT_HR = 3
MEMORY_GB = 80

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "torch>=2.4.0",
        "transformers>=4.45.0",
        "trl>=0.12.0",
        "peft>=0.13.0",
        "accelerate>=0.34.0",
        "datasets>=3.0.0",
        "bitsandbytes",
        "trackio",
        "huggingface_hub>=0.25.0",
        "sentencepiece",
        "protobuf",
    )
)

app = modal.App("atc-parser-sft", image=image)

HF_DATASET = "kinglyai/atc-parser-canonical-v0"  # canonical conventions, FMM-relevant slots
HF_DATASET_LARGE = "kinglyai/atc-parser-spike-v0"  # broad UPPERCASE corpus (19k rows)
HF_OUTPUT = "kinglyai/qwen3-32b-atc-parser-v1"
BASE_MODEL = "Qwen/Qwen3-32B-Instruct-2507"  # or Qwen2.5-14B-Instruct for cheaper test
RUN_NAME = "qwen3-32b-canonical-v1"
TRACKIO_PROJECT = "atc-parser"


@app.function(
    gpu=GPU,
    timeout=TIMEOUT_HR * 3600,
    secrets=[modal.Secret.from_name("huggingface")],
    memory=MEMORY_GB * 1024,
)
def train():
    import os
    import torch
    from datasets import load_dataset
    from peft import LoraConfig
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from trl import SFTConfig, SFTTrainer
    import trackio

    print(f"Modal container starting · GPU: {GPU} · base: {BASE_MODEL}")

    trackio.init(project=TRACKIO_PROJECT, name=RUN_NAME)

    print(f"loading tokenizer + model from {BASE_MODEL}")
    tok = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )

    print(f"loading dataset from {HF_DATASET}")
    ds = load_dataset(HF_DATASET, data_files={
        "train": "train.jsonl",
        "valid": "valid.jsonl",
        "test": "test.jsonl",
    })

    lora = LoraConfig(
        r=32,
        lora_alpha=64,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        bias="none",
        lora_dropout=0.05,
        task_type="CAUSAL_LM",
    )

    cfg = SFTConfig(
        output_dir="/tmp/sft_out",
        num_train_epochs=10,  # canonical dataset is small (~417 rows); more epochs to learn conventions
        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=200,
        save_strategy="steps",
        save_steps=400,
        save_total_limit=3,
        logging_steps=20,
        report_to="trackio",
        run_name=RUN_NAME,
        push_to_hub=True,
        hub_model_id=HF_OUTPUT,
        hub_strategy="every_save",
        hub_private_repo=True,
        save_safetensors=True,
        max_length=1024,
        gradient_checkpointing=True,
    )

    print(f"starting training · {len(ds['train'])} train rows · {len(ds['valid'])} valid")
    trainer = SFTTrainer(
        model=model,
        tokenizer=tok,
        train_dataset=ds["train"],
        eval_dataset=ds["valid"],
        peft_config=lora,
        args=cfg,
    )

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

    print(f"DONE · adapter pushed to {HF_OUTPUT}")
    trackio.finish()


@app.local_entrypoint()
def main():
    """Launch from local: `modal run poc/llm-finetune/training/train_modal.py`"""
    print(f"submitting Modal job · {GPU} · {TIMEOUT_HR}h timeout · base={BASE_MODEL}")
    train.remote()
    print("Modal job complete. Check HF Hub for adapter:")
    print(f"  https://huggingface.co/{HF_OUTPUT}")