Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
267 lines
11 KiB
Python
267 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert NVIDIA Nemotron 3 Diarization from its original NeMo checkpoint.
|
|
|
|
The converter reads the official ``.nemo`` archive directly, stages canonical
|
|
Safetensors plus the model/frontend configuration, and can package a
|
|
self-contained GGUF with audio.cpp's converter.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import subprocess
|
|
import tarfile
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
SOURCE_REPO = "nvidia/Nemotron-3-Diarization"
|
|
SOURCE_REVISION = "723e19c601d99b7e58fba6a14e32153e0afe48d9"
|
|
SOURCE_LICENSE = "openmdw-1.1"
|
|
SOURCE_LICENSE_URL = "https://openmdw.ai/license/1-1/"
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def load_checkpoint(source: Path) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
import torch
|
|
import yaml
|
|
|
|
if not source.is_file():
|
|
raise SystemExit(f"checkpoint not found: {source}")
|
|
with tarfile.open(source, mode="r:*") as archive:
|
|
config_member = archive.extractfile("model_config.yaml")
|
|
weights_member = archive.extractfile("model_weights.ckpt")
|
|
if config_member is None or weights_member is None:
|
|
raise SystemExit(f"NeMo archive is missing model_config.yaml or model_weights.ckpt: {source}")
|
|
config = yaml.safe_load(config_member.read().decode("utf-8"))
|
|
# delete=False: Windows cannot reopen a NamedTemporaryFile while it is open.
|
|
with tempfile.NamedTemporaryFile(suffix=".ckpt", delete=False) as checkpoint:
|
|
while block := weights_member.read(16 * 1024 * 1024):
|
|
checkpoint.write(block)
|
|
try:
|
|
loaded = torch.load(checkpoint.name, map_location="cpu", weights_only=False)
|
|
finally:
|
|
Path(checkpoint.name).unlink(missing_ok=True)
|
|
state = loaded.get("state_dict", loaded)
|
|
if not isinstance(state, dict) or not state:
|
|
raise ValueError("checkpoint did not contain a non-empty state dictionary")
|
|
return config, state
|
|
|
|
|
|
def assert_model_shape(config: dict[str, Any]) -> None:
|
|
encoder = config["encoder"]
|
|
modules = config["sortformer_modules"]
|
|
checks = {
|
|
"sample_rate": (config["sample_rate"], 16000),
|
|
"max_num_of_spks": (config["max_num_of_spks"], 8),
|
|
"encoder.n_layers": (encoder["n_layers"], 31),
|
|
"encoder.d_model": (encoder["d_model"], 512),
|
|
"encoder.n_heads": (encoder["n_heads"], 8),
|
|
"encoder.subsampling": (encoder["subsampling"], "feature_stacking"),
|
|
"encoder.subsampling_factor": (encoder["subsampling_factor"], 8),
|
|
"encoder.self_attention_model": (encoder["self_attention_model"], "rope"),
|
|
"sortformer_modules.num_spks": (modules["num_spks"], 8),
|
|
"sortformer_modules.tf_d_model": (modules["tf_d_model"], 192),
|
|
}
|
|
mismatches = [
|
|
f"{name}={actual!r}, expected {expected!r}"
|
|
for name, (actual, expected) in checks.items()
|
|
if actual != expected
|
|
]
|
|
if mismatches:
|
|
raise ValueError("not the expected Nemotron 3 Diarization architecture: " + "; ".join(mismatches))
|
|
|
|
|
|
def processor_config(config: dict[str, Any]) -> dict[str, Any]:
|
|
pre = config["preprocessor"]
|
|
return {
|
|
"feature_extractor": {
|
|
"feature_size": int(pre["features"]),
|
|
"sampling_rate": int(pre["sample_rate"]),
|
|
"n_fft": int(pre["n_fft"]),
|
|
"win_length": int(round(float(pre["window_size"]) * int(pre["sample_rate"]))),
|
|
"hop_length": int(round(float(pre["window_stride"]) * int(pre["sample_rate"]))),
|
|
"window": str(pre["window"]),
|
|
"dither": float(pre["dither"]),
|
|
"normalize": str(pre["normalize"]),
|
|
"preemphasis": float(pre.get("preemph", 0.97)),
|
|
"return_attention_mask": True,
|
|
}
|
|
}
|
|
|
|
|
|
def model_config(config: dict[str, Any], source_hash: str) -> dict[str, Any]:
|
|
encoder = config["encoder"]
|
|
modules = config["sortformer_modules"]
|
|
pre = config["preprocessor"]
|
|
return {
|
|
"model_type": "nemotron_3_diar",
|
|
"architectures": ["Nemotron3Diarization"],
|
|
"audiocpp_family": "nemotron_3_diar",
|
|
"source_model": SOURCE_REPO,
|
|
"source_revision": SOURCE_REVISION,
|
|
"source_sha256": source_hash,
|
|
"sample_rate": int(config["sample_rate"]),
|
|
"num_speakers": int(config["max_num_of_spks"]),
|
|
"high_resolution": bool(config["high_resolution"]),
|
|
"output_subsampling_factor": int(config["output_subsampling_factor"]),
|
|
"frontend": {
|
|
"num_mels": int(pre["features"]),
|
|
"n_fft": int(pre["n_fft"]),
|
|
"win_length": int(round(float(pre["window_size"]) * int(pre["sample_rate"]))),
|
|
"hop_length": int(round(float(pre["window_stride"]) * int(pre["sample_rate"]))),
|
|
"window": str(pre["window"]),
|
|
"dither": float(pre["dither"]),
|
|
"normalize": str(pre["normalize"]),
|
|
"preemphasis": float(pre.get("preemph", 0.97)),
|
|
},
|
|
"encoder": {
|
|
"feat_in": int(encoder["feat_in"]),
|
|
"hidden_size": int(encoder["d_model"]),
|
|
"intermediate_size": int(float(encoder["ff_expansion"]) * int(encoder["d_model"])),
|
|
"num_attention_heads": int(encoder["n_heads"]),
|
|
"num_hidden_layers": int(encoder["n_layers"]),
|
|
"subsampling_factor": int(encoder["subsampling_factor"]),
|
|
"pre_block_norm": bool(encoder["pre_block_norm"]),
|
|
"qkv_bias": bool(encoder["qkv_bias"]),
|
|
"qk_norm": bool(encoder["qk_norm"]),
|
|
"rope_theta": float(encoder.get("rope_base", 10000.0)),
|
|
"rotary_fraction": float(encoder.get("rotary_fraction", 1.0)),
|
|
"layer_norm_eps": 1.0e-5,
|
|
},
|
|
"head": {
|
|
"hidden_size": int(modules["tf_d_model"]),
|
|
"upsample_factor": int(encoder["subsampling_factor"]),
|
|
"upsample_kernel_size": 3,
|
|
},
|
|
"streaming": {
|
|
"spkcache_len": int(modules["spkcache_len"]),
|
|
"fifo_len": int(modules["fifo_len"]),
|
|
"chunk_len": int(modules["chunk_len"]),
|
|
"chunk_left_context": int(modules["chunk_left_context"]),
|
|
"chunk_right_context": int(modules["chunk_right_context"]),
|
|
"spkcache_update_period": int(modules["spkcache_update_period"]),
|
|
"spkcache_sil_frames_per_spk": int(modules["spkcache_sil_frames_per_spk"]),
|
|
"pred_score_threshold": float(modules["pred_score_threshold"]),
|
|
"scores_boost_latest": float(modules["scores_boost_latest"]),
|
|
"sil_threshold": float(modules["sil_threshold"]),
|
|
"strong_boost_rate": float(modules["strong_boost_rate"]),
|
|
"weak_boost_rate": float(modules["weak_boost_rate"]),
|
|
"min_pos_scores_rate": float(modules["min_pos_scores_rate"]),
|
|
"max_index": int(modules["max_index"]),
|
|
},
|
|
}
|
|
|
|
|
|
def staging_spec() -> dict[str, Any]:
|
|
return json.loads(Path("model_specs/nemotron_3_diar.json").read_text(encoding="utf-8"))
|
|
|
|
|
|
def convert(checkpoint: Path, output_dir: Path) -> None:
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
config, state = load_checkpoint(checkpoint)
|
|
assert_model_shape(config)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
source_hash = sha256(checkpoint)
|
|
|
|
tensors: dict[str, torch.Tensor] = {}
|
|
manifest: list[dict[str, Any]] = []
|
|
for source_name, value in sorted(state.items()):
|
|
if source_name == "preprocessor.featurizer.window":
|
|
continue
|
|
if source_name == "preprocessor.featurizer.fb":
|
|
value = value.squeeze(0)
|
|
destination = "preprocessor.fb"
|
|
else:
|
|
destination = source_name
|
|
if not isinstance(value, torch.Tensor):
|
|
raise TypeError(f"expected tensor for {source_name}, got {type(value).__name__}")
|
|
if value.dtype not in (torch.float32, torch.float16, torch.bfloat16):
|
|
raise ValueError(f"unsupported tensor dtype for {source_name}: {value.dtype}")
|
|
tensor = value.detach().cpu().contiguous()
|
|
tensors[destination] = tensor
|
|
manifest.append({
|
|
"source": source_name,
|
|
"destination": destination,
|
|
"shape": list(tensor.shape),
|
|
"dtype": str(tensor.dtype),
|
|
})
|
|
|
|
save_file(tensors, str(output_dir / "model.safetensors"), metadata={
|
|
"format": "nemotron-3-diarization-audiocpp-staging",
|
|
"source_repo": SOURCE_REPO,
|
|
"source_revision": SOURCE_REVISION,
|
|
})
|
|
(output_dir / "config.json").write_text(
|
|
json.dumps(model_config(config, source_hash), indent=2) + "\n", encoding="utf-8")
|
|
(output_dir / "processor_config.json").write_text(
|
|
json.dumps(processor_config(config), indent=2) + "\n", encoding="utf-8")
|
|
(output_dir / "model_spec.json").write_text(
|
|
json.dumps(staging_spec(), indent=2) + "\n", encoding="utf-8")
|
|
(output_dir / "tensor_manifest.json").write_text(json.dumps({
|
|
"source": SOURCE_REPO,
|
|
"revision": SOURCE_REVISION,
|
|
"source_sha256": source_hash,
|
|
"source_tensor_count": len(state),
|
|
"emitted_tensor_count": len(tensors),
|
|
"tensors": manifest,
|
|
}, indent=2) + "\n", encoding="utf-8")
|
|
(output_dir / "provenance.json").write_text(json.dumps({
|
|
"source_repo": SOURCE_REPO,
|
|
"source_revision": SOURCE_REVISION,
|
|
"source_license": SOURCE_LICENSE,
|
|
"source_license_url": SOURCE_LICENSE_URL,
|
|
"source_sha256": source_hash,
|
|
}, indent=2) + "\n", encoding="utf-8")
|
|
print(f"source tensors: {len(state)}")
|
|
print(f"emitted tensors: {len(tensors)}")
|
|
print(f"wrote staging package: {output_dir}")
|
|
|
|
|
|
def package(args: argparse.Namespace) -> None:
|
|
if args.converter is None or args.gguf_output is None:
|
|
return
|
|
command = [
|
|
str(args.converter.resolve()),
|
|
"--input", str((args.output_dir / "model.safetensors").resolve()),
|
|
"--root", str(args.output_dir.resolve()),
|
|
"--family", "nemotron_3_diar",
|
|
"--model-spec", str((args.output_dir / "model_spec.json").resolve()),
|
|
"--type", args.type,
|
|
"--keep-type", "preprocessor.fb=f32",
|
|
"--output", str(args.gguf_output.resolve()),
|
|
"--overwrite",
|
|
]
|
|
args.gguf_output.resolve().parent.mkdir(parents=True, exist_ok=True)
|
|
print("+", " ".join(command))
|
|
subprocess.run(command, check=True)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--checkpoint", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--converter", type=Path)
|
|
parser.add_argument("--gguf-output", type=Path)
|
|
parser.add_argument("--type", choices=["orig", "f16", "q8_0"], default="orig")
|
|
args = parser.parse_args()
|
|
args.output_dir = args.output_dir.resolve()
|
|
convert(args.checkpoint.resolve(), args.output_dir)
|
|
package(args)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|