# Derived from NVIDIA LocateAnything-3B (modeling_vit.py / modeling_locateanything.py, Copyright (c) 2026 NVIDIA
# CORPORATION) under the NVIDIA License (see ../LICENSE); MoonViT is MIT-licensed by Moonshot AI. Modified.
import json, shutil, os, torch
from safetensors import safe_open
from safetensors.torch import save_file
from transformers import Qwen2Config, Qwen2ForCausalLM
SRC="model"; DST="qwen25_3b_la"
os.makedirs(DST, exist_ok=True)
idx=json.load(open(f"{SRC}/model.safetensors.index.json"))["weight_map"]
llm={}; extra={}
files={}
for k,f in idx.items(): files.setdefault(f,[]).append(k)
for f,keys in files.items():
    with safe_open(f"{SRC}/{f}","pt") as sf:
        for k in keys:
            if k.startswith("language_model."): llm[k[len("language_model."):]]=sf.get_tensor(k)
            elif k.startswith("mlp1."): extra[k]=sf.get_tensor(k)
tied = torch.equal(llm["lm_head.weight"], llm["model.embed_tokens.weight"])
print("lm_head == embed_tokens:", tied, "| vocab rows:", llm["model.embed_tokens.weight"].shape)
cfg=json.load(open(f"{SRC}/config.json"))["text_config"]
qc=Qwen2Config(vocab_size=cfg["vocab_size"],hidden_size=cfg["hidden_size"],intermediate_size=cfg["intermediate_size"],
  num_hidden_layers=cfg["num_hidden_layers"],num_attention_heads=cfg["num_attention_heads"],num_key_value_heads=cfg["num_key_value_heads"],
  max_position_embeddings=cfg["max_position_embeddings"],rms_norm_eps=cfg["rms_norm_eps"],rope_theta=cfg["rope_theta"],
  tie_word_embeddings=tied,use_sliding_window=False,sliding_window=None,bos_token_id=cfg["bos_token_id"],eos_token_id=cfg["eos_token_id"],torch_dtype="bfloat16")
qc.architectures=["Qwen2ForCausalLM"]
qc.save_pretrained(DST)
if tied: llm.pop("lm_head.weight")
save_file({k:v.contiguous() for k,v in llm.items()}, f"{DST}/model.safetensors", metadata={"format":"pt"})
save_file({k:v.contiguous() for k,v in extra.items()}, f"{DST}/mlp1.safetensors")
for f in ["vocab.json","merges.txt","tokenizer_config.json","added_tokens.json","special_tokens_map.json","generation_config.json"]:
    shutil.copy(f"{SRC}/{f}", f"{DST}/{f}")
# sanity: reload as plain Qwen2 and compare to original LLM logits
m=Qwen2ForCausalLM.from_pretrained(DST, torch_dtype=torch.float32).eval()
ids=torch.randint(0,150000,(1,32))
with torch.no_grad(): lg=m(ids).logits
print("reloaded ok, logits", tuple(lg.shape), "finite:", torch.isfinite(lg).all().item())
torch.save({"ids":ids,"logits":lg}, "ref_logits_plain.pt")