#!/usr/bin/env python
"""Extract the small projection heads the pipeline/solver need from model.safetensors.
All are Linear/LayerNorm (weight (out,in)) - no transpose. CPU/numpy, sandbox-safe."""
from safetensors import safe_open
from safetensors.numpy import save_file

SNAP = ("/Users/samm/.cache/huggingface/hub/models--rednote-hilab--dots.tts-soar/"
        "snapshots/1fd9452e55c2c9f38fe1a8ee09eaf7448c222d35/model.safetensors")
OUT = "/Users/samm/git/sammcj/dots.tts-soar-mlx/heads/model.safetensors"
PREFIXES = ("coordinate_proj.", "hidden_proj.", "latent_proj.", "xvec_proj.", "eos_proj.")

out = {}
with safe_open(SNAP, framework="numpy") as f:
    for k in f.keys():
        if k.startswith(PREFIXES):
            out[k] = f.get_tensor(k)

import os
os.makedirs(os.path.dirname(OUT), exist_ok=True)
save_file(out, OUT)
for k, v in sorted(out.items()):
    print(k, v.shape, v.dtype)
print(f"wrote {OUT} ({len(out)} tensors)")