| |
| |
| """Diarization in one line, at full speed. |
| |
| from anuvani_diar import Diarizer |
| diarize = Diarizer() # every visible GPU (or the CPU); Diarizer(devices=2): two GPUs |
| result = diarize("meeting.wav") # -> Diarization |
| result.save_rttm("meeting.rttm") |
| for r in diarize.run(["a.wav", "b.flac"]): # many files; results arrive as each one finishes |
| print(r.uri, r.speakers) |
| |
| Everything except the two networks runs in the Rust engine (`anuvani_diar._core`), on its own threads: decoding, |
| windowing, batching across all GPUs, masks, clustering, the decoder and the RTTM text. The output is the release's |
| (the original program, `reference/diarize.py` in the model repository): byte-identical on the setup the release's |
| outputs were produced with, and near-identical, checked at start-up, on any other GPU or PyTorch version. |
| """ |
| import json |
| import os |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Iterable, Iterator, List, Optional, Tuple, Union |
|
|
| from . import _core |
|
|
| _HERE = Path(__file__).resolve().parent |
|
|
|
|
| def default_bundle() -> Path: |
| """`$ANUVANI_DIAR_BUNDLE`; else, in a source tree, the bundle exported beside it (artifacts/bundle, made by |
| tools/export_bundle.py from the release's checkpoints: safetensors, the Indic model only); else the model folder of |
| the Hugging Face repository, downloaded once and sha256-checked (anuvani_diar.hub).""" |
| env = os.environ.get("ANUVANI_DIAR_BUNDLE") |
| if env: |
| return Path(env) |
| p = _HERE.parents[1] / "artifacts" / "bundle" |
| if (p / "MANIFEST.json").is_file(): |
| return p |
| from .hub import bundle_dir |
| return bundle_dir() |
|
|
|
|
| def devices_arg(text: str) -> Union[str, int]: |
| """A --devices value as `Diarizer(devices=...)` takes it: a bare number is a count of GPUs.""" |
| text = text.strip() |
| return int(text) if text.isdigit() else text |
|
|
|
|
| def default_golden() -> Optional[Path]: |
| env = os.environ.get("ANUVANI_DIAR_GOLDEN") |
| if env: |
| return Path(env) |
| from .hub import assets_dir |
| p = assets_dir() / "nn_golden.json" |
| return p if p.is_file() else None |
|
|
|
|
| @dataclass |
| class Diarization: |
| """One file's result. `rttm` is the exact text the release writes; `segments` are read back from it.""" |
| uri: str |
| path: str |
| rttm: Optional[str] |
| error: Optional[str] = None |
| duration_s: Optional[float] = None |
| timing: dict = field(default_factory=dict) |
| record: Optional[dict] = None |
| index: int = -1 |
|
|
| @property |
| def ok(self) -> bool: |
| return self.error is None and self.rttm is not None |
|
|
| @property |
| def segments(self) -> List[Tuple[float, float, str]]: |
| """(start, end, speaker) per RTTM line, in the file's order.""" |
| out = [] |
| for line in (self.rttm or "").splitlines(): |
| f = line.split() |
| out.append((float(f[3]), float(f[3]) + float(f[4]), f[7])) |
| return out |
|
|
| @property |
| def speakers(self) -> List[str]: |
| return sorted({s for _, _, s in self.segments}) |
|
|
| def save_rttm(self, path: Union[str, Path]) -> Path: |
| if not self.ok: |
| raise RuntimeError(self.error or f"no output for {self.path}") |
| path = Path(path) |
| path.write_bytes(self.rttm.encode()) |
| return path |
|
|
|
|
| class Diarizer: |
| """The release's diarization. |
| |
| devices: where the networks run. "auto" (default): every visible GPU, or the CPU when there is none; a number N: |
| the first N GPUs; a list of GPU indices or device names ([0, 3], ["cuda:1"], "0,3"); or "cpu". The engine |
| spreads every file's windows over all of them. |
| workers_per_device: kept at 1; one CUDA stream per GPU is what keeps the outputs deterministic. |
| precision: "exact" (default: the release's FP32 networks and settings), or a faster reduced-precision preset whose |
| outputs differ slightly: "ffn16", "tf32conv" and "mixed"; "fp32", "tf32" and "fp16" are for measurement only. |
| strict: byte-identical outputs need the setup the release's outputs were produced with. On any other GPU or |
| PyTorch version, with strict=False (default), the exact lane runs if its start-up self-test stays within |
| near-identical bounds of the release's outputs (a warning says so, and `info()["exactness"]` gives the numbers); |
| strict=True refuses to start instead. |
| """ |
|
|
| def __init__(self, bundle: Union[str, Path, None] = None, devices: Union[str, int, List, None] = "auto", |
| workers_per_device: int = 1, threads: int = 2, repack: bool = True, keep_records: bool = False, |
| decode_threads: int = 4, post_threads: Optional[int] = None, max_ready: Optional[int] = None, golden=None, |
| precision: str = "exact", strict: bool = False): |
| from .nn.torch_backend import TorchNN |
| self.bundle = Path(bundle) if bundle else default_bundle() |
| manifest = json.loads((self.bundle / "MANIFEST.json").read_text()) |
| self.manifest = manifest |
| model = manifest["modes"]["indic"] |
| if golden is None: |
| golden = default_golden() |
| if golden is None: |
| raise FileNotFoundError( |
| "assets/nn_golden.json not found: every GPU checks the networks against the release's own outputs " |
| "before use. Set ANUVANI_DIAR_GOLDEN to the file (or pass golden=False to run without the check).") |
| elif golden is False: |
| golden = None |
| self.nn = TorchNN(self.bundle, devices, workers_per_device, threads, golden=golden, precision=precision, |
| strict=strict) |
| self.engine = _core.Engine( |
| self.nn, clustering=model["clustering"], decoder_path=str(self.bundle / model["decoder"]), |
| decode_threads=decode_threads, post_threads=post_threads or max(4, (os.cpu_count() or 8) // 2), |
| repack=repack, max_ready=max_ready, keep_records=keep_records) |
|
|
| def __call__(self, path: Union[str, Path]) -> Diarization: |
| return next(iter(self.run([path]))) |
|
|
| def run(self, paths: Iterable[Union[str, Path]], out_dir: Union[str, Path, None] = None) -> Iterator[Diarization]: |
| """Diarise many files; yields results in completion order. With `out_dir`, also writes `<stem>.rttm` there.""" |
| paths = [Path(p) for p in paths] |
| jobs = [(p.stem, str(p), str(p)) for p in paths] |
| out = Path(out_dir) if out_dir else None |
| for r in self.engine.run(jobs): |
| d = Diarization(uri=r["uri"], path=r["name"], rttm=r["rttm"], error=r["error"], duration_s=r["duration_s"], |
| timing=r["timing"], record=r.get("record"), index=r["index"]) |
| if out is not None and d.ok: |
| out.mkdir(parents=True, exist_ok=True) |
| d.save_rttm(out / f"{d.uri}.rttm") |
| yield d |
|
|
| def info(self) -> dict: |
| return dict(self.nn.info(), version=self.manifest["version"], bundle=str(self.bundle)) |
|
|
