"""Modal Tier 2 — Q3B-32B + LoRA r=32 on canonical-v1 multi-segment corpus.

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

Schema target: {"segments": [{"intent","slots","text"}, ...], "abstain_reason": null|str}
51-intent enum (50 canonical + "unknown" for abstention)

Cost: ~$5-8 on A100-80GB for ~600 rows × 5 epochs = ~3000 steps.
"""
import modal

GPU = "A100-80GB"
TIMEOUT_HR = 3
MEMORY_GB = 100

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-v9", image=image)

HF_DATASET = "kinglyai/atc-parser-canonical-v1"
HF_OUTPUT = "kinglyai/qwen3-32b-atc-parser-v9"
BASE_MODEL = "Qwen/Qwen3-32B"  # Qwen3-32B is instruct-tuned (no separate -Instruct-2507 variant exists at 32B)
RUN_NAME = "qwen3-32b-canonical-v1-multisegment"
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 V9 starting · GPU: {GPU} · base: {BASE_MODEL} · dataset: {HF_DATASET}")
    # trackio.init removed — static Space conflicts with private repo flag

    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,
    )

    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=8,  # small dataset, more epochs
        per_device_train_batch_size=2,
        per_device_eval_batch_size=2,
        gradient_accumulation_steps=8,  # effective bs=16
        learning_rate=1.5e-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=20,
        report_to="none",  # trackio Space privacy conflict; tail logs via `modal logs` instead
        run_name=RUN_NAME,
        push_to_hub=True,
        hub_model_id=HF_OUTPUT,
        hub_strategy="every_save",
        # hub_private_repo removed — incompatible with static Trackio space
        max_length=2048,  # longer for compound transmissions
        gradient_checkpointing=True,
    )

    trainer = SFTTrainer(
        model=model, processing_class=tok,  # newer TRL API (was `tokenizer`)
        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}")


@app.local_entrypoint()
def main():
    print(f"submitting Modal V9 · {GPU} · {TIMEOUT_HR}h · {BASE_MODEL}")
    train.remote()
    print(f"complete. Adapter: https://huggingface.co/{HF_OUTPUT}")