| |
| """Convert a merged fish-native model dir (model.pth + dual_ar config) back to the HF fish_qwen3_omni |
| layout (safetensors shards + original index/config/tokenizer/codec) for sglang-omni / vLLM-omni. |
| Exact inverse of _remap_fish_qwen3_omni_keys, replayed over the ORIGINAL s2-pro index. |
| |
| python scripts/export_hf.py --src checkpoints/s2-pro --merged checkpoints/s2pro-egy-v2 --dst checkpoints/s2pro-egy-v2-hf |
| """ |
| import argparse |
| import json |
| import shutil |
| from collections import OrderedDict, defaultdict |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import save_file |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--src", default="checkpoints/s2-pro") |
| ap.add_argument("--merged", required=True) |
| ap.add_argument("--dst", required=True) |
| a = ap.parse_args() |
| SRC, MERGED, DST = Path(a.src), Path(a.merged), Path(a.dst) |
| DST.mkdir(parents=True, exist_ok=True) |
|
|
|
|
| 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"] |
| inv = {} |
| for hf_key in weight_map: |
| ik = forward_map(hf_key) |
| assert ik not in inv, f"collision: {ik}" |
| inv[ik] = hf_key |
|
|
| sd = torch.load(MERGED / "model.pth", map_location="cpu", mmap=True, weights_only=True) |
| if "state_dict" in sd: |
| sd = sd["state_dict"] |
| sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items() if "lora" not in k} |
| 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(f"index keys {len(weight_map)} | merged keys {len(sd)} | not in index {len(missing_in_index)} | missing from merged {len(missing_in_model)}") |
| assert not missing_in_index, missing_in_index[:5] |
| assert not missing_in_model, missing_in_model[:5] |
|
|
| shards = defaultdict(OrderedDict) |
| for ik, t in sd.items(): |
| shards[weight_map[inv[ik]]][inv[ik]] = t.to(torch.bfloat16).contiguous() |
| total = 0 |
| for name, tensors in sorted(shards.items()): |
| total += sum(t.numel() * t.element_size() for t in tensors.values()) |
| save_file(tensors, str(DST / name), metadata={"format": "pt"}) |
| print(f"wrote {name} ({len(tensors)} tensors)") |
| (DST / "model.safetensors.index.json").write_text(json.dumps({"metadata": {"total_size": total}, "weight_map": weight_map}, indent=2)) |
| for f in ["config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "chat_template.jinja", "LICENSE.md", "codec.pth"]: |
| if (SRC / f).exists(): |
| shutil.copy2(SRC / f, DST / f) |
| print(f"EXPORT_DONE {DST} total_size={total/1e9:.2f}GB") |
|
|
