# finetune/train_lora.py
"""Stage 2 — LoRA fine-tune MiniCPM-V 4.5 on the distilled Bengali dataset.

Uses **ms-SWIFT** (modelscope/ms-swift) rather than OpenBMB's official finetune
scripts. The official path hit a hard dependency wall: its requirements pin torch
2.1.2, but MiniCPM-V 4.5's remote code needs torch>=2.4 (`torch.library.custom_op`),
and the repo's finetune.py targets older MiniCPM releases (4.0/2.6/2.5), not 4.5.
SWIFT ships a maintained recipe for this exact model — `model_type=minicpmv4_5`,
`template=minicpmv4_5`, deps `timm, transformers>=4.36, decord` — so the version
matrix is solved for us.

Vision encoder frozen (`--freeze_vit true`); LoRA on the LLM self-attention layers
only (`q/k/v/o_proj`) — that's the weak-Bengali part. ViT and aligner stay frozen.

Inputs  (Modal Volume `rupkotha-finetune`):  /train.json  +  /labelset/*.jpg
        (train.json 'image' fields are container-absolute: /data/labelset/<name>)
        Converted in-container to SWIFT's messages/images JSONL schema.
Output  (same volume):  /out/lora-bengali/   (PEFT adapter — feed to finetune/merge.py)

Run a SHORT validation first to surface integration errors cheaply:
    uv run modal run finetune/train_lora.py --max-steps 4
Then the full run:
    uv run modal run finetune/train_lora.py
"""

import modal

from core.model_config import STUDENT_BASE_REPO  # openbmb/MiniCPM-V-4_5

app = modal.App("rupkotha-finetune")

_vol = modal.Volume.from_name("rupkotha-finetune", create_if_missing=True)
_hf = modal.Volume.from_name("rupkotha-hf", create_if_missing=True)

# ms-SWIFT brings its own pinned transformers/peft/accelerate; we only add the
# MiniCPM-V 4.5 model extras (timm, decord) and force torch>=2.4 for its remote
# code. No deepspeed/nvcc needed — a single A100-80GB fits an 8B base + LoRA, and
# attn defaults to sdpa (no flash-attn build), so debian_slim + pip wheels suffice.
# Pin ms-swift 3.12.6 (last 3.x): it lists model_type `minicpmv4_5` AND pins
# transformers>=4.33,<4.58. The 4.x line pulls transformers 5.x, whose remote-code
# loader follows the HF-cache symlink into blobs/ and fails to resolve MiniCPM-V's
# relative imports (modeling_navit_siglip.py). 4.x transformers loads it cleanly.
_train_image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install("git")
    .pip_install(
        "ms-swift==3.12.6",
        "torch>=2.4",
        "timm",
        "decord",
        "pillow",
        "sentencepiece",
    )
    # train_lora.py imports core.model_config at module load; `modal run <file>`
    # doesn't auto-mount the project, so make `core` importable in the container.
    .add_local_python_source("core")
)


@app.function(
    gpu="A100-80GB",  # headroom for an 8B base + LoRA on one GPU
    image=_train_image,
    volumes={"/data": _vol, "/root/.cache/huggingface": _hf},
    secrets=[modal.Secret.from_name("algaeguard-secrets")],  # HF_TOKEN for the base
    timeout=60 * 60 * 6,
)
def train(max_steps: int = 0) -> str:
    import glob
    import json
    import os
    import shutil
    import subprocess

    # SWIFT defaults to ModelScope; force HuggingFace so it pulls openbmb/MiniCPM-V-4_5
    # (and reuses the HF_TOKEN from the algaeguard-secrets secret).
    os.environ["USE_HF"] = "1"
    # Mirror the original max_slice_nums=9 (SWIFT reads this env for MiniCPM-V).
    os.environ.setdefault("MAX_SLICE_NUMS", "9")

    # Repair the shared HF cache: the base model was cached weights-first (vLLM),
    # leaving the MiniCPM-V remote-code .py files as dangling symlinks (e.g.
    # modeling_navit_siglip.py). Refetch the small code/config files cleanly so
    # trust_remote_code loads; the .safetensors weights stay cached untouched.
    from huggingface_hub import snapshot_download
    snapshot_download(
        STUDENT_BASE_REPO,
        allow_patterns=["*.py", "*.json", "*.txt", "*.model", "tokenizer*"],
        force_download=True,
    )

    # ── Convert MiniCPM conversations format → SWIFT messages/images JSONL ──
    # Source rows already carry role/content turns with a leading <image> in the
    # user content and a container-absolute image path. SWIFT wants `messages` +
    # an `images` list (one path per <image> placeholder).
    src = "/data/train.json"  # volume copy: image paths already /data/labelset/<name>
    swift_data = "/data/train_swift.jsonl"
    rows = json.load(open(src))
    with open(swift_data, "w") as f:
        for r in rows:
            f.write(json.dumps(
                {"messages": r["conversations"], "images": [r["image"]]},
                ensure_ascii=False,
            ) + "\n")
    print(f"Converted {len(rows)} examples → {swift_data}")

    adapter_dir = "/data/out/lora-bengali"     # canonical path merge.py expects
    swift_out = "/data/out/swift-runs"         # SWIFT writes vX-<ts>/checkpoint-N here

    # Mirror finetune_lora.sh intent via SWIFT's CLI: vision frozen, LoRA r=16 on
    # the LLM self-attention projections only.
    cmd = [
        "swift", "sft",
        "--model", STUDENT_BASE_REPO,
        "--model_type", "minicpmv4_5",
        "--train_type", "lora",
        "--dataset", swift_data,
        "--freeze_vit", "true",
        "--target_modules", "q_proj", "k_proj", "v_proj", "o_proj",
        "--lora_rank", "16", "--lora_alpha", "32", "--lora_dropout", "0.05",
        "--torch_dtype", "bfloat16",
        # 4096, not 2048: the verbose Bengali prompt + MiniCPM image slices
        # (max_slice_nums=9) push some rows to ~2083 tokens. At 2048 SWIFT raises
        # MaxLengthError per over-long row and silently drops it from training;
        # 4096 keeps all 389 samples (peak mem was only ~28 GiB of 80).
        "--max_length", "4096",
        "--per_device_train_batch_size", "1",
        "--gradient_accumulation_steps", "8",
        "--learning_rate", "1e-4",
        "--gradient_checkpointing", "true",
        "--save_strategy", "steps", "--save_steps", "200", "--save_total_limit", "2",
        "--logging_steps", "5", "--report_to", "none",
        "--dataloader_num_workers", "4",
        "--output_dir", swift_out,
    ]
    # Short validation run vs full run.
    if max_steps and max_steps > 0:
        cmd += ["--max_steps", str(max_steps)]
    else:
        cmd += ["--num_train_epochs", "3"]

    print("Running:", " ".join(cmd))
    subprocess.run(cmd, check=True)

    # SWIFT nests output under output_dir/<version>/checkpoint-<step>/. Find the
    # latest dir that actually holds a PEFT adapter and copy it to the canonical
    # path so finetune/merge.py (which loads /data/out/lora-bengali) works unchanged.
    adapters = glob.glob(f"{swift_out}/**/adapter_config.json", recursive=True)
    if not adapters:
        raise RuntimeError(f"No adapter produced under {swift_out}")
    final_ckpt = max(
        (os.path.dirname(p) for p in adapters), key=os.path.getmtime
    )
    print(f"Final adapter checkpoint: {final_ckpt}")
    if os.path.exists(adapter_dir):
        shutil.rmtree(adapter_dir)
    shutil.copytree(final_ckpt, adapter_dir)

    _vol.commit()
    return adapter_dir


@app.local_entrypoint()
def main(max_steps: int = 0):
    path = train.remote(max_steps=max_steps)
    print(f"LoRA adapter written to volume rupkotha-finetune at {path}")
    print("Next: finetune/merge.py to fold the adapter into full weights.")