#!/usr/bin/env python3
"""Physically remove tensors not referenced by the index from /data/glm52.

vLLM's local-dir loader iterates every tensor in every *.safetensors file,
ignoring the index, so de-indexed tensors must actually be deleted. Rewrites
any shard containing unreferenced keys; deletes shards left empty.
"""
import json
import os
import sys

from safetensors import safe_open
from safetensors.torch import save_file

DST = sys.argv[1] if len(sys.argv) > 1 else "/data/glm52"

idx = json.load(open(os.path.join(DST, "model.safetensors.index.json")))
wm = idx["weight_map"]

files = sorted(f for f in os.listdir(DST) if f.endswith(".safetensors"))
total_removed = 0
for f in files:
    path = os.path.join(DST, f)
    with safe_open(path, framework="pt") as sf:
        keys = list(sf.keys())
        keep = [k for k in keys if wm.get(k) == f]
        if len(keep) == len(keys):
            continue
        removed = len(keys) - len(keep)
        if not keep:
            print(f"{f}: all {len(keys)} tensors stale -> deleting file")
            tensors = None
        else:
            print(f"{f}: removing {removed}/{len(keys)} stale tensors")
            tensors = {k: sf.get_tensor(k) for k in keep}
    total_removed += removed
    if tensors is None:
        os.remove(path)
    else:
        tmp = path + ".tmp"
        save_file(tensors, tmp)
        os.replace(tmp, path)

print(f"done: removed {total_removed} stale tensors")