| |
| |
| """Build the model bundle `anuvani_diar` loads (no pickles, the Indic model only) from the frozen release bundle. |
| |
| python -s tools/export_bundle.py --out artifacts/bundle |
| |
| Layout: |
| MANIFEST.json version, the model's clustering config and decoder (`modes.indic`), licences, the |
| sha256 of every file, and the sha256 of the reference checkpoints they were made from |
| segmentation/config.json SegmentationNet arguments |
| segmentation/model.safetensors the release's hub/pytorch_model.bin state dict, tensors unchanged |
| voiceprint/model.safetensors the release's embedding/pytorch_model.bin Lightning state dict, tensors unchanged |
| decoder.json the release's decoder, without its internal recipe label |
| |
| Every tensor is checked equal (bit for bit, same dtype and shape) to the release checkpoint it came from, and each new |
| network is loaded with strict=True. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file, save_file |
|
|
| PROJ = Path(__file__).resolve().parents[1] |
| REL = PROJ.parent / "anuvani_diar_v1" / "bundle" |
| |
| MODEL_VERSION = "1.0.0" |
| SEG_ARGS = dict(wavlm_src="wavlm_large_s80_md", wavlm_layer_num=25, wavlm_feat_dim=1024, attention_in=256, |
| ffn_hidden=1024, num_head=4, num_layer=4, kernel_size=31, dropout=0.1, use_posi=False, |
| output_activate_function=False, num_classes=16) |
|
|
|
|
| def sha256(p): |
| h = hashlib.sha256() |
| with open(p, "rb") as f: |
| for b in iter(lambda: f.read(1 << 24), b""): |
| h.update(b) |
| return h.hexdigest() |
|
|
|
|
| def same(a, b): |
| return a.dtype == b.dtype and a.shape == b.shape and torch.equal(a.view(torch.uint8) if a.dtype == torch.bool else a, b) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--out", type=Path, required=True) |
| a = ap.parse_args() |
| sys.path.insert(0, str(PROJ / "python")) |
| from anuvani_diar.nn.seg import SegmentationNet |
| from anuvani_diar.nn.voiceprint import VoiceprintNet |
| rel = json.loads((REL / "MANIFEST.json").read_text()) |
| for f, want in rel["files"].items(): |
| if sha256(REL / f) != want: |
| raise SystemExit(f"release bundle file {f} does not match its MANIFEST sha256") |
| out = a.out |
| (out / "segmentation").mkdir(parents=True, exist_ok=True) |
| (out / "voiceprint").mkdir(parents=True, exist_ok=True) |
|
|
| |
| seg_state = torch.load(REL / "hub/pytorch_model.bin", map_location="cpu") |
| save_file({k: v.contiguous() for k, v in seg_state.items()}, str(out / "segmentation/model.safetensors"), |
| metadata={"format": "pt", "source": "reference/bundle/hub/pytorch_model.bin"}) |
| (out / "segmentation/config.json").write_text(json.dumps(dict(architecture="SegmentationNet", model_args=SEG_ARGS), indent=1) + "\n") |
| ck = torch.load(REL / "embedding/pytorch_model.bin", map_location="cpu") |
| vp_state = {k: v for k, v in ck["state_dict"].items()} |
| save_file({k: v.contiguous() for k, v in vp_state.items()}, str(out / "voiceprint/model.safetensors"), |
| metadata={"format": "pt", "source": "reference/bundle/embedding/pytorch_model.bin (pyannote/wespeaker-voxceleb-resnet34-LM)"}) |
| decoder = (REL / rel["modes"]["indic"]["decoder"]).read_text() |
| label = '{"recipe": "cd1", ' |
| if not decoder.startswith(label): |
| raise SystemExit("the release's decoder does not start with its recipe label") |
| (out / "decoder.json").write_text("{" + decoder[len(label):]) |
| if json.loads((out / "decoder.json").read_text()) != {k: v for k, v in json.loads(decoder).items() if k != "recipe"}: |
| raise SystemExit("decoder.json differs from the release's decoder beyond the recipe label") |
|
|
| |
| for name, state in (("segmentation", seg_state), ("voiceprint", vp_state)): |
| back = load_file(str(out / name / "model.safetensors")) |
| assert back.keys() == state.keys(), name |
| bad = [k for k in state if not same(back[k], state[k])] |
| assert not bad, (name, bad[:5]) |
| seg = SegmentationNet.from_pretrained(out / "segmentation") |
| vp = VoiceprintNet.from_pretrained(out / "voiceprint") |
| n_seg = sum(p.numel() for p in seg.parameters()) |
| n_vp = sum(p.numel() for p in vp.parameters()) |
|
|
| files = {} |
| for p in sorted(out.rglob("*")): |
| if p.is_file() and p.name != "MANIFEST.json": |
| files[str(p.relative_to(out))] = sha256(p) |
| manifest = dict( |
| version=MODEL_VERSION, model="Anuvani-Diarization", |
| modes={"indic": dict(clustering=rel["modes"]["indic"]["clustering"], decoder="decoder.json")}, |
| networks=dict(segmentation=dict(dir="segmentation", parameters=n_seg, source_sha256=rel["files"]["hub/pytorch_model.bin"]), |
| voiceprint=dict(dir="voiceprint", parameters=n_vp, source_sha256=rel["files"]["embedding/pytorch_model.bin"])), |
| licence=rel["licence"], files=files) |
| (out / "MANIFEST.json").write_text(json.dumps(manifest, indent=1) + "\n") |
| print(json.dumps(dict(out=str(out), segmentation_parameters=n_seg, voiceprint_parameters=n_vp, files=files), indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
