Anuvani-Diarization / src /tests /test_front.py
StephennFernandes's picture
Anuvani-Diarization 1.0.0
63dd946
Raw History Blame Contribute Delete
6.84 kB
# 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)