# Copyright (c) 2026 The Anuvani-Diarization authors # SPDX-License-Identifier: MIT OR Apache-2.0 """Front-end plan, reference batches, masks, voiceprint rows and float16 rounding (crate diar-core::front) against pyannote's own `Inference.slide` / `get_embeddings` (traced with recording stand-ins for the networks) and the release's `powerset_masks`, bit for bit.""" import json import subprocess import sys import numpy as np import pytest from hypothesis import given, settings, strategies as st from conftest import PROJ WIN, STEP = 256000, 25600 LENGTHS = [1, 2, 399, 400, 401, 25599, 25600, 25601, 160000, 255999, 256000, 256001, 281599, 281600, 281601, WIN + 7 * STEP, WIN + 8 * STEP - 1, WIN + 8 * STEP, WIN + 8 * STEP + 1, WIN + 9 * STEP - 1, WIN + 15 * STEP, WIN + 16 * STEP, 480000, 1_000_000, 2_000_003, 3_200_000 + 12345] @pytest.fixture(scope="module") def trace(tmp_path_factory): out = tmp_path_factory.mktemp("trace") / "trace.json" subprocess.run([sys.executable, "-s", "-W", "ignore", str(PROJ / "tests/ref_front_trace.py"), json.dumps(LENGTHS), "7", str(out)], capture_output=True, text=True, check=True) return json.loads(out.read_text()) def window_rows(n, first, count): rows = [] for i in range(first, first + count): s = i * STEP valid = max(0, min(WIN, n - s)) rows.append([s if valid else -1, valid, True]) return rows def test_model_constants(trace): from pyannote.audio.utils.powerset import Powerset # -> MIN_NUM_FRAMES = ceil(799 * 400 / 256000) = 2 assert trace["min_num_samples"] == 400 assert trace["receptive_field"] == [-0.00753125, 0.025, 0.02] assert np.array_equal(np.array(trace["mapping"], np.float32), Powerset(4, 4).mapping.numpy()) def test_segmentation_batches(core, trace): for run in trace["runs"]: n = run["n"] complete, has_last, seg, _ = core.front_plan(n) assert run["chunks"] == complete + has_last and run["sw"] == [0.0, 16.0, 1.6] ours = [dict(B=ln, rows=window_rows(n, first, ln)) for first, ln, _ in seg] assert run["seg"] == ours, n kinds = [k for _, _, k in seg] assert kinds.count("Last") == int(has_last) and kinds.count("Remainder") == int(complete % 8 > 0) def test_voiceprint_batches_and_masks(core, trace): for run in trace["runs"]: n, c = run["n"], run["chunks"] masks = np.array(run["masks"], np.uint8) used = core.front_used_masks(masks) _, _, _, vp = core.front_plan(n) ours = [] for first, nch in vp: rows, ms = [], [] for ch in range(first, first + nch): for s in range(4): rows += window_rows(n, ch, 1) ms.append(used[ch, s].astype(np.uint8).tolist()) ours.append(dict(B=4 * nch, rows=rows, masks=ms)) assert run["emb"] == ours, n assert run["emb_shape"] == [c, 4, 256] # ---------------------------------------------------------------- powerset masks vs the release function @pytest.fixture(scope="module") def ref_masks(): sys.path.insert(0, str(PROJ / "tests")) import refpkg refpkg.load() from anuvani_diar_ref.masks import powerset_masks from pyannote.audio.utils.powerset import Powerset mapping = Powerset(4, 4).mapping.numpy() return lambda p: powerset_masks(p, mapping, 11) def make_p(seed, chunks, temp, ties, invalid): rng = np.random.default_rng(seed) logits = rng.standard_normal((chunks, 799, 16)) * temp # long stretches of silence or speech logits[:, :, 0] += rng.standard_normal((chunks, 1)) * 2 p = np.exp(logits - logits.max(-1, keepdims=True)) p = (p / p.sum(-1, keepdims=True)).astype(np.float32) if ties: f = rng.integers(0, 799, size=60) k1, k2 = rng.integers(0, 16, 60), rng.integers(0, 16, 60) # exact ties at the maximum p[0, f, k1] = 0.5; p[0, f, k2] = 0.5 if invalid == "neg": p[0, 5, 3] = -1e-30 elif invalid == "nan": p[-1, 798, 15] = np.nan elif invalid == "inf": p[0, 0, 0] = np.inf elif invalid == "sum_hi": p[0, 100] = p[0, 100] * np.float32(1.0002) elif invalid == "sum_edge": p[0, 100] = p[0, 100] * np.float32(1.00010) elif invalid == "zero_row": p[0, 7] = 0 return p @settings(max_examples=300) @given(st.integers(0, 2**32 - 1), st.integers(1, 3), st.sampled_from([0.1, 1.0, 3.0, 30.0]), st.booleans(), st.sampled_from([None, None, None, "neg", "nan", "inf", "sum_hi", "sum_edge", "zero_row"])) def test_powerset_masks_match_release(core, ref_masks, seed, chunks, temp, ties, invalid): p = make_p(seed, chunks, temp, ties, invalid) try: ref, ref_err = ref_masks(p).astype(np.uint8), None except ValueError as e: ref, ref_err = None, f"ValueError: {e}" got, got_err = core.front_masks(p) assert got_err == ref_err if ref is not None: assert got.dtype == ref.dtype and got.tobytes() == ref.tobytes() def test_masks_of_real_posteriors(core, ref_masks): """The cached records' float16 posteriors, widened back to float32, through both implementations.""" import records for _, path in records.sample(25): p = np.load(path)["P"].astype(np.float32) try: ref, ref_err = ref_masks(p).astype(np.uint8), None except ValueError as e: ref, ref_err = None, f"ValueError: {e}" got, got_err = core.front_masks(p) assert got_err == ref_err and (ref is None or got.tobytes() == ref.tobytes()), path # ---------------------------------------------------------------- float16 rounding def test_f16_specials(core): x = np.array([0.0, -0.0, 1e-8, 6e-8, 5.96e-8, 2.98e-8, 65504, 65520, 65519.99, 1e5, np.inf, -np.inf, 1 / 3, 0.1, 2 ** -24, 2 ** -25, 2 ** -25 * 1.0000001, 2 ** -14, 2 ** -14 * (1 - 2 ** -11)], np.float32) assert core.f16_bits(x).tobytes() == x.astype(np.float16).view(np.uint16).tobytes() @pytest.mark.slow def test_f16_every_float32(): """All 2^32 float32 bit patterns: bits identical to numpy for every non-NaN input, NaN for every NaN input.""" from anuvani_diar import _core as core step = 1 << 26 nan_bits_differ = 0 for lo in range(0, 1 << 32, step): x = np.arange(lo, lo + step, dtype=np.uint64).astype(np.uint32).view(np.float32) ref = x.astype(np.float16).view(np.uint16) got = core.f16_bits_range(lo, lo + step) nan = np.isnan(x) assert np.array_equal(ref[~nan], got[~nan]), hex(lo) assert np.isnan(got[nan].view(np.float16)).all() nan_bits_differ += int((ref[nan] != got[nan]).sum()) print("NaN inputs whose float16 NaN bits differ from numpy:", nan_bits_differ)