File size: 6,836 Bytes
63dd946 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | # 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)
|