"""Speaker-similarity measurements on top of the ECAPA-TDNN embedding worker.
Two uses, driven by ``scripts/speaker_sim.py``:
* **say**: cosine similarity of every generated chunk of a ``nida say`` output directory to
the embedding of its speaker prompt (:func:`read_say_chunks`, :func:`resolve_references`,
:func:`score_say`, :func:`summarize_say`);
* **groups**: does a dataset ``group`` (speaker / source video) really hold one voice? Intra-
group versus inter-group cosine similarity of sampled segments (:func:`read_index`,
:func:`select_candidates`, :func:`sample_groups`, :func:`summarize_groups`).
Embedding is a side effect at the edge (:func:`embed_with_adapter`: the worker protocol with
its resumable results file, run in-process); everything else is pure and needs only numpy.
Embeddings are 192-d float32 unit vectors, so a cosine similarity is a dot product. Thresholds
are arguments (the script's flags), never constants of the measurement.
``say`` layout relied upon (``nida.infer.say``): ``say.jsonl`` (``prompt_id``, ``index`` per
chunk), ``chunks/<prompt_id>__<index:02d>.16k.wav`` (the 16 kHz generated chunk),
``settings.json`` (``external_prompts``: external prompt id -> source wav path as given to
``--prompt-list``, with ``external_prompt_texts`` and ``external_prompt_sha256`` beside it in newer
runs; all absent or empty for dataset prompts) and ``audio/<chunk>.16k.wav`` (the decoder's 16 kHz
output for ``prompt + sentence``: ``say`` cuts the generated chunk from it with
``prompt_frames`` / ``num_frames`` of ``say.jsonl``, and the same cut's first part is the
decoded prompt).
"""
from __future__ import annotations
import hashlib
import io
import json
import logging
import math
import os
import random
from collections import defaultdict
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from math import gcd
from pathlib import Path
from typing import Literal
import numpy as np
import numpy.typing as npt
from nida.bench.protocol import (
WorkRequest,
drop_results,
load_requests,
load_results,
read_summary,
write_requests,
)
from nida.bench.workers.runner import Adapter, ModelRef, WorkerOptions, run_worker
from nida.data.workers.spk_ecapa import EMBEDDING_DIM, sha256_file
from nida.text.turkish import nfc
from nida.utils.io import atomic_write_bytes
__all__ = [
"DEFAULT_DRIFT_THRESHOLD",
"ChunkRef",
"ClipEmbedding",
"ClipFailure",
"EmbedOutcome",
"Embedder",
"ExternalPromptInfo",
"IndexRow",
"PromptListEntry",
"ReferenceMode",
"ReferenceRef",
"SpeakerSimInputError",
"cosine",
"cut_prompt_audio",
"describe",
"embed_with_adapter",
"find_dataset_dir",
"list_audio_stems",
"prompt_id_for",
"read_external_prompts",
"read_index",
"read_prompt_list",
"read_say_chunks",
"read_say_context",
"resolve_references",
"sample_groups",
"score_say",
"select_candidates",
"summarize_groups",
"summarize_say",
"unmatched_prompt_entries",
"validate_group_options",
"validate_threshold",
"write_16k_mono_wav",
"write_decoded_prompt_wav",
"write_if_changed",
]
log = logging.getLogger(__name__)
#: Default for ``--threshold`` of the ``say`` subcommand (the user's choice, 2026-10-15).
DEFAULT_DRIFT_THRESHOLD = 0.5
#: Sample rate the embedding worker requires.
WORK_RATE_HZ = 16000
#: Unit-norm tolerance for embeddings read back from disk.
NORM_TOLERANCE = 1e-3
SAY_MANIFEST = "say.jsonl"
SAY_SETTINGS = "settings.json"
LOWEST_CHUNKS = 5
ReferenceMode = Literal["auto", "source", "decoded"]
FloatVec = npt.NDArray[np.float32]
class SpeakerSimInputError(ValueError):
"""An input directory, manifest or option is missing, malformed or inconsistent."""
# ------------------------------------------------------------------------------ embedding
@dataclass(frozen=True)
class ClipEmbedding:
"""A unit-norm embedding ``(192,)`` and the duration of the clip it came from."""
vector: FloatVec
duration_ms: float
@dataclass(frozen=True)
class ClipFailure:
"""A clip the worker could not embed (``"<ExceptionType>: <message>"``)."""
error: str
EmbedOutcome = ClipEmbedding | ClipFailure
#: Embeds the clips of the requests, keeping its files (results, embeddings) in the given work
#: directory (:func:`embed_with_adapter`): ``clip_id -> outcome``.
Embedder = Callable[[Sequence[WorkRequest], Path], Mapping[str, EmbedOutcome]]
def embed_with_adapter(
adapter: Adapter,
model: ModelRef,
requests: Sequence[WorkRequest],
work_dir: Path,
*,
device: str,
batch_size: int,
) -> dict[str, EmbedOutcome]:
"""Run ``adapter`` over ``requests`` through the worker protocol and read the embeddings.
``work_dir`` holds ``requests.jsonl``, ``results.jsonl``, ``results.summary.json`` and
``embeddings/<clip_id>.npy``. The run is resumable: results of an earlier invocation are
reused when their embedding file exists and the clip's audio is unchanged (each request
carries an ``audio_fingerprint`` param, size and mtime of the file, so a ``say`` rerun into
the same directory is never served stale embeddings, and a ``model`` param ``id@revision``);
failed clips are tried again, results of clips no longer requested, changed or lost are
dropped. Results of another model / revision are refused. Each request's ``out_path`` is set
here.
Raises:
SpeakerSimInputError: results of another model, no result for a request, or an
embedding file that is not a unit vector of the expected shape.
TooManyFailuresError: (from the worker loop) the model or device is broken.
"""
work_dir.mkdir(parents=True, exist_ok=True)
emb_dir = work_dir / "embeddings"
prepared = [
WorkRequest(
r.clip_id,
r.audio_path,
{
**r.params,
"out_path": str(emb_dir / f"{r.clip_id}.npy"),
"audio_fingerprint": _fingerprint(r.audio_path),
"model": f"{model.model_id}@{model.revision}",
},
)
for r in requests
]
requests_path = work_dir / "requests.jsonl"
results_path = work_dir / "results.jsonl"
summary_path = work_dir / "results.summary.json"
if results_path.exists() and summary_path.exists():
previous = read_summary(summary_path)
if (previous.model_id, previous.revision) != (model.model_id, model.revision):
raise SpeakerSimInputError(
f"{work_dir}: existing results come from {previous.model_id}@{previous.revision}, "
f"not {model.model_id}@{model.revision}; move them away first"
)
if results_path.exists():
# Without the previous requests nothing proves a result belongs to the current audio.
before = (
{r.clip_id: r for r in load_requests(requests_path)} if requests_path.exists() else {}
)
current = {r.clip_id: r for r in prepared}
existing = load_results(results_path, repair=True)
doomed = [
cid
for cid, res in existing.items()
if cid not in current
or before.get(cid) != current[cid]
or res.output is None
or not Path(str(res.output.get("out_path", ""))).is_file()
]
if doomed:
log.info("dropping %d stray / failed / lost results before the run", len(doomed))
drop_results(results_path, doomed)
write_requests(requests_path, prepared)
run_worker(
adapter,
WorkerOptions(
requests=requests_path,
results=results_path,
summary=summary_path,
model=model,
device=device,
batch_size=batch_size,
retry_failed=False,
),
)
results = load_results(results_path)
outcomes: dict[str, EmbedOutcome] = {}
for request in prepared:
result = results.get(request.clip_id)
if result is None:
raise SpeakerSimInputError(f"{request.clip_id}: the worker produced no result")
if result.output is None:
outcomes[request.clip_id] = ClipFailure(str(result.error))
continue
outcomes[request.clip_id] = _read_embedding(request.clip_id, result.output)
return outcomes
def _fingerprint(audio_path: str) -> str:
"""``<size>:<mtime_ns>`` of an audio file: changes whenever the file is rewritten."""
try:
stat = Path(audio_path).stat()
except OSError as exc:
raise SpeakerSimInputError(f"{audio_path}: cannot stat the audio ({exc})") from exc
return f"{stat.st_size}:{stat.st_mtime_ns}"
def _read_embedding(clip_id: str, output: Mapping[str, object]) -> ClipEmbedding:
path = Path(str(output["out_path"]))
vector = np.load(path, allow_pickle=False)
if vector.shape != (EMBEDDING_DIM,) or vector.dtype != np.float32:
raise SpeakerSimInputError(
f"{clip_id}: {path} has shape {vector.shape} / dtype {vector.dtype}, "
f"expected ({EMBEDDING_DIM},) float32"
)
norm = float(np.linalg.norm(vector))
if not math.isfinite(norm) or abs(norm - 1.0) > NORM_TOLERANCE:
raise SpeakerSimInputError(f"{clip_id}: {path} is not a unit vector (norm {norm})")
duration = output.get("duration_ms")
if isinstance(duration, bool) or not isinstance(duration, int | float):
raise SpeakerSimInputError(f"{clip_id}: worker output lacks duration_ms")
return ClipEmbedding(vector, float(duration))
# --------------------------------------------------------------------------- numeric helpers
def cosine(a: npt.ArrayLike, b: npt.ArrayLike) -> float:
"""Cosine similarity of two non-zero finite vectors of equal shape.
Raises:
ValueError: on a shape mismatch, non-finite values or a zero vector.
"""
x, y = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64)
if x.shape != y.shape or x.ndim != 1:
raise ValueError(f"cosine needs two 1-D vectors of one length, got {x.shape} and {y.shape}")
if not (np.isfinite(x).all() and np.isfinite(y).all()):
raise ValueError("cosine: non-finite values")
nx, ny = float(np.linalg.norm(x)), float(np.linalg.norm(y))
if nx == 0.0 or ny == 0.0:
raise ValueError("cosine: zero vector")
return float(np.dot(x, y) / (nx * ny))
def describe(values: Sequence[float]) -> dict[str, float | int | None]:
"""``n``, ``mean``, ``std`` (sample, ``ddof=1``; 0.0 for one value), ``min``, ``max`` and the
5/25/50/75/95th percentiles (linear interpolation) of ``values``; ``None`` fields if empty.
Raises:
ValueError: on non-finite values.
"""
keys = ("mean", "std", "min", "max", "p05", "p25", "p50", "p75", "p95")
if not values:
return {"n": 0, **dict.fromkeys(keys)}
data = np.asarray(values, dtype=np.float64)
if not np.isfinite(data).all():
raise ValueError("describe: non-finite values")
p05, p25, p50, p75, p95 = (float(x) for x in np.percentile(data, [5, 25, 50, 75, 95]))
return {
"n": int(data.size),
"mean": float(data.mean()),
"std": float(data.std(ddof=1)) if data.size > 1 else 0.0,
"min": float(np.min(data)),
"max": float(np.max(data)),
"p05": p05,
"p25": p25,
"p50": p50,
"p75": p75,
"p95": p95,
}
def write_if_changed(path: Path, data: bytes) -> bool:
"""Atomically write ``data`` to ``path`` unless it already holds exactly these bytes.
Returns ``True`` if the file was written. An identical file is left alone so that its
mtime stays put: the embedding loop fingerprints audio by size and mtime, and a rewrite
would make a rerun redo work that is still valid.
"""
if path.is_file() and path.read_bytes() == data:
return False
atomic_write_bytes(path, data)
return True
def _wav_bytes(audio: FloatVec) -> bytes:
import soundfile # type: ignore[import-untyped] # bench extra
buffer = io.BytesIO()
soundfile.write(buffer, np.clip(audio, -1.0, 1.0), WORK_RATE_HZ, format="WAV", subtype="PCM_16")
return buffer.getvalue()
def write_16k_mono_wav(source: Path, target: Path) -> float:
"""Write a mono 16 kHz 16-bit PCM copy of ``source`` (unless identical); returns its
duration (ms).
Channels are averaged and the signal resampled with a polyphase filter, exactly as
``nida.infer.external_prompt`` does for its prompts. Needs soundfile and scipy.
Raises:
SpeakerSimInputError: if the file is empty.
"""
import soundfile
audio, rate = soundfile.read(str(source), dtype="float32", always_2d=True)
mono = audio.mean(axis=1).astype(np.float32)
if mono.size == 0:
raise SpeakerSimInputError(f"{source}: empty audio")
if int(rate) != WORK_RATE_HZ:
from scipy import signal # type: ignore[import-untyped] # bench extra
g = gcd(int(rate), WORK_RATE_HZ)
mono = signal.resample_poly(mono, WORK_RATE_HZ // g, int(rate) // g).astype(np.float32)
write_if_changed(target, _wav_bytes(mono))
return float(1000.0 * mono.shape[0] / WORK_RATE_HZ)
def cut_prompt_audio(audio: FloatVec, prompt_frames: int, num_frames: int) -> FloatVec:
"""The prompt part of a decoded ``prompt + sentence`` item: the first
``round(prompt_frames * len(audio) / num_frames)`` samples, the cut of ``nida.infer.say._cut``
(a test pins equality) that also produced the generated chunks from the same file.
Raises:
SpeakerSimInputError: unless ``0 < prompt_frames <= num_frames``.
"""
if not 0 < prompt_frames <= num_frames:
raise SpeakerSimInputError(
f"need 0 < prompt_frames <= num_frames, got {prompt_frames} / {num_frames}"
)
at = round(prompt_frames * audio.shape[0] / num_frames)
return audio[:at]
def write_decoded_prompt_wav(
audio_path: Path, prompt_frames: int, num_frames: int, target: Path
) -> float:
"""Write the prompt part of a decoded ``say`` item (``audio/<chunk>.16k.wav``, 16 kHz) as
the reference wav (unless identical); returns its duration (ms).
Raises:
SpeakerSimInputError: not 16 kHz audio, bad frame counts, or an empty prompt part.
"""
import soundfile
audio, rate = soundfile.read(str(audio_path), dtype="float32")
if int(rate) != WORK_RATE_HZ:
raise SpeakerSimInputError(f"{audio_path}: expected {WORK_RATE_HZ} Hz, got {rate} Hz")
if audio.ndim != 1:
audio = audio.mean(axis=1).astype(np.float32)
prompt = cut_prompt_audio(audio, prompt_frames, num_frames)
if prompt.size == 0:
raise SpeakerSimInputError(f"{audio_path}: the prompt part is empty")
write_if_changed(target, _wav_bytes(prompt))
return float(1000.0 * prompt.shape[0] / WORK_RATE_HZ)
def validate_threshold(name: str, value: float) -> None:
"""Raise unless ``value`` is a finite number (a cosine threshold).
Raises:
SpeakerSimInputError: on NaN or infinity.
"""
if not math.isfinite(value):
raise SpeakerSimInputError(f"{name} must be a finite number, got {value}")
def validate_group_options(
*,
n_groups: int,
per_group: int,
inter_pairs: int,
min_duration_ms: float,
intra_threshold: float,
) -> None:
"""Check the ``groups`` options before any model work.
Raises:
SpeakerSimInputError: ``n_groups < 1``, ``per_group < 2`` (one segment has no intra
pair), ``inter_pairs < 1``, a negative / non-finite ``min_duration_ms`` or a
non-finite ``intra_threshold``.
"""
if n_groups < 1:
raise SpeakerSimInputError(f"--groups must be >= 1, got {n_groups}")
if per_group < 2:
raise SpeakerSimInputError(f"--per-group must be >= 2, got {per_group}")
if inter_pairs < 1:
raise SpeakerSimInputError(f"--inter-pairs must be >= 1, got {inter_pairs}")
if not math.isfinite(min_duration_ms) or min_duration_ms < 0:
raise SpeakerSimInputError(f"--min-duration-ms must be >= 0, got {min_duration_ms}")
validate_threshold("--intra-threshold", intra_threshold)
# ------------------------------------------------------------------------------- say: input
@dataclass(frozen=True)
class ChunkRef:
"""One generated chunk of a ``say`` run.
``prompt_frames`` / ``num_frames`` are the item's frame counts from ``say.jsonl`` (``None``
in a manifest that does not have them); they locate the prompt inside the decoded audio.
"""
chunk_id: str
prompt_id: str
index: int
wav_path: Path
prompt_frames: int | None = None
num_frames: int | None = None
@dataclass(frozen=True)
class ReferenceRef:
"""The audio a prompt's chunks are compared with.
Attributes:
kind: ``"source"`` (the recording the prompt came from) or ``"decoded"`` (the prompt
part of the decoded audio the model's output was cut from).
wav_path: The file: a source recording (any rate / channel count), or for ``decoded``
the 16 kHz ``audio/<chunk>.16k.wav`` of the prompt's first chunk.
prompt_frames: ``decoded`` only: latent frames of the prompt ...
num_frames: ... and of the whole item; the prompt is the first
``prompt_frames / num_frames`` of the file (:func:`cut_prompt_audio`).
"""
kind: Literal["source", "decoded"]
wav_path: Path
prompt_frames: int | None = None
num_frames: int | None = None
def __post_init__(self) -> None:
if self.kind == "decoded" and (self.prompt_frames is None or self.num_frames is None):
raise ValueError("a decoded reference needs prompt_frames and num_frames")
@dataclass(frozen=True)
class PromptListEntry:
"""One line of a ``--prompt-list`` file with the id ``nida say`` gave the prompt."""
prompt_id: str
wav_path: Path
transcript: str
@dataclass(frozen=True)
class ExternalPromptInfo:
"""What ``settings.json`` of a say run records about one external prompt.
``source`` is the path as given to ``--prompt-list``; ``text`` (the normalised transcript)
and ``sha256`` (of the recording when ``say`` ran) are ``None`` for runs from before they
were recorded, or if the file was not readable then.
"""
source: str
text: str | None = None
sha256: str | None = None
def chunk_id(prompt_id: str, index: int) -> str:
"""Chunk name used by ``say`` for its files: ``<prompt_id>__<index:02d>``."""
return f"{prompt_id}__{index:02d}"
def _optional_count(record: Mapping[str, object], key: str, where: str) -> int | None:
value = record.get(key)
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise SpeakerSimInputError(f"{where}: bad {key} {value!r}")
return value
def read_say_chunks(say_dir: Path) -> list[ChunkRef]:
"""The chunks listed in ``<say_dir>/say.jsonl`` with their 16 kHz wav paths.
Raises:
SpeakerSimInputError: missing / empty / unparseable manifest, a record without a
valid ``prompt_id`` or ``index`` (or a malformed ``prompt_frames`` /
``num_frames``), a duplicate chunk, or missing chunk audio (the first missing
path is named).
"""
path = say_dir / SAY_MANIFEST
if not path.is_file():
raise SpeakerSimInputError(f"{path}: no say manifest (a `nida say` output directory?)")
refs: list[ChunkRef] = []
seen: set[str] = set()
for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
if not line.strip():
continue
where = f"{path} line {line_no}"
try:
record = json.loads(line)
except json.JSONDecodeError as exc:
raise SpeakerSimInputError(f"{where}: invalid JSON ({exc})") from exc
if not isinstance(record, dict):
raise SpeakerSimInputError(f"{where}: expected a JSON object")
prompt_id, index = record.get("prompt_id"), record.get("index")
if not isinstance(prompt_id, str) or not prompt_id:
raise SpeakerSimInputError(f"{where}: bad prompt_id {prompt_id!r}")
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
raise SpeakerSimInputError(f"{where}: bad index {index!r}")
cid = chunk_id(prompt_id, index)
if cid in seen:
raise SpeakerSimInputError(f"{where}: duplicate chunk {cid!r}")
seen.add(cid)
refs.append(
ChunkRef(
cid,
prompt_id,
index,
say_dir / "chunks" / f"{cid}.16k.wav",
_optional_count(record, "prompt_frames", where),
_optional_count(record, "num_frames", where),
)
)
if not refs:
raise SpeakerSimInputError(f"{path}: no chunks")
missing = [str(r.wav_path) for r in refs if not r.wav_path.is_file()]
if missing:
raise SpeakerSimInputError(
f"{len(missing)} chunk audio file(s) missing, first: {missing[0]} "
"(`say` writes chunks/<chunk>.16k.wav when it decodes the audio)"
)
return refs
def _read_settings(say_dir: Path) -> dict[str, object]:
path = say_dir / SAY_SETTINGS
if not path.is_file():
return {}
try:
settings = json.loads(path.read_text(encoding="utf-8"))
except json.JSONDecodeError as exc:
raise SpeakerSimInputError(f"{path}: invalid JSON ({exc})") from exc
if not isinstance(settings, dict):
raise SpeakerSimInputError(f"{path}: expected a JSON object")
return settings
def _string_map(settings: Mapping[str, object], key: str, where: Path) -> dict[str, str]:
value = settings.get(key, {})
if not isinstance(value, dict) or not all(
isinstance(k, str) and isinstance(v, str) for k, v in value.items()
):
raise SpeakerSimInputError(f"{where}: {key!r} must map ids to strings")
return dict(value)
def read_external_prompts(say_dir: Path) -> dict[str, ExternalPromptInfo]:
"""The external prompts recorded in ``settings.json`` (``{}`` if the file or the key is
absent: a run with dataset prompts only, or an older run).
Raises:
SpeakerSimInputError: unparseable settings, a malformed mapping, a transcript or hash
for an id that is not an external prompt, or a hash that is not 64 hex digits.
"""
path = say_dir / SAY_SETTINGS
settings = _read_settings(say_dir)
sources = _string_map(settings, "external_prompts", path)
texts = _string_map(settings, "external_prompt_texts", path)
digests = _string_map(settings, "external_prompt_sha256", path)
stray = sorted((set(texts) | set(digests)) - set(sources))
if stray:
raise SpeakerSimInputError(
f"{path}: transcript / hash for unknown external prompts {stray}"
)
for pid, digest in digests.items():
if len(digest) != 64 or set(digest) - set("0123456789abcdef"):
raise SpeakerSimInputError(f"{path}: external_prompt_sha256[{pid!r}] is not a sha256")
return {
pid: ExternalPromptInfo(src, texts.get(pid), digests.get(pid))
for pid, src in sources.items()
}
def read_say_context(say_dir: Path) -> dict[str, object]:
"""Checkpoint, step, ``w_emo`` and seed of the say run, when ``settings.json`` has them."""
settings = _read_settings(say_dir)
return {k: settings[k] for k in ("checkpoint", "step", "w_emo", "seed") if k in settings}
def prompt_id_for(wav: Path, text: str) -> str:
"""The id ``nida say`` gives an external prompt: ``ext-`` + 12 hex of a sha1 over the
resolved wav path and the (normalised) transcript. A copy of
``nida.infer.external_prompt._prompt_id`` (that module needs torch; a test pins equality).
The id depends on the path: a moved recording gets another id, so prompts are matched to
recordings by transcript and content hash (:func:`resolve_references`), not by this id alone.
"""
digest = hashlib.sha1(f"{wav.resolve()}\x1f{text}".encode()).hexdigest()
return f"ext-{digest[:12]}"
def read_prompt_list(path: Path) -> list[PromptListEntry]:
"""``<wav>\\t<transcript>`` lines (blank lines and ``#`` comments skipped), with prompt ids.
Same format and normalisation as ``nida.infer.external_prompt.read_prompt_list`` (a test
pins equality of the ids).
Raises:
SpeakerSimInputError: a line without a tab, an empty transcript, a missing wav, an
empty list, or two lines with the same prompt id.
"""
entries: list[PromptListEntry] = []
for number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
if not line.strip() or line.lstrip().startswith("#"):
continue
wav, sep, text = line.partition("\t")
if not sep or not text.strip():
raise SpeakerSimInputError(f"{path}:{number}: expected '<wav>\\t<transcript>'")
wav_path = Path(wav.strip())
if not wav_path.is_file():
raise SpeakerSimInputError(f"{path}:{number}: {wav_path} does not exist")
transcript = " ".join(nfc(text).split())
entries.append(PromptListEntry(prompt_id_for(wav_path, transcript), wav_path, transcript))
if not entries:
raise SpeakerSimInputError(f"{path}: no prompt lines")
ids = [e.prompt_id for e in entries]
if len(set(ids)) != len(ids):
raise SpeakerSimInputError(f"{path}: the same (wav, transcript) is listed twice")
return entries
def _match_prompt_list(
pid: str, info: ExternalPromptInfo | None, prompt_list: Sequence[PromptListEntry]
) -> int | None:
"""Index of the prompt-list entry for external prompt ``pid``, or ``None``.
By id first (the recording did not move); else by transcript, narrowed by the recorded
content hash when several lines share the transcript.
Raises:
SpeakerSimInputError: ambiguous transcript match, or a matched file whose content
differs from the recording ``say`` used.
"""
found = [i for i, e in enumerate(prompt_list) if e.prompt_id == pid]
if not found and info is not None and info.text is not None:
found = [i for i, e in enumerate(prompt_list) if e.transcript == info.text]
if len(found) > 1 and info.sha256 is not None:
found = [i for i in found if sha256_file(prompt_list[i].wav_path) == info.sha256]
if len(found) > 1:
raise SpeakerSimInputError(
f"prompt {pid}: {len(found)} prompt-list lines have its transcript and cannot "
"be told apart (no recorded recording hash)"
)
if not found:
return None
index = found[0]
entry = prompt_list[index]
if info is not None and info.sha256 is not None and sha256_file(entry.wav_path) != info.sha256:
raise SpeakerSimInputError(
f"prompt {pid}: {entry.wav_path} differs from the recording `say` used (sha256)"
)
return index
def _source_for(
pid: str,
info: ExternalPromptInfo | None,
prompt_list: Sequence[PromptListEntry],
matched: set[int],
) -> Path:
index = _match_prompt_list(pid, info, prompt_list)
if index is not None:
matched.add(index)
return prompt_list[index].wav_path
if info is None:
raise SpeakerSimInputError(
f"prompt {pid}: an external prompt whose source recording is not recorded in "
"settings.json; pass the --prompt-list used for the run, or --reference decoded"
)
source = Path(info.source)
if not source.is_absolute():
raise SpeakerSimInputError(
f"prompt {pid}: the recorded source path {info.source!r} is relative, so it cannot "
"be found reliably; pass the --prompt-list used for the run, or --reference decoded"
)
if not source.is_file():
raise SpeakerSimInputError(
f"prompt {pid}: source recording {source} does not exist (pass the --prompt-list "
"used for the run, or --reference decoded)"
)
if info.sha256 is not None and sha256_file(source) != info.sha256:
raise SpeakerSimInputError(
f"prompt {pid}: {source} differs from the recording `say` used (sha256)"
)
return source
def _decoded_reference(say_dir: Path, pid: str, chunks: Sequence[ChunkRef]) -> ReferenceRef:
first = min(chunks, key=lambda c: c.index)
if first.prompt_frames is None or first.num_frames is None:
raise SpeakerSimInputError(
f"prompt {pid}: say.jsonl has no prompt_frames / num_frames (an older run), so the "
"decoded prompt cannot be located; use --reference source"
)
if not 0 < first.prompt_frames <= first.num_frames:
raise SpeakerSimInputError(
f"prompt {pid}: bad frame counts {first.prompt_frames} / {first.num_frames}"
)
audio = say_dir / "audio" / f"{first.chunk_id}.16k.wav"
if not audio.is_file():
raise SpeakerSimInputError(
f"prompt {pid}: decoded audio {audio} is missing (`say` writes audio/<chunk>.16k.wav)"
)
return ReferenceRef("decoded", audio, first.prompt_frames, first.num_frames)
def resolve_references(
say_dir: Path,
chunks: Sequence[ChunkRef],
*,
mode: ReferenceMode,
external: Mapping[str, ExternalPromptInfo],
prompt_list: Sequence[PromptListEntry] = (),
) -> tuple[dict[str, ReferenceRef], set[int]]:
"""The reference audio of every prompt of a say run, and the prompt-list lines it used.
An *external* prompt is one recorded in ``settings.json`` or with an ``ext-`` id. In
``auto`` / ``source`` mode it is compared with its **source recording**, found in
``prompt_list`` (by id, else by transcript, narrowed and verified by the recorded content
hash) or at the absolute path recorded in ``settings.json`` (verified by the hash when
recorded). A dataset prompt has no recording in the say directory: ``auto`` compares it with
the **decoded prompt** (the prompt part of ``audio/<chunk>.16k.wav``, cut like the chunks),
``source`` refuses it. ``decoded`` uses the decoded prompt for every prompt. An external
prompt whose source cannot be found is an error in ``auto`` and ``source`` mode, never a
silent fall-back to the decoded prompt.
Returns:
``(prompt id -> reference, indices into prompt_list that matched a prompt)``.
Raises:
SpeakerSimInputError: a needed file or frame count is missing, a recorded path is
relative, a hash differs, a transcript match is ambiguous, or ``source`` mode meets a
dataset prompt.
"""
by_prompt: dict[str, list[ChunkRef]] = defaultdict(list)
for c in chunks:
by_prompt[c.prompt_id].append(c)
refs: dict[str, ReferenceRef] = {}
matched: set[int] = set()
no_source: list[str] = []
for pid in sorted(by_prompt):
info = external.get(pid)
is_external = info is not None or pid.startswith("ext-")
if mode != "decoded" and is_external:
refs[pid] = ReferenceRef("source", _source_for(pid, info, prompt_list, matched))
elif mode == "source":
no_source.append(pid)
else:
refs[pid] = _decoded_reference(say_dir, pid, by_prompt[pid])
if no_source:
raise SpeakerSimInputError(
f"--reference source: no source recording known for prompt(s) {no_source} "
"(dataset prompts have none; use --reference auto or decoded)"
)
return refs, matched
def unmatched_prompt_entries(
prompt_list: Sequence[PromptListEntry], matched: set[int]
) -> list[PromptListEntry]:
"""The prompt-list lines that no prompt of any say run matched (``matched`` is the union of
the index sets :func:`resolve_references` returned): a stale or wrong list."""
return [e for i, e in enumerate(prompt_list) if i not in matched]
# ----------------------------------------------------------------------------- say: scoring
def score_say(
chunks: Sequence[ChunkRef],
references: Mapping[str, ReferenceRef],
reference_outcomes: Mapping[str, EmbedOutcome],
chunk_outcomes: Mapping[str, EmbedOutcome],
threshold: float,
) -> list[dict[str, object]]:
"""One record per chunk: similarity to its prompt's reference, or the reason it has none.
``reference_outcomes`` is keyed by prompt id, ``chunk_outcomes`` by chunk id. A record
is ``{chunk, prompt_id, index, reference_kind, similarity, below_threshold, duration_ms}``
or ``{chunk, prompt_id, index, reference_kind, error}`` (the chunk's own failure, or the
failure of its reference). ``below_threshold`` is ``similarity < threshold``.
Raises:
SpeakerSimInputError: an outcome is missing for a chunk or prompt.
"""
records: list[dict[str, object]] = []
for chunk in chunks:
head: dict[str, object] = {
"chunk": chunk.chunk_id,
"prompt_id": chunk.prompt_id,
"index": chunk.index,
"reference_kind": references[chunk.prompt_id].kind,
}
ref = reference_outcomes.get(chunk.prompt_id)
out = chunk_outcomes.get(chunk.chunk_id)
if ref is None or out is None:
raise SpeakerSimInputError(f"{chunk.chunk_id}: no embedding outcome")
if isinstance(ref, ClipFailure):
records.append({**head, "error": f"reference embedding failed: {ref.error}"})
elif isinstance(out, ClipFailure):
records.append({**head, "error": out.error})
else:
sim = cosine(ref.vector, out.vector)
records.append(
{
**head,
"similarity": sim,
"below_threshold": sim < threshold,
"duration_ms": out.duration_ms,
}
)
return records
def _similarity_stats(
records: Sequence[Mapping[str, object]], threshold: float
) -> dict[str, object]:
sims = [float(str(r["similarity"])) for r in records if "similarity" in r]
stats = describe(sims)
below = sum(1 for s in sims if s < threshold)
return {
"n_scored": len(sims),
"mean": stats["mean"],
"min": stats["min"],
"max": stats["max"],
"std": stats["std"],
"spread": (max(sims) - min(sims)) if sims else None,
"n_below_threshold": below,
"drift_fraction": (below / len(sims)) if sims else None,
}
def summarize_say(
records: Sequence[Mapping[str, object]],
references: Mapping[str, ReferenceRef],
reference_outcomes: Mapping[str, EmbedOutcome],
threshold: float,
) -> dict[str, object]:
"""Per-reference-kind and per-prompt statistics of :func:`score_say` records.
``overall_by_reference_kind`` maps ``"source"`` / ``"decoded"`` (the kinds that occur) to the
statistics of the chunks compared with that kind of reference. They are never pooled into one
number: a similarity to the original recording and one to the VAE-decoded prompt are not on the
same scale. Those entries and every ``per_prompt`` entry hold: ``n_scored``; ``mean`` /
``min`` / ``max`` / ``std`` of the similarity; ``spread`` (max - min across the chunks: the
within-paragraph spread); ``n_below_threshold`` and ``drift_fraction`` (the share of scored
chunks below ``threshold``). A kind adds ``n_chunks`` and ``n_failed``; a prompt adds
``reference_kind``, ``reference_path``, ``n_chunks``, ``n_failed`` and
``reference_duration_ms`` (``None`` if its reference failed). Failed chunks are listed with
their error, never dropped; ``lowest_chunks`` are the weakest few (with their kind).
"""
per_prompt: dict[str, dict[str, object]] = {}
by_prompt: dict[str, list[Mapping[str, object]]] = defaultdict(list)
by_kind: dict[str, list[Mapping[str, object]]] = defaultdict(list)
for record in records:
by_prompt[str(record["prompt_id"])].append(record)
by_kind[str(record["reference_kind"])].append(record)
for pid in sorted(by_prompt):
mine = by_prompt[pid]
ref = reference_outcomes.get(pid)
per_prompt[pid] = {
"reference_kind": references[pid].kind,
"reference_path": str(references[pid].wav_path),
"reference_duration_ms": ref.duration_ms if isinstance(ref, ClipEmbedding) else None,
"n_chunks": len(mine),
"n_failed": sum(1 for r in mine if "error" in r),
**_similarity_stats(mine, threshold),
}
overall = {
kind: {
"n_chunks": len(mine),
"n_failed": sum(1 for r in mine if "error" in r),
**_similarity_stats(mine, threshold),
}
for kind, mine in sorted(by_kind.items())
}
failed = [{"chunk": r["chunk"], "error": r["error"]} for r in records if "error" in r]
scored = sorted(
(r for r in records if "similarity" in r), key=lambda r: float(str(r["similarity"]))
)
return {
"threshold": threshold,
"n_chunks": len(records),
"n_failed": len(failed),
"failed_chunks": failed,
"overall_by_reference_kind": overall,
"per_prompt": per_prompt,
"lowest_chunks": [
{
"chunk": r["chunk"],
"similarity": r["similarity"],
"reference_kind": r["reference_kind"],
}
for r in scored[:LOWEST_CHUNKS]
],
}
# ---------------------------------------------------------------------------- groups: input
@dataclass(frozen=True)
class IndexRow:
"""The fields of a latent dataset ``index.jsonl`` row that the group check needs."""
segment_id: str
group: str
split: str
num_frames: int
def read_index(dataset_dir: Path) -> tuple[list[IndexRow], float]:
"""Rows of ``<dataset_dir>/index.jsonl`` and the latent frame hop (ms) from ``meta.json``.
Raises:
SpeakerSimInputError: a missing / empty / malformed index, a duplicate segment id, or
a ``meta.json`` without a positive ``frame_hop_ms``.
"""
index_path, meta_path = dataset_dir / "index.jsonl", dataset_dir / "meta.json"
if not index_path.is_file() or not meta_path.is_file():
raise SpeakerSimInputError(
f"{dataset_dir}: needs index.jsonl and meta.json (a latent dataset directory)"
)
try:
meta = json.loads(meta_path.read_text(encoding="utf-8"))
except json.JSONDecodeError as exc:
raise SpeakerSimInputError(f"{meta_path}: invalid JSON ({exc})") from exc
hop = meta.get("frame_hop_ms") if isinstance(meta, dict) else None
if isinstance(hop, bool) or not isinstance(hop, int | float) or not hop > 0:
raise SpeakerSimInputError(f"{meta_path}: frame_hop_ms must be a positive number")
rows: list[IndexRow] = []
seen: set[str] = set()
for line_no, line in enumerate(index_path.read_text(encoding="utf-8").splitlines(), start=1):
if not line.strip():
continue
where = f"{index_path} line {line_no}"
try:
record = json.loads(line)
except json.JSONDecodeError as exc:
raise SpeakerSimInputError(f"{where}: invalid JSON ({exc})") from exc
if not isinstance(record, dict):
raise SpeakerSimInputError(f"{where}: expected a JSON object")
sid, group, split = record.get("segment_id"), record.get("group"), record.get("split")
frames = record.get("num_frames")
if not isinstance(sid, str) or not sid or not isinstance(group, str) or not group:
raise SpeakerSimInputError(f"{where}: bad segment_id / group")
if not isinstance(split, str) or isinstance(frames, bool) or not isinstance(frames, int):
raise SpeakerSimInputError(f"{where}: bad split / num_frames")
if sid in seen:
raise SpeakerSimInputError(f"{where}: duplicate segment {sid!r}")
seen.add(sid)
rows.append(IndexRow(sid, group, split, frames))
if not rows:
raise SpeakerSimInputError(f"{index_path}: no rows")
return rows, float(hop)
def find_dataset_dir(root: Path) -> Path:
"""The one finished ``dataset-*`` directory (with ``meta.json``, written last) of an importer
root.
Raises:
SpeakerSimInputError: none, or more than one (pass the directory explicitly).
"""
found = sorted(p for p in root.glob("dataset-*") if (p / "meta.json").is_file())
if len(found) != 1:
raise SpeakerSimInputError(
f"{root}: expected exactly one finished dataset-*/ directory, found "
f"{[p.name for p in found]} (use --dataset-dir)"
)
return found[0]
def list_audio_stems(audio_root: Path, suffix: str = ".wav") -> set[str]:
"""File stems (segment ids) of ``<audio_root>/*<suffix>``, from one directory listing.
Raises:
SpeakerSimInputError: if ``audio_root`` is not a directory.
"""
if not audio_root.is_dir():
raise SpeakerSimInputError(
f"{audio_root}: not a directory (the importer's 16 kHz wavs live in its local "
"cache, <cache>/audio16k; rerun `nida data-import --stage audio` to rebuild them)"
)
with os.scandir(audio_root) as entries:
return {e.name[: -len(suffix)] for e in entries if e.name.endswith(suffix) and e.is_file()}
def select_candidates(
rows: Sequence[IndexRow],
available: set[str],
*,
hop_ms: float,
min_duration_ms: float,
split: str | None = None,
) -> tuple[dict[str, list[str]], dict[str, int]]:
"""Segment ids per group that can be embedded, and why the others cannot.
A row is dropped (and counted, never silently) for: another ``split`` than requested,
``num_frames * hop_ms < min_duration_ms`` (``too_short``: an embedding of a very short
clip is unreliable), or no wav in ``available`` (``missing_audio``).
Returns:
``(group -> sorted segment ids, counts)`` with the counts ``n_rows``, ``n_other_split``,
``n_too_short``, ``n_missing_audio`` and ``n_candidates``.
"""
counts = {
"n_rows": len(rows),
"n_other_split": 0,
"n_too_short": 0,
"n_missing_audio": 0,
"n_candidates": 0,
}
groups: dict[str, list[str]] = defaultdict(list)
for row in rows:
if split is not None and row.split != split:
counts["n_other_split"] += 1
elif row.num_frames * hop_ms < min_duration_ms:
counts["n_too_short"] += 1
elif row.segment_id not in available:
counts["n_missing_audio"] += 1
else:
groups[row.group].append(row.segment_id)
counts["n_candidates"] += 1
return {g: sorted(ids) for g, ids in sorted(groups.items())}, counts
def sample_groups(
candidates: Mapping[str, Sequence[str]], *, n_groups: int, per_group: int, seed: int
) -> list[tuple[str, tuple[str, ...]]]:
"""``n_groups`` groups with at least ``per_group`` candidates, ``per_group`` segments each.
Deterministic: a pure function of the candidate sets, ``seed`` and the arguments (groups and
segments are sorted before sampling). Fewer eligible groups than ``n_groups`` yields all of
them. The result is ordered by group id, segments by id.
Raises:
ValueError: ``n_groups < 1`` or ``per_group < 2`` (one segment has no intra pair).
"""
if n_groups < 1:
raise ValueError(f"n_groups must be >= 1, got {n_groups}")
if per_group < 2:
raise ValueError(f"per_group must be >= 2, got {per_group}")
rng = random.Random(seed)
eligible = sorted(g for g, ids in candidates.items() if len(set(ids)) >= per_group)
chosen = sorted(rng.sample(eligible, min(n_groups, len(eligible))))
return [(g, tuple(sorted(rng.sample(sorted(set(candidates[g])), per_group)))) for g in chosen]
# --------------------------------------------------------------------------- groups: scoring
def _pair_similarities(vectors: Sequence[FloatVec]) -> list[float]:
matrix = np.stack(vectors).astype(np.float64)
gram = matrix @ matrix.T
upper = np.triu_indices(len(vectors), k=1)
return [float(x) for x in gram[upper]]
def _auc(positives: Sequence[float], negatives: Sequence[float]) -> float:
"""P(positive > negative) with ties counting half (the area under the ROC curve)."""
neg = np.sort(np.asarray(negatives, dtype=np.float64))
pos = np.asarray(positives, dtype=np.float64)
below = np.searchsorted(neg, pos, side="left")
at_most = np.searchsorted(neg, pos, side="right")
return float(np.mean((below + 0.5 * (at_most - below)) / neg.size))
def _inter_similarities(
embedded: Mapping[str, Sequence[tuple[str, FloatVec]]], pairs: int, seed: int
) -> list[float]:
"""Cosine of ``pairs`` random segment pairs drawn from two different groups."""
groups = sorted(embedded)
if len(groups) < 2:
return []
rng = random.Random(seed)
sims: list[float] = []
for _ in range(pairs):
first, second = rng.sample(groups, 2)
a = rng.choice(embedded[first])[1].astype(np.float64)
b = rng.choice(embedded[second])[1].astype(np.float64)
sims.append(float(np.dot(a, b)))
return sims
def summarize_groups(
samples: Sequence[tuple[str, tuple[str, ...]]],
outcomes: Mapping[str, EmbedOutcome],
*,
intra_threshold: float,
inter_pairs: int,
seed: int,
) -> tuple[list[dict[str, object]], dict[str, object]]:
"""Per-group records and the intra- vs inter-group summary of the sampled segments.
Per group (``outcomes`` is keyed by segment id): ``group``, ``segments``, ``n_embedded``,
``errors`` (segment -> error, never dropped) and, when at least two segments were embedded,
``intra_mean`` / ``intra_min`` (cosine over all pairs) and ``below_threshold``
(``intra_mean < intra_threshold``).
Summary: ``intra.pairs`` / ``intra.group_means`` / ``inter`` (:func:`describe`), the counts,
``n_groups_below`` / ``fraction_groups_below`` (over groups with an intra mean),
``separation`` (mean intra pair minus mean inter pair) and ``auc`` (probability that a random
intra pair is more similar than a random inter pair). ``inter`` uses ``inter_pairs``
random pairs of embedded segments from two different groups (seeded; empty stats if fewer
than two groups have an embedded segment).
Raises:
ValueError: ``inter_pairs < 1``.
SpeakerSimInputError: a sampled segment has no outcome.
"""
if inter_pairs < 1:
raise ValueError(f"inter_pairs must be >= 1, got {inter_pairs}")
records: list[dict[str, object]] = []
embedded: dict[str, list[tuple[str, FloatVec]]] = {}
intra_pair_sims: list[float] = []
group_means: list[float] = []
n_failed = 0
for group, segments in samples:
good: list[tuple[str, FloatVec]] = []
errors: dict[str, str] = {}
for sid in segments:
outcome = outcomes.get(sid)
if outcome is None:
raise SpeakerSimInputError(f"{sid}: no embedding outcome")
if isinstance(outcome, ClipFailure):
errors[sid] = outcome.error
else:
good.append((sid, outcome.vector))
n_failed += len(errors)
record: dict[str, object] = {
"group": group,
"segments": list(segments),
"n_embedded": len(good),
"errors": errors,
}
if good:
embedded[group] = good
if len(good) >= 2:
sims = _pair_similarities([v for _, v in good])
mean = float(np.mean(sims))
record["intra_mean"] = mean
record["intra_min"] = min(sims)
record["below_threshold"] = mean < intra_threshold
intra_pair_sims.extend(sims)
group_means.append(mean)
records.append(record)
inter_sims = _inter_similarities(embedded, inter_pairs, seed)
below = sum(1 for m in group_means if m < intra_threshold)
intra_mean_all = float(np.mean(intra_pair_sims)) if intra_pair_sims else None
inter_mean_all = float(np.mean(inter_sims)) if inter_sims else None
summary: dict[str, object] = {
"n_groups_sampled": len(samples),
"n_groups_with_intra": len(group_means),
"n_segments_failed": n_failed,
"intra_threshold": intra_threshold,
"n_groups_below": below,
"fraction_groups_below": (below / len(group_means)) if group_means else None,
"intra": {"pairs": describe(intra_pair_sims), "group_means": describe(group_means)},
"inter": describe(inter_sims),
"separation": (
intra_mean_all - inter_mean_all
if intra_mean_all is not None and inter_mean_all is not None
else None
),
"auc": _auc(intra_pair_sims, inter_sims) if intra_pair_sims and inter_sims else None,
}
return records, summary
"""Speaker-similarity measurements on top of the ECAPA-TDNN embedding worker.
Two uses, driven by ``scripts/speaker_sim.py``:
* **say**: cosine similarity of every generated chunk of a ``nida say`` output directory to
the embedding of its speaker prompt (:func:`read_say_chunks`, :func:`resolve_references`,
:func:`score_say`, :func:`summarize_say`);
* **groups**: does a dataset ``group`` (speaker / source video) really hold one voice? Intra-
group versus inter-group cosine similarity of sampled segments (:func:`read_index`,
:func:`select_candidates`, :func:`sample_groups`, :func:`summarize_groups`).
Embedding is a side effect at the edge (:func:`embed_with_adapter`: the worker protocol with
its resumable results file, run in-process); everything else is pure and needs only numpy.
Embeddings are 192-d float32 unit vectors, so a cosine similarity is a dot product. Thresholds
are arguments (the script's flags), never constants of the measurement.
``say`` layout relied upon (``nida.infer.say``): ``say.jsonl`` (``prompt_id``, ``index`` per
chunk), ``chunks/<prompt_id>__<index:02d>.16k.wav`` (the 16 kHz generated chunk),
``settings.json`` (``external_prompts``: external prompt id -> source wav path as given to
``--prompt-list``, with ``external_prompt_texts`` and ``external_prompt_sha256`` beside it in newer
runs; all absent or empty for dataset prompts) and ``audio/<chunk>.16k.wav`` (the decoder's 16 kHz
output for ``prompt + sentence``: ``say`` cuts the generated chunk from it with
``prompt_frames`` / ``num_frames`` of ``say.jsonl``, and the same cut's first part is the
decoded prompt).
"""
from __future__ import annotations
import hashlib
import io
import json
import logging
import math
import os
import random
from collections import defaultdict
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from math import gcd
from pathlib import Path
from typing import Literal
import numpy as np
import numpy.typing as npt
from nida.bench.protocol import (
WorkRequest,
drop_results,
load_requests,
load_results,
read_summary,
write_requests,
)
from nida.bench.workers.runner import Adapter, ModelRef, WorkerOptions, run_worker
from nida.data.workers.spk_ecapa import EMBEDDING_DIM, sha256_file
from nida.text.turkish import nfc
from nida.utils.io import atomic_write_bytes
__all__ = [
"DEFAULT_DRIFT_THRESHOLD",
"ChunkRef",
"ClipEmbedding",
"ClipFailure",
"EmbedOutcome",
"Embedder",
"ExternalPromptInfo",
"IndexRow",
"PromptListEntry",
"ReferenceMode",
"ReferenceRef",
"SpeakerSimInputError",
"cosine",
"cut_prompt_audio",
"describe",
"embed_with_adapter",
"find_dataset_dir",
"list_audio_stems",
"prompt_id_for",
"read_external_prompts",
"read_index",
"read_prompt_list",
"read_say_chunks",
"read_say_context",
"resolve_references",
"sample_groups",
"score_say",
"select_candidates",
"summarize_groups",
"summarize_say",
"unmatched_prompt_entries",
"validate_group_options",
"validate_threshold",
"write_16k_mono_wav",
"write_decoded_prompt_wav",
"write_if_changed",
]
log = logging.getLogger(__name__)
#: Default for ``--threshold`` of the ``say`` subcommand (the user's choice, 2026-10-15).