"""为视觉塔单独搭一个「仅含 shard24」的模型目录，便于在剩余分片下载完成前先导出 mmproj。

做法：只把含 model.visual.* 的分片软链进来，并把 index.json 过滤成只列该分片，
这样 convert_hf_to_gguf.py 的 index_tensors 不会去打开还没下完的分片。
"""
import json
import os
import shutil
from pathlib import Path

ROOT = Path(__file__).resolve().parent
SRC = ROOT / "models" / "Youtu-LLM-A7B-Base"
DST = ROOT / "models" / "youtu-vision-only"

DST.mkdir(parents=True, exist_ok=True)

idx = json.loads((SRC / "model.safetensors.index.json").read_text())
wm = idx["weight_map"]
vis = {k: v for k, v in wm.items() if k.startswith("model.visual.")}
if not vis:
    raise SystemExit("没找到 model.visual.* 张量")
shards = sorted(set(vis.values()))
print(f"visual 张量 {len(vis)} 个，位于 {shards}")
if len(shards) != 1:
    print("[warn] 视觉张量跨多个分片，需全部就绪")

# 过滤后的 index 必须把「就绪分片里的全部张量」都列上，
# 否则 index_tensors 的 weight_map / part 一致性校验会报 Extra tensors。
keep = {k: v for k, v in wm.items() if v in shards}
print(f"索引保留 {len(keep)} 个张量（含同分片的文本张量，转换时按名过滤）")

for s in shards:
    src = SRC / s
    if not src.exists():
        raise SystemExit(f"缺少分片 {s}，尚未下载完成")
    dst = DST / s
    if dst.exists() or dst.is_symlink():
        dst.unlink()
    os.symlink(os.path.relpath(src, DST), dst)

for f in SRC.iterdir():
    if f.is_file() and f.suffix in (".json", ".jinja", ".txt", ".py") and not f.name.endswith(".parts.json"):
        shutil.copy2(f, DST / f.name)

(DST / "model.safetensors.index.json").write_text(
    json.dumps({"metadata": {"total_size": 0}, "weight_map": keep}, indent=2)
)
print(f"已准备 {DST}")