#!/usr/bin/env python3
"""Splice the 15 BF16 mtp.* tensors into the modelopt NVFP4 quant output → the servable seat.

Runs AFTER quant_modelopt.py. Copies heretic2-modelopt-nvfp4 → heretic2-modelopt-nvfp4-mtp,
splices the grafted BF16 mtp head into the single shard (transformers never builds an mtp module
at load, so mtp is always post-quant-spliced — same as AEON/pantheon), and adds the mtp module
names to config.json exclude_modules for tidiness. NOTE: the actual thing that keeps the mtp head
BF16 at serve time is the sitecustomize MTP workaround (see runbook landmine #4); the config
exclude here is belt-and-suspenders and does NOT by itself prevent the drafter-quant crash.

Run in a vLLM container (root; /tank/aimodels files are root-owned):
  docker run --rm -v /tank/aimodels:/tank/aimodels -v /home/lkraven:/lk \
    --entrypoint python3 vllm/vllm-openai:v0.24.0 /lk/finalize_modelopt_mtp.py
"""
import json
import os
import shutil

from safetensors import safe_open
from safetensors.torch import save_file

WORK = "/tank/aimodels/heretic2-nvfp4-work"
SRC = f"{WORK}/heretic2-modelopt-nvfp4"
DST = f"{WORK}/heretic2-modelopt-nvfp4-mtp"
GRAFT = f"{WORK}/heretic2-mtp-bf16"

if os.path.exists(DST):
    shutil.rmtree(DST)
print(f"copying {SRC} -> {DST}", flush=True)
shutil.copytree(SRC, DST)

out_st = f"{DST}/model.safetensors"   # single shard (quant_modelopt.py forces max_shard_size huge)
mtp_st = f"{GRAFT}/model-mtp.safetensors"
tensors = {}
with safe_open(out_st, framework="pt") as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)
n_main = len(tensors)
with safe_open(mtp_st, framework="pt") as f:
    mtp_keys = list(f.keys())
    for k in mtp_keys:
        tensors[k] = f.get_tensor(k)
assert not any("mtp" in k.lower() for k in list(tensors)[:n_main]), "output already had mtp?"
save_file(tensors, out_st, metadata={"format": "pt"})
print(f"spliced {len(mtp_keys)} bf16 mtp tensors -> {len(tensors)} total", flush=True)

cfgp = f"{DST}/config.json"
cfg = json.load(open(cfgp))
qc = cfg.setdefault("quantization_config", {})
exc = qc.setdefault("exclude_modules", [])
mtp_mods = sorted({k.rsplit(".", 1)[0] for k in mtp_keys})
exc.extend(m for m in mtp_mods if m not in exc)
json.dump(cfg, open(cfgp, "w"), indent=2)
print(f"config exclude_modules += {len(mtp_mods)} mtp modules; total {len(exc)}", flush=True)
print(f"DONE: {DST}", flush=True)