File size: 2,454 Bytes
5c2beba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | #!/usr/bin/env python3
"""Convert merged internal-format model.pth back to HF fish_qwen3_omni layout.
Builds the exact inverse of _remap_fish_qwen3_omni_keys by replaying the
forward mapping over the ORIGINAL s2-pro index, then writes safetensors
shards with the original shard assignment + index.json. Fidelity by
construction: same keys, same shards, byte-identical layout semantics.
"""
import json
from collections import OrderedDict, defaultdict
from pathlib import Path
import torch
from safetensors.torch import save_file
SRC = Path("/opt/work/checkpoints/s2-pro")
DST = Path("/opt/work/checkpoints/s2pro-egy-merged")
def forward_map(hf_key: str) -> str:
if hf_key.startswith("text_model.model."):
return hf_key[len("text_model.model."):]
if hf_key.startswith("audio_decoder."):
suffix = hf_key[len("audio_decoder."):]
return suffix if suffix.startswith("codebook_embeddings.") else "fast_" + suffix
return hf_key
index = json.loads((SRC / "model.safetensors.index.json").read_text())
weight_map = index["weight_map"] # hf_key -> shard filename
inv = {}
for hf_key in weight_map:
ik = forward_map(hf_key)
assert ik not in inv, f"collision: {ik}"
inv[ik] = hf_key
print(f"index has {len(weight_map)} keys")
sd = torch.load(DST / "model.pth", map_location="cpu", mmap=True, weights_only=True)
print(f"merged model.pth has {len(sd)} keys")
missing_in_index = [k for k in sd if k not in inv]
missing_in_model = [k for k in inv if k not in sd]
print("merged keys not in index:", missing_in_index[:5], f"({len(missing_in_index)})")
print("index keys not in merged:", missing_in_model[:5], f"({len(missing_in_model)})")
assert not missing_in_index, "unexpected keys in merged model"
assert not missing_in_model, "merged model is missing weights"
shards = defaultdict(OrderedDict)
for ik, tensor in sd.items():
hf_key = inv[ik]
shards[weight_map[hf_key]][hf_key] = tensor.contiguous()
total = 0
for shard_name, tensors in sorted(shards.items()):
total += sum(t.numel() * t.element_size() for t in tensors.values())
save_file(tensors, str(DST / shard_name), metadata={"format": "pt"})
print(f"wrote {shard_name} ({len(tensors)} tensors)")
out_index = {"metadata": {"total_size": total}, "weight_map": weight_map}
(DST / "model.safetensors.index.json").write_text(json.dumps(out_index, indent=2))
print(f"index written, total_size={total/1e9:.2f}GB")
print("CONVERT_DONE")
|