personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4"""Differential harness for the local speaker pipeline.
5
6The bundle schema says ``statements`` everywhere because
7``solstone.observe.transcribe.main`` owns the production pipeline vocabulary.
8At the single boundary into ``solstone.observe.transcribe.diarize``, those same
9records are passed as that module's local ``sentences`` parameter.
10"""
11
12from __future__ import annotations
13
14import argparse
15import contextlib
16import hashlib
17import importlib
18import json
19import logging
20import platform
21import sys
22from collections.abc import Iterator, Sequence
23from datetime import UTC, datetime
24from pathlib import Path
25from typing import Any
26
27import numpy as np
28
29from solstone.observe.transcribe import diarize, overlap
30from solstone.observe.vad import AudioReduction, restore_statement_timestamps
31from solstone.think.utils import get_rev
32from tests._speaker_differential_fixtures import (
33 COMPARATOR_THRESHOLDS,
34 COSINE_NORM_FLOOR,
35 EMBEDDING_MAX_ABS_TOLERANCE,
36 EMBEDDING_MEDIAN_COSINE_SIMILARITY,
37 EMBEDDING_MIN_COSINE_SIMILARITY,
38 EVIDENCE_FLOAT_ABS_TOLERANCE,
39 LOGPROB_ARGMAX_AGREEMENT_MIN,
40 LOGPROB_MAX_ABS_TOLERANCE,
41 LOGPROB_MEDIAN_ABS_TOLERANCE,
42 SAMPLE_RATE,
43 STATEMENT_DURATION_ABS_TOLERANCE,
44)
45
46transcribe_main = importlib.import_module("solstone.observe.transcribe.main")
47
48logger = logging.getLogger(__name__)
49
50ROOT = Path(__file__).resolve().parent.parent
51BUNDLE_SCHEMA = "solstone-speaker-differential-bundle"
52REPORT_SCHEMA = "solstone-speaker-differential-report"
53SCHEMA_VERSION = 1
54MANIFEST_KEY = "__speaker_differential_manifest_json__"
55PLACEHOLDER_WAV_PATH = Path("__speaker_differential_audio_supplied_no_wav_read.wav")
56
57PRESENT = "present"
58ABSENT_NONE = "absent-none"
59NOT_EVALUATED = "not-evaluated"
60FIELD_STATES = frozenset({PRESENT, ABSENT_NONE, NOT_EVALUATED})
61
62EQUAL = "equal"
63FUNCTIONALLY_EQUAL = "functionally-equal"
64UNEXPECTED_DIFFERS = "unexpected-differs"
65HARNESS_ERROR = "harness_error"
66
67COMPARE = "compare"
68STATE_PAIR_VERDICT = {
69 (PRESENT, PRESENT): COMPARE,
70 (PRESENT, ABSENT_NONE): UNEXPECTED_DIFFERS,
71 (PRESENT, NOT_EVALUATED): UNEXPECTED_DIFFERS,
72 (ABSENT_NONE, PRESENT): UNEXPECTED_DIFFERS,
73 (ABSENT_NONE, ABSENT_NONE): NOT_EVALUATED,
74 (ABSENT_NONE, NOT_EVALUATED): UNEXPECTED_DIFFERS,
75 (NOT_EVALUATED, PRESENT): UNEXPECTED_DIFFERS,
76 (NOT_EVALUATED, ABSENT_NONE): UNEXPECTED_DIFFERS,
77 (NOT_EVALUATED, NOT_EVALUATED): NOT_EVALUATED,
78}
79
80LABEL_NULL_SENTINEL = np.int32(-1)
81
82INPUT_STATEMENT_IDS = "inputs.statement_embedding.statement_ids"
83INPUT_STATEMENT_SPANS = "inputs.statement_embedding.spans_s"
84INPUT_DIARIZATION_IDS = "inputs.diarization.statement_ids"
85INPUT_DIARIZATION_SPANS = "inputs.diarization.spans_s"
86
87PYANNOTE_LOGPROBS = "pyannote.avg_log_probs"
88PYANNOTE_WINDOW_STATS = "pyannote.window_stats"
89
90EVIDENCE_SPEAKER = "evidence.speaker_evidence"
91EVIDENCE_MULTI_FRACTION = "evidence.multi_window_fraction"
92EVIDENCE_MEAN_OVERLAP = "evidence.mean_window_overlap_share"
93EVIDENCE_OVERLAP_FRACTION = "evidence.overlap_fraction"
94
95STATEMENT_EMBEDDINGS = "statement_embeddings.embeddings"
96STATEMENT_EMBEDDING_IDS = "statement_embeddings.statement_ids"
97STATEMENT_DURATIONS = "statement_embeddings.durations_s"
98STATEMENT_ENCODER = "statement_embeddings.encoder"
99
100DIARIZATION_INTERVALS = "diarization.intervals"
101DIARIZATION_VALID_INTERVALS = "diarization.valid_intervals"
102DIARIZATION_INTERVAL_EMBEDDINGS = "diarization.interval_embeddings"
103DIARIZATION_CLUSTER_LABELS = "diarization.cluster_labels"
104DIARIZATION_STATEMENT_LABELS = "diarization.statement_labels"
105DIARIZATION_SILHOUETTE_K = "diarization.silhouette_k"
106DIARIZATION_EFFECTIVE_K = "diarization.effective_k"
107
108GATE_DECLINED_FIELDS = frozenset(
109 {
110 DIARIZATION_INTERVALS,
111 DIARIZATION_VALID_INTERVALS,
112 DIARIZATION_INTERVAL_EMBEDDINGS,
113 DIARIZATION_CLUSTER_LABELS,
114 DIARIZATION_STATEMENT_LABELS,
115 DIARIZATION_SILHOUETTE_K,
116 DIARIZATION_EFFECTIVE_K,
117 }
118)
119
120COMPONENT_FIELDS = {
121 "pyannote_log_probs": (PYANNOTE_LOGPROBS,),
122 "evidence": (
123 EVIDENCE_SPEAKER,
124 EVIDENCE_MULTI_FRACTION,
125 EVIDENCE_MEAN_OVERLAP,
126 EVIDENCE_OVERLAP_FRACTION,
127 PYANNOTE_WINDOW_STATS,
128 ),
129 "intervals": (DIARIZATION_INTERVALS,),
130 "interval_embeddings": (
131 DIARIZATION_VALID_INTERVALS,
132 DIARIZATION_INTERVAL_EMBEDDINGS,
133 ),
134 "statement_embeddings": (
135 STATEMENT_EMBEDDINGS,
136 STATEMENT_EMBEDDING_IDS,
137 STATEMENT_DURATIONS,
138 STATEMENT_ENCODER,
139 ),
140 "clustering": (
141 DIARIZATION_CLUSTER_LABELS,
142 DIARIZATION_SILHOUETTE_K,
143 DIARIZATION_EFFECTIVE_K,
144 ),
145 "statement_labels": (DIARIZATION_STATEMENT_LABELS,),
146}
147
148COMPONENT_ORDER = (
149 "pyannote_log_probs",
150 "evidence",
151 "statement_embeddings",
152 "intervals",
153 "interval_embeddings",
154 "clustering",
155 "statement_labels",
156)
157
158PYANNOTE_THRESHOLD_IDS = {
159 "max_abs_diff": "LOGPROB_MAX_ABS_TOLERANCE",
160 "median_abs_diff": "LOGPROB_MEDIAN_ABS_TOLERANCE",
161 "per_frame_argmax_agreement_fraction": "LOGPROB_ARGMAX_AGREEMENT_MIN",
162}
163
164EVIDENCE_THRESHOLD_IDS = {
165 EVIDENCE_OVERLAP_FRACTION: "EVIDENCE_FLOAT_ABS_TOLERANCE",
166 EVIDENCE_MULTI_FRACTION: "EVIDENCE_FLOAT_ABS_TOLERANCE",
167 EVIDENCE_MEAN_OVERLAP: "EVIDENCE_FLOAT_ABS_TOLERANCE",
168 EVIDENCE_SPEAKER: "exact",
169 PYANNOTE_WINDOW_STATS: "exact",
170}
171
172STATEMENT_EMBEDDING_THRESHOLD_IDS = {
173 "max_abs_component_diff": "EMBEDDING_MAX_ABS_TOLERANCE",
174 "min_cosine_similarity": "EMBEDDING_MIN_COSINE_SIMILARITY",
175 "median_cosine_similarity": "EMBEDDING_MEDIAN_COSINE_SIMILARITY",
176 "duration_max_abs_diff": "STATEMENT_DURATION_ABS_TOLERANCE",
177 STATEMENT_ENCODER: "exact",
178}
179
180INTERVAL_EMBEDDING_THRESHOLD_IDS = {
181 "max_abs_component_diff": "EMBEDDING_MAX_ABS_TOLERANCE",
182 "min_cosine_similarity": "EMBEDDING_MIN_COSINE_SIMILARITY",
183 "median_cosine_similarity": "EMBEDDING_MEDIAN_COSINE_SIMILARITY",
184 DIARIZATION_VALID_INTERVALS: "exact",
185}
186
187
188class HarnessError(RuntimeError):
189 """Raised when the harness cannot produce or compare trustworthy data."""
190
191
192Bundle = dict[str, Any]
193
194
195def _json_default(value: object) -> object:
196 if isinstance(value, Path):
197 return str(value)
198 if isinstance(value, np.generic):
199 return value.item()
200 raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")
201
202
203def _render_report(report: dict[str, Any]) -> str:
204 return json.dumps(report, default=_json_default, indent=2, sort_keys=True) + "\n"
205
206
207def _session_providers(session: object | None) -> list[str] | None:
208 if session is None:
209 return None
210 get_providers = getattr(session, "get_providers", None)
211 if get_providers is None:
212 return None
213 try:
214 return list(get_providers())
215 except Exception:
216 logger.exception("failed to read ONNX execution providers")
217 return None
218
219
220def _module_version(module_name: str) -> str | None:
221 try:
222 module = importlib.import_module(module_name)
223 except Exception:
224 logger.exception("failed to import %s for provenance", module_name)
225 return None
226 version = getattr(module, "__version__", None)
227 return str(version) if version is not None else None
228
229
230def _provenance(producer: str) -> dict[str, Any]:
231 return {
232 "generated_at": datetime.now(UTC).isoformat().replace("+00:00", "Z"),
233 "implementation": {"name": producer},
234 "harness": {
235 "name": "tests.verify_speaker_differential",
236 "repo_commit": get_rev(),
237 "schema_version": SCHEMA_VERSION,
238 },
239 "onnx_execution_providers": {
240 "statement_encoder": _session_providers(
241 getattr(transcribe_main, "_embedder_session", None)
242 ),
243 "interval_encoder": _session_providers(
244 getattr(diarize, "_wespeaker_session", None)
245 ),
246 "pyannote": _session_providers(getattr(overlap, "_overlap_session", None)),
247 },
248 "versions": {
249 "onnxruntime": _module_version("onnxruntime"),
250 "kaldi_native_fbank": _module_version("kaldi_native_fbank"),
251 },
252 "host": {
253 "platform": platform.platform(),
254 "python": platform.python_version(),
255 },
256 }
257
258
259def new_bundle(*, producer: str = "production-python") -> Bundle:
260 return {
261 "manifest": {
262 "schema": BUNDLE_SCHEMA,
263 "schema_version": SCHEMA_VERSION,
264 "fields": {},
265 "inputs": {},
266 "counters": {},
267 "stage_errors": [],
268 "provenance": _provenance(producer),
269 "label_null_sentinel": int(LABEL_NULL_SENTINEL),
270 "gate_declined_fields": sorted(GATE_DECLINED_FIELDS),
271 },
272 "arrays": {},
273 }
274
275
276def _refresh_provenance(bundle: Bundle, producer: str) -> None:
277 bundle["manifest"]["provenance"] = _provenance(producer)
278
279
280def copy_bundle(bundle: Bundle) -> Bundle:
281 return {
282 "manifest": json.loads(json.dumps(bundle["manifest"], default=_json_default)),
283 "arrays": {
284 str(key): np.array(value, copy=True)
285 for key, value in bundle["arrays"].items()
286 },
287 }
288
289
290def _set_array(bundle: Bundle, key: str, value: np.ndarray, component: str) -> None:
291 array = np.asarray(value)
292 bundle["arrays"][key] = np.array(array, copy=True)
293 bundle["manifest"]["fields"][key] = {
294 "state": PRESENT,
295 "kind": "array",
296 "component": component,
297 "dtype": str(bundle["arrays"][key].dtype),
298 "shape": list(bundle["arrays"][key].shape),
299 }
300
301
302def _set_scalar(bundle: Bundle, key: str, value: Any, component: str) -> None:
303 if isinstance(value, np.generic):
304 value = value.item()
305 bundle["manifest"]["fields"][key] = {
306 "state": PRESENT,
307 "kind": "scalar",
308 "component": component,
309 "value": value,
310 }
311
312
313def _set_state(bundle: Bundle, key: str, state: str, component: str) -> None:
314 if state not in FIELD_STATES:
315 raise HarnessError(f"invalid field state {state!r} for {key}")
316 if state == PRESENT:
317 raise HarnessError(f"{key} needs a value for present state")
318 bundle["arrays"].pop(key, None)
319 bundle["manifest"]["fields"][key] = {
320 "state": state,
321 "kind": "state",
322 "component": component,
323 }
324
325
326def replace_array(bundle: Bundle, key: str, value: np.ndarray) -> None:
327 field = _require_field(bundle, key)
328 component = str(field["component"])
329 _set_array(bundle, key, np.asarray(value), component)
330
331
332def set_scalar_value(bundle: Bundle, key: str, value: Any) -> None:
333 field = _require_field(bundle, key)
334 component = str(field["component"])
335 _set_scalar(bundle, key, value, component)
336
337
338def _require_field(bundle: Bundle, key: str) -> dict[str, Any]:
339 try:
340 field = bundle["manifest"]["fields"][key]
341 except KeyError as exc:
342 raise HarnessError(f"bundle missing manifest field {key}") from exc
343 state = field.get("state")
344 if state not in FIELD_STATES:
345 raise HarnessError(f"bundle field {key} has invalid state {state!r}")
346 return field
347
348
349def _field_state(bundle: Bundle, key: str) -> str:
350 return str(_require_field(bundle, key)["state"])
351
352
353def _array(bundle: Bundle, key: str) -> np.ndarray:
354 field = _require_field(bundle, key)
355 if field["state"] != PRESENT:
356 raise HarnessError(f"bundle field {key} is {field['state']}, not present")
357 if field.get("kind") != "array":
358 raise HarnessError(f"bundle field {key} is not an array")
359 try:
360 return bundle["arrays"][key]
361 except KeyError as exc:
362 raise HarnessError(f"bundle missing payload array {key}") from exc
363
364
365def _scalar(bundle: Bundle, key: str) -> Any:
366 field = _require_field(bundle, key)
367 if field["state"] != PRESENT:
368 raise HarnessError(f"bundle field {key} is {field['state']}, not present")
369 if field.get("kind") != "scalar":
370 raise HarnessError(f"bundle field {key} is not a scalar")
371 return field.get("value")
372
373
374def _refuse_repo_destination(path: Path) -> Path:
375 resolved = path.expanduser().resolve()
376 if resolved == ROOT or resolved.is_relative_to(ROOT):
377 raise HarnessError(
378 f"speaker differential refuses in-repo destination: {resolved}"
379 )
380 return resolved
381
382
383def write_bundle(bundle: Bundle, path: Path) -> None:
384 destination = _refuse_repo_destination(path)
385 destination.parent.mkdir(parents=True, exist_ok=True)
386 manifest = bundle["manifest"]
387 payload = {
388 MANIFEST_KEY: np.array(
389 json.dumps(manifest, default=_json_default, sort_keys=True)
390 )
391 }
392 payload.update(bundle["arrays"])
393 with destination.open("wb") as fh:
394 np.savez(fh, **payload)
395
396
397def load_bundle(path: Path) -> Bundle:
398 try:
399 with np.load(path, allow_pickle=False) as payload:
400 if MANIFEST_KEY not in payload.files:
401 raise HarnessError(f"bundle {path} missing {MANIFEST_KEY}")
402 manifest_array = payload[MANIFEST_KEY]
403 if manifest_array.shape != ():
404 raise HarnessError(f"bundle {path} manifest is not a 0-d array")
405 if manifest_array.dtype.kind not in {"U", "S"}:
406 raise HarnessError(f"bundle {path} manifest is not a string array")
407 manifest = json.loads(str(manifest_array.item()))
408 _validate_manifest(manifest, payload, path)
409 arrays = {
410 key: np.array(payload[key], copy=True)
411 for key in payload.files
412 if key != MANIFEST_KEY
413 }
414 except HarnessError:
415 raise
416 except Exception as exc:
417 raise HarnessError(f"failed to load speaker bundle {path}: {exc}") from exc
418 return {"manifest": manifest, "arrays": arrays}
419
420
421def _validate_manifest(manifest: dict[str, Any], payload: Any, path: Path) -> None:
422 if manifest.get("schema") != BUNDLE_SCHEMA:
423 raise HarnessError(
424 f"bundle {path} has unsupported schema {manifest.get('schema')!r}"
425 )
426 if manifest.get("schema_version") != SCHEMA_VERSION:
427 raise HarnessError(
428 f"bundle {path} has unsupported schema_version "
429 f"{manifest.get('schema_version')!r}"
430 )
431 fields = manifest.get("fields")
432 if not isinstance(fields, dict):
433 raise HarnessError(f"bundle {path} manifest fields must be an object")
434
435 expected_payload_keys = {MANIFEST_KEY}
436 for key, field in fields.items():
437 state = field.get("state")
438 if state not in FIELD_STATES:
439 raise HarnessError(f"bundle {path} field {key} has invalid state {state!r}")
440 if state != PRESENT:
441 if key in payload.files:
442 raise HarnessError(
443 f"bundle {path} stores absent field {key} as a payload array"
444 )
445 continue
446 if field.get("kind") == "array":
447 if key not in payload.files:
448 raise HarnessError(f"bundle {path} missing payload array {key}")
449 array = payload[key]
450 declared_dtype = str(field.get("dtype"))
451 declared_shape = list(field.get("shape", []))
452 if str(array.dtype) != declared_dtype:
453 raise HarnessError(
454 f"bundle {path} field {key} dtype mismatch: "
455 f"{declared_dtype} != {array.dtype}"
456 )
457 if list(array.shape) != declared_shape:
458 raise HarnessError(
459 f"bundle {path} field {key} shape mismatch: "
460 f"{declared_shape} != {list(array.shape)}"
461 )
462 expected_payload_keys.add(key)
463 elif field.get("kind") == "scalar":
464 if "value" not in field:
465 raise HarnessError(f"bundle {path} scalar field {key} lacks value")
466 else:
467 raise HarnessError(f"bundle {path} present field {key} has invalid kind")
468
469 extras = set(payload.files) - expected_payload_keys
470 if extras:
471 raise HarnessError(
472 f"bundle {path} has undeclared payload arrays: {sorted(extras)}"
473 )
474
475
476def _audio_identity(audio: np.ndarray, sample_rate: int) -> dict[str, Any]:
477 samples = np.ascontiguousarray(np.asarray(audio, dtype="<f4"))
478 digest = hashlib.sha256(samples.view(np.uint8).tobytes()).hexdigest()
479 return {
480 "sample_count": int(samples.shape[0]),
481 "sample_rate": int(sample_rate),
482 "dtype": "float32-le",
483 "content_hash": f"sha256:{digest}",
484 }
485
486
487def _statement_ids_and_spans(
488 statements: Sequence[dict[str, Any]],
489) -> tuple[np.ndarray, np.ndarray]:
490 ids: list[int] = []
491 spans: list[tuple[float, float]] = []
492 for index, statement in enumerate(statements):
493 statement_id = statement.get("id", index)
494 ids.append(int(statement_id))
495 start = statement.get("start")
496 end = statement.get("end")
497 if start is None or end is None:
498 spans.append((np.nan, np.nan))
499 else:
500 spans.append((float(start), float(end)))
501 return np.asarray(ids, dtype=np.int32), np.asarray(spans, dtype=np.float64)
502
503
504def _record_input_identity(
505 bundle: Bundle,
506 plane: str,
507 audio: np.ndarray,
508 sample_rate: int,
509 statements: Sequence[dict[str, Any]],
510) -> None:
511 ids, spans = _statement_ids_and_spans(statements)
512 ids_key = (
513 INPUT_STATEMENT_IDS if plane == "statement_embedding" else INPUT_DIARIZATION_IDS
514 )
515 spans_key = (
516 INPUT_STATEMENT_SPANS
517 if plane == "statement_embedding"
518 else INPUT_DIARIZATION_SPANS
519 )
520 component = f"inputs.{plane}"
521 _set_array(bundle, ids_key, ids, component)
522 _set_array(bundle, spans_key, spans, component)
523 identity = _audio_identity(audio, sample_rate)
524 identity["statement_ids_field"] = ids_key
525 identity["statement_spans_field"] = spans_key
526 bundle["manifest"]["inputs"][plane] = identity
527
528
529def _window_stats_array(stats: Sequence[overlap.SpeakerWindowStats]) -> np.ndarray:
530 return np.asarray(
531 [
532 (row.speech_frames, row.active_slot_count, row.overlap_frames)
533 for row in stats
534 ],
535 dtype=np.int32,
536 ).reshape((-1, 3))
537
538
539def _intervals_array(intervals: Sequence[tuple[float, float, int]]) -> np.ndarray:
540 if not intervals:
541 return np.zeros((0, 3), dtype=np.float64)
542 return np.asarray(
543 [(float(start), float(end), int(local)) for start, end, local in intervals],
544 dtype=np.float64,
545 )
546
547
548def encode_statement_labels(labels: Sequence[int | None]) -> np.ndarray:
549 encoded: list[int] = []
550 for label in labels:
551 if label is None:
552 encoded.append(int(LABEL_NULL_SENTINEL))
553 continue
554 label_int = int(label)
555 if label_int <= 0:
556 raise HarnessError(
557 "production diarization label must be 1-indexed and positive; "
558 f"got {label_int}"
559 )
560 encoded.append(label_int)
561 return np.asarray(encoded, dtype=np.int32)
562
563
564@contextlib.contextmanager
565def record_diarize_private_helpers() -> Iterator[dict[str, Any]]:
566 """Record private diarizer stage values without changing orchestration."""
567 records: dict[str, Any] = {}
568 original_find = diarize._find_intervals
569 original_embed_all = diarize._embed_all_intervals
570 original_pick_k = diarize._pick_k_silhouette
571 original_cluster = diarize._cluster_intervals
572 original_assign = diarize._assign_sentences
573
574 def find_wrapper(avg_log_probs: np.ndarray, audio_len_samples: int) -> list[tuple]:
575 result = original_find(avg_log_probs, audio_len_samples)
576 records["intervals"] = list(result)
577 return result
578
579 def embed_all_wrapper(
580 audio: np.ndarray,
581 intervals: list[tuple[float, float, int]],
582 ) -> tuple[list[tuple[float, float, int]], np.ndarray]:
583 valid, embs = original_embed_all(audio, intervals)
584 records["valid_intervals"] = list(valid)
585 records["interval_embeddings"] = np.array(embs, copy=True)
586 return valid, embs
587
588 def pick_k_wrapper(embs_n: np.ndarray, max_k: int) -> int:
589 result = int(original_pick_k(embs_n, max_k))
590 records["silhouette_k"] = result
591 return result
592
593 def cluster_wrapper(embs: np.ndarray, n_speakers: int | None) -> np.ndarray:
594 labels = original_cluster(embs, n_speakers)
595 if n_speakers is None:
596 selected = records.get("silhouette_k")
597 else:
598 selected = n_speakers
599 records["silhouette_k"] = None
600 if selected is None:
601 records["effective_k"] = None
602 else:
603 upper = len(embs) - 1 if len(embs) > 1 else 1
604 records["effective_k"] = int(max(1, min(int(selected), upper)))
605 records["cluster_labels"] = np.asarray(labels, dtype=np.int32).copy()
606 return labels
607
608 def assign_wrapper(
609 sentences: list[dict],
610 intervals: list[tuple[float, float, int]],
611 global_labels: np.ndarray,
612 ) -> list[int | None]:
613 result = original_assign(sentences, intervals, global_labels)
614 records["statement_labels"] = list(result)
615 return result
616
617 diarize._find_intervals = find_wrapper
618 diarize._embed_all_intervals = embed_all_wrapper
619 diarize._pick_k_silhouette = pick_k_wrapper
620 diarize._cluster_intervals = cluster_wrapper
621 diarize._assign_sentences = assign_wrapper
622 try:
623 yield records
624 finally:
625 diarize._find_intervals = original_find
626 diarize._embed_all_intervals = original_embed_all
627 diarize._pick_k_silhouette = original_pick_k
628 diarize._cluster_intervals = original_cluster
629 diarize._assign_sentences = original_assign
630
631
632def _record_statement_embedding(
633 bundle: Bundle,
634 embedding_result: dict[str, np.ndarray] | None,
635) -> None:
636 if embedding_result is None:
637 for key in (
638 STATEMENT_EMBEDDINGS,
639 STATEMENT_EMBEDDING_IDS,
640 STATEMENT_DURATIONS,
641 STATEMENT_ENCODER,
642 ):
643 _set_state(bundle, key, ABSENT_NONE, "statement_embeddings")
644 return
645
646 _set_array(
647 bundle,
648 STATEMENT_EMBEDDINGS,
649 np.asarray(embedding_result["embeddings"], dtype=np.float32),
650 "statement_embeddings",
651 )
652 _set_array(
653 bundle,
654 STATEMENT_EMBEDDING_IDS,
655 np.asarray(embedding_result["statement_ids"], dtype=np.int32),
656 "statement_embeddings",
657 )
658 _set_array(
659 bundle,
660 STATEMENT_DURATIONS,
661 np.asarray(embedding_result.get("durations_s", []), dtype=np.float32),
662 "statement_embeddings",
663 )
664 _set_array(
665 bundle,
666 STATEMENT_ENCODER,
667 np.asarray(embedding_result["encoder"]),
668 "statement_embeddings",
669 )
670
671
672def _mark_fields(
673 bundle: Bundle,
674 fields: Sequence[str],
675 state: str,
676 component: str,
677) -> None:
678 for key in fields:
679 _set_state(bundle, key, state, component)
680
681
682def _record_pyannote_and_evidence(
683 bundle: Bundle,
684 result: overlap.OverlapInferenceResult,
685 evidence: overlap.SpeakerEvidenceDecision,
686) -> None:
687 _set_array(
688 bundle,
689 PYANNOTE_LOGPROBS,
690 np.asarray(result.avg_log_probs, dtype=np.float32),
691 "pyannote_log_probs",
692 )
693 _set_array(
694 bundle,
695 PYANNOTE_WINDOW_STATS,
696 _window_stats_array(result.window_stats),
697 "evidence",
698 )
699 _set_scalar(
700 bundle, EVIDENCE_OVERLAP_FRACTION, float(result.overlap_fraction), "evidence"
701 )
702 _set_scalar(bundle, EVIDENCE_SPEAKER, evidence.speaker_evidence, "evidence")
703 _set_scalar(
704 bundle,
705 EVIDENCE_MULTI_FRACTION,
706 float(evidence.multi_window_fraction),
707 "evidence",
708 )
709 _set_scalar(
710 bundle,
711 EVIDENCE_MEAN_OVERLAP,
712 float(evidence.mean_window_overlap_share),
713 "evidence",
714 )
715
716
717def _record_diarization(
718 bundle: Bundle,
719 records: dict[str, Any],
720 labels: Sequence[int | None],
721) -> None:
722 intervals = records.get("intervals")
723 if intervals is None:
724 _mark_fields(
725 bundle,
726 tuple(GATE_DECLINED_FIELDS),
727 NOT_EVALUATED,
728 "diarization",
729 )
730 return
731
732 _set_array(
733 bundle,
734 DIARIZATION_INTERVALS,
735 _intervals_array(intervals),
736 "intervals",
737 )
738 if not intervals:
739 _mark_fields(
740 bundle,
741 (
742 DIARIZATION_VALID_INTERVALS,
743 DIARIZATION_INTERVAL_EMBEDDINGS,
744 DIARIZATION_CLUSTER_LABELS,
745 DIARIZATION_STATEMENT_LABELS,
746 DIARIZATION_SILHOUETTE_K,
747 DIARIZATION_EFFECTIVE_K,
748 ),
749 NOT_EVALUATED,
750 "diarization",
751 )
752 return
753
754 valid_intervals = records.get("valid_intervals", [])
755 interval_embeddings = np.asarray(
756 records.get("interval_embeddings", np.zeros((0, 256), dtype=np.float32)),
757 dtype=np.float32,
758 )
759 _set_array(
760 bundle,
761 DIARIZATION_VALID_INTERVALS,
762 _intervals_array(valid_intervals),
763 "interval_embeddings",
764 )
765 _set_array(
766 bundle,
767 DIARIZATION_INTERVAL_EMBEDDINGS,
768 interval_embeddings,
769 "interval_embeddings",
770 )
771 if len(interval_embeddings) == 0:
772 _mark_fields(
773 bundle,
774 (
775 DIARIZATION_CLUSTER_LABELS,
776 DIARIZATION_STATEMENT_LABELS,
777 DIARIZATION_SILHOUETTE_K,
778 DIARIZATION_EFFECTIVE_K,
779 ),
780 NOT_EVALUATED,
781 "diarization",
782 )
783 return
784
785 _set_array(
786 bundle,
787 DIARIZATION_CLUSTER_LABELS,
788 np.asarray(records["cluster_labels"], dtype=np.int32),
789 "clustering",
790 )
791 _set_scalar(
792 bundle,
793 DIARIZATION_SILHOUETTE_K,
794 records.get("silhouette_k"),
795 "clustering",
796 )
797 _set_scalar(
798 bundle,
799 DIARIZATION_EFFECTIVE_K,
800 records.get("effective_k"),
801 "clustering",
802 )
803 _set_array(
804 bundle,
805 DIARIZATION_STATEMENT_LABELS,
806 encode_statement_labels(labels),
807 "statement_labels",
808 )
809
810
811def emit_speaker_bundle(
812 *,
813 audio_buffer: np.ndarray,
814 statements: Sequence[dict[str, Any]],
815 sample_rate: int = SAMPLE_RATE,
816 reduced_audio: np.ndarray | None = None,
817 reduction: AudioReduction | None = None,
818 producer: str = "production-python",
819) -> Bundle:
820 """Run production speaker stages and return a versioned differential bundle."""
821 if sample_rate != SAMPLE_RATE:
822 raise HarnessError(f"speaker pipeline requires {SAMPLE_RATE} Hz audio")
823
824 bundle = new_bundle(producer=producer)
825 working_statements = [dict(statement) for statement in statements]
826 stt_buffer = (
827 np.asarray(reduced_audio, dtype=np.float32)
828 if reduced_audio is not None
829 else np.asarray(audio_buffer, dtype=np.float32)
830 )
831 full_buffer = np.asarray(audio_buffer, dtype=np.float32)
832
833 _record_input_identity(
834 bundle,
835 "statement_embedding",
836 stt_buffer,
837 sample_rate,
838 working_statements,
839 )
840
841 if not working_statements:
842 _mark_fields(
843 bundle,
844 (
845 STATEMENT_EMBEDDINGS,
846 STATEMENT_EMBEDDING_IDS,
847 STATEMENT_DURATIONS,
848 STATEMENT_ENCODER,
849 ),
850 NOT_EVALUATED,
851 "statement_embeddings",
852 )
853 for component, fields in COMPONENT_FIELDS.items():
854 if component != "statement_embeddings":
855 _mark_fields(bundle, fields, NOT_EVALUATED, component)
856 _record_input_identity(
857 bundle,
858 "diarization",
859 full_buffer,
860 sample_rate,
861 working_statements,
862 )
863 _refresh_provenance(bundle, producer)
864 return bundle
865
866 embedding_result = transcribe_main._embed_statements(
867 stt_buffer,
868 working_statements,
869 sample_rate,
870 )
871 _record_statement_embedding(bundle, embedding_result)
872
873 overlap_result = overlap.compute_overlap_and_logprobs(full_buffer, sample_rate)
874 evidence = overlap.decide_speaker_evidence(
875 overlap_result.overlap_fraction,
876 overlap_result.window_stats,
877 )
878 _record_pyannote_and_evidence(bundle, overlap_result, evidence)
879
880 restored_statements = (
881 restore_statement_timestamps(working_statements, reduction)
882 if reduction is not None
883 else [dict(statement) for statement in working_statements]
884 )
885 _record_input_identity(
886 bundle,
887 "diarization",
888 full_buffer,
889 sample_rate,
890 restored_statements,
891 )
892
893 if evidence.speaker_evidence != "multi":
894 _mark_fields(bundle, tuple(GATE_DECLINED_FIELDS), NOT_EVALUATED, "diarization")
895 _refresh_provenance(bundle, producer)
896 return bundle
897
898 bundle["manifest"]["counters"]["overlap_session_run_count_before_diarize"] = (
899 _session_run_count(getattr(overlap, "_overlap_session", None))
900 )
901 try:
902 with record_diarize_private_helpers() as records:
903 # ``wav_path`` is inert when ``audio=`` is supplied; production reads
904 # it only in the audio-is-None branch.
905 labels = diarize.diarize_auto_k(
906 PLACEHOLDER_WAV_PATH,
907 restored_statements,
908 avg_log_probs=overlap_result.avg_log_probs,
909 audio=full_buffer,
910 )
911 except Exception as exc:
912 logger.exception("production diarization failed during speaker differential")
913 bundle["manifest"]["stage_errors"].append(
914 {
915 "stage": "diarization",
916 "class": type(exc).__name__,
917 "message": str(exc),
918 }
919 )
920 _mark_fields(bundle, tuple(GATE_DECLINED_FIELDS), NOT_EVALUATED, "diarization")
921 _refresh_provenance(bundle, producer)
922 return bundle
923 finally:
924 bundle["manifest"]["counters"]["overlap_session_run_count_after_diarize"] = (
925 _session_run_count(getattr(overlap, "_overlap_session", None))
926 )
927
928 _record_diarization(bundle, records, labels)
929 _refresh_provenance(bundle, producer)
930 return bundle
931
932
933def _session_run_count(session: object | None) -> int | None:
934 if session is None:
935 return None
936 count = getattr(session, "run_count", None)
937 if count is None:
938 return None
939 return int(count)
940
941
942def _component_state_verdict(
943 left: Bundle,
944 right: Bundle,
945 fields: Sequence[str],
946) -> str:
947 verdict = COMPARE
948 for key in fields:
949 pair = (_field_state(left, key), _field_state(right, key))
950 field_verdict = STATE_PAIR_VERDICT[pair]
951 if field_verdict == UNEXPECTED_DIFFERS:
952 return UNEXPECTED_DIFFERS
953 if field_verdict == NOT_EVALUATED:
954 verdict = NOT_EVALUATED
955 return verdict
956
957
958def _threshold_block(mapping: dict[str, str]) -> dict[str, Any]:
959 thresholds: dict[str, Any] = {}
960 for metric, constant_name in mapping.items():
961 if constant_name == "exact":
962 thresholds[metric] = "exact"
963 else:
964 thresholds[metric] = {
965 "constant": constant_name,
966 "value": COMPARATOR_THRESHOLDS[constant_name],
967 }
968 return thresholds
969
970
971def _component_report(
972 classification: str,
973 *,
974 thresholds: dict[str, Any] | str,
975 **extra: Any,
976) -> dict[str, Any]:
977 report = {"classification": classification, "thresholds": thresholds}
978 report.update(extra)
979 return report
980
981
982def _float_metrics(left_array: np.ndarray, right_array: np.ndarray) -> dict[str, Any]:
983 if left_array.shape != right_array.shape:
984 return {"shape_mismatch": [list(left_array.shape), list(right_array.shape)]}
985 diff = np.abs(left_array.astype(np.float64) - right_array.astype(np.float64))
986 if diff.size == 0:
987 return {"shape": list(left_array.shape), "empty": True}
988 return {
989 "shape": list(left_array.shape),
990 "max_abs_diff": float(diff.max()),
991 "median_abs_diff": float(np.median(diff)),
992 }
993
994
995def _compare_pyannote(left: Bundle, right: Bundle) -> dict[str, Any]:
996 thresholds = _threshold_block(PYANNOTE_THRESHOLD_IDS)
997 state_verdict = _component_state_verdict(
998 left, right, COMPONENT_FIELDS["pyannote_log_probs"]
999 )
1000 if state_verdict != COMPARE:
1001 return _component_report(state_verdict, thresholds=thresholds)
1002 left_array = _array(left, PYANNOTE_LOGPROBS)
1003 right_array = _array(right, PYANNOTE_LOGPROBS)
1004 metrics = _float_metrics(left_array, right_array)
1005 if "shape_mismatch" in metrics:
1006 return _component_report(
1007 UNEXPECTED_DIFFERS, thresholds=thresholds, metrics=metrics
1008 )
1009 if left_array.size == 0 and right_array.size == 0:
1010 return _component_report(NOT_EVALUATED, thresholds=thresholds, metrics=metrics)
1011 argmax_agreement = float(
1012 (left_array.argmax(axis=-1) == right_array.argmax(axis=-1)).mean()
1013 )
1014 metrics["per_frame_argmax_agreement_fraction"] = argmax_agreement
1015 if np.array_equal(left_array, right_array):
1016 classification = EQUAL
1017 elif (
1018 metrics["max_abs_diff"] <= LOGPROB_MAX_ABS_TOLERANCE
1019 and metrics["median_abs_diff"] <= LOGPROB_MEDIAN_ABS_TOLERANCE
1020 and argmax_agreement >= LOGPROB_ARGMAX_AGREEMENT_MIN
1021 ):
1022 classification = FUNCTIONALLY_EQUAL
1023 else:
1024 classification = UNEXPECTED_DIFFERS
1025 return _component_report(classification, thresholds=thresholds, metrics=metrics)
1026
1027
1028def _compare_evidence(left: Bundle, right: Bundle) -> dict[str, Any]:
1029 thresholds = _threshold_block(EVIDENCE_THRESHOLD_IDS)
1030 state_verdict = _component_state_verdict(left, right, COMPONENT_FIELDS["evidence"])
1031 if state_verdict != COMPARE:
1032 return _component_report(state_verdict, thresholds=thresholds)
1033 diffs: list[dict[str, Any]] = []
1034 left_speaker = _scalar(left, EVIDENCE_SPEAKER)
1035 right_speaker = _scalar(right, EVIDENCE_SPEAKER)
1036 if left_speaker != right_speaker:
1037 diffs.append(
1038 {
1039 "field": EVIDENCE_SPEAKER,
1040 "left": left_speaker,
1041 "right": right_speaker,
1042 }
1043 )
1044 metrics: dict[str, Any] = {}
1045 for key in (
1046 EVIDENCE_OVERLAP_FRACTION,
1047 EVIDENCE_MULTI_FRACTION,
1048 EVIDENCE_MEAN_OVERLAP,
1049 ):
1050 delta = abs(float(_scalar(left, key)) - float(_scalar(right, key)))
1051 metrics[key] = {"abs_diff": delta}
1052 if delta > EVIDENCE_FLOAT_ABS_TOLERANCE:
1053 diffs.append({"field": key, "abs_diff": delta})
1054 left_stats = _array(left, PYANNOTE_WINDOW_STATS)
1055 right_stats = _array(right, PYANNOTE_WINDOW_STATS)
1056 if not np.array_equal(left_stats, right_stats):
1057 diffs.append(
1058 {
1059 "field": PYANNOTE_WINDOW_STATS,
1060 "left": left_stats.tolist(),
1061 "right": right_stats.tolist(),
1062 }
1063 )
1064 metrics["window_stats_equal"] = bool(np.array_equal(left_stats, right_stats))
1065 if diffs:
1066 return _component_report(
1067 UNEXPECTED_DIFFERS,
1068 thresholds=thresholds,
1069 metrics=metrics,
1070 differences=diffs,
1071 )
1072 exact = all(
1073 metrics[key]["abs_diff"] == 0.0
1074 for key in metrics
1075 if key != "window_stats_equal"
1076 )
1077 return _component_report(
1078 EQUAL if exact else FUNCTIONALLY_EQUAL,
1079 thresholds=thresholds,
1080 metrics=metrics,
1081 )
1082
1083
1084def _interval_set(array: np.ndarray) -> set[tuple[float, float, int]]:
1085 return {
1086 (float(row[0]), float(row[1]), int(row[2]))
1087 for row in np.asarray(array).reshape((-1, 3))
1088 }
1089
1090
1091def _assert_unique_statement_ids(ids: np.ndarray, field: str) -> None:
1092 values = [int(value) for value in ids.tolist()]
1093 if len(values) != len(set(values)):
1094 raise HarnessError(f"duplicate statement ids in {field}: {values}")
1095
1096
1097def _assert_parallel_length(ids: np.ndarray, values: np.ndarray, field: str) -> None:
1098 if len(values) != len(ids):
1099 raise HarnessError(
1100 f"{field} length {len(values)} does not match statement id count {len(ids)}"
1101 )
1102
1103
1104def _scalar_string_array(bundle: Bundle, key: str) -> str:
1105 value = _array(bundle, key)
1106 if value.shape != ():
1107 raise HarnessError(f"{key} must be a 0-d string array")
1108 return str(value.item())
1109
1110
1111def _compare_intervals(left: Bundle, right: Bundle) -> dict[str, Any]:
1112 state_verdict = _component_state_verdict(left, right, COMPONENT_FIELDS["intervals"])
1113 if state_verdict != COMPARE:
1114 return _component_report(state_verdict, thresholds="exact")
1115 left_intervals = _array(left, DIARIZATION_INTERVALS)
1116 right_intervals = _array(right, DIARIZATION_INTERVALS)
1117 if left_intervals.shape != right_intervals.shape:
1118 return _component_report(
1119 UNEXPECTED_DIFFERS,
1120 thresholds="exact",
1121 metrics={
1122 "shape_mismatch": [
1123 list(left_intervals.shape),
1124 list(right_intervals.shape),
1125 ]
1126 },
1127 )
1128 if len(left_intervals) == 0 and len(right_intervals) == 0:
1129 return _component_report(
1130 NOT_EVALUATED, thresholds="exact", metrics={"interval_count": 0}
1131 )
1132 left_set = _interval_set(left_intervals)
1133 right_set = _interval_set(right_intervals)
1134 if left_set == right_set:
1135 return _component_report(
1136 EQUAL, thresholds="exact", metrics={"interval_count": len(left_set)}
1137 )
1138 return _component_report(
1139 UNEXPECTED_DIFFERS,
1140 thresholds="exact",
1141 metrics={"left_count": len(left_set), "right_count": len(right_set)},
1142 differences={
1143 "added": sorted(right_set - left_set),
1144 "removed": sorted(left_set - right_set),
1145 },
1146 )
1147
1148
1149def _cosine_metrics(left_array: np.ndarray, right_array: np.ndarray) -> dict[str, Any]:
1150 diff = np.abs(left_array.astype(np.float64) - right_array.astype(np.float64))
1151 left_norm = np.linalg.norm(left_array, axis=1)
1152 right_norm = np.linalg.norm(right_array, axis=1)
1153 denom = np.maximum(left_norm * right_norm, COSINE_NORM_FLOOR)
1154 cosine = np.sum(left_array * right_array, axis=1) / denom
1155 return {
1156 "max_abs_component_diff": float(diff.max()) if diff.size else 0.0,
1157 "min_cosine_similarity": float(cosine.min()) if cosine.size else 1.0,
1158 "median_cosine_similarity": float(np.median(cosine)) if cosine.size else 1.0,
1159 "max_cosine_similarity": float(cosine.max()) if cosine.size else 1.0,
1160 }
1161
1162
1163def _embedding_classification(
1164 left_array: np.ndarray,
1165 right_array: np.ndarray,
1166 metrics: dict[str, Any],
1167) -> str:
1168 if np.array_equal(left_array, right_array):
1169 return EQUAL
1170 if (
1171 metrics["max_abs_component_diff"] <= EMBEDDING_MAX_ABS_TOLERANCE
1172 and metrics["min_cosine_similarity"] >= EMBEDDING_MIN_COSINE_SIMILARITY
1173 and metrics["median_cosine_similarity"] >= EMBEDDING_MEDIAN_COSINE_SIMILARITY
1174 ):
1175 return FUNCTIONALLY_EQUAL
1176 return UNEXPECTED_DIFFERS
1177
1178
1179def _compare_statement_embeddings(left: Bundle, right: Bundle) -> dict[str, Any]:
1180 thresholds = _threshold_block(STATEMENT_EMBEDDING_THRESHOLD_IDS)
1181 state_verdict = _component_state_verdict(
1182 left, right, COMPONENT_FIELDS["statement_embeddings"]
1183 )
1184 if state_verdict != COMPARE:
1185 return _component_report(state_verdict, thresholds=thresholds)
1186 left_ids = _array(left, STATEMENT_EMBEDDING_IDS)
1187 right_ids = _array(right, STATEMENT_EMBEDDING_IDS)
1188 _assert_unique_statement_ids(left_ids, STATEMENT_EMBEDDING_IDS)
1189 _assert_unique_statement_ids(right_ids, STATEMENT_EMBEDDING_IDS)
1190 if set(left_ids.tolist()) != set(right_ids.tolist()):
1191 return _component_report(
1192 UNEXPECTED_DIFFERS,
1193 thresholds=thresholds,
1194 differences={
1195 "left_ids": left_ids.tolist(),
1196 "right_ids": right_ids.tolist(),
1197 },
1198 )
1199 left_embs = _array(left, STATEMENT_EMBEDDINGS)
1200 right_embs = _array(right, STATEMENT_EMBEDDINGS)
1201 left_durations = _array(left, STATEMENT_DURATIONS)
1202 right_durations = _array(right, STATEMENT_DURATIONS)
1203 _assert_parallel_length(left_ids, left_embs, STATEMENT_EMBEDDINGS)
1204 _assert_parallel_length(right_ids, right_embs, STATEMENT_EMBEDDINGS)
1205 _assert_parallel_length(left_ids, left_durations, STATEMENT_DURATIONS)
1206 _assert_parallel_length(right_ids, right_durations, STATEMENT_DURATIONS)
1207 if len(left_ids) == 0 and len(right_ids) == 0:
1208 return _component_report(
1209 NOT_EVALUATED, thresholds=thresholds, metrics={"statement_count": 0}
1210 )
1211 if left_embs.shape[1:] != right_embs.shape[1:]:
1212 return _component_report(
1213 UNEXPECTED_DIFFERS,
1214 thresholds=thresholds,
1215 metrics={"shape_mismatch": [list(left_embs.shape), list(right_embs.shape)]},
1216 )
1217 left_index = {int(statement_id): idx for idx, statement_id in enumerate(left_ids)}
1218 right_index = {int(statement_id): idx for idx, statement_id in enumerate(right_ids)}
1219 ordered_ids = sorted(left_index)
1220 left_ordered = np.stack(
1221 [left_embs[left_index[statement_id]] for statement_id in ordered_ids]
1222 )
1223 right_ordered = np.stack(
1224 [right_embs[right_index[statement_id]] for statement_id in ordered_ids]
1225 )
1226 duration_diffs = np.asarray(
1227 [
1228 abs(
1229 float(left_durations[left_index[statement_id]])
1230 - float(right_durations[right_index[statement_id]])
1231 )
1232 for statement_id in ordered_ids
1233 ],
1234 dtype=np.float64,
1235 )
1236 left_encoder = _scalar_string_array(left, STATEMENT_ENCODER)
1237 right_encoder = _scalar_string_array(right, STATEMENT_ENCODER)
1238 metrics = _cosine_metrics(left_ordered, right_ordered)
1239 metrics["statement_count"] = len(ordered_ids)
1240 metrics["duration_max_abs_diff"] = (
1241 float(duration_diffs.max()) if duration_diffs.size else 0.0
1242 )
1243 vector_classification = _embedding_classification(
1244 left_ordered, right_ordered, metrics
1245 )
1246 differences: list[dict[str, Any]] = []
1247 if left_encoder != right_encoder:
1248 differences.append(
1249 {
1250 "field": STATEMENT_ENCODER,
1251 "left": left_encoder,
1252 "right": right_encoder,
1253 }
1254 )
1255 if metrics["duration_max_abs_diff"] > STATEMENT_DURATION_ABS_TOLERANCE:
1256 differences.append(
1257 {
1258 "field": STATEMENT_DURATIONS,
1259 "max_abs_diff": metrics["duration_max_abs_diff"],
1260 }
1261 )
1262 if vector_classification == UNEXPECTED_DIFFERS or differences:
1263 return _component_report(
1264 UNEXPECTED_DIFFERS,
1265 thresholds=thresholds,
1266 metrics=metrics,
1267 differences=differences,
1268 )
1269 classification = (
1270 FUNCTIONALLY_EQUAL
1271 if vector_classification == FUNCTIONALLY_EQUAL
1272 or metrics["duration_max_abs_diff"] > 0.0
1273 else EQUAL
1274 )
1275 return _component_report(classification, thresholds=thresholds, metrics=metrics)
1276
1277
1278def _compare_interval_embeddings(left: Bundle, right: Bundle) -> dict[str, Any]:
1279 thresholds = _threshold_block(INTERVAL_EMBEDDING_THRESHOLD_IDS)
1280 state_verdict = _component_state_verdict(
1281 left, right, COMPONENT_FIELDS["interval_embeddings"]
1282 )
1283 if state_verdict != COMPARE:
1284 return _component_report(state_verdict, thresholds=thresholds)
1285 left_valid = _array(left, DIARIZATION_VALID_INTERVALS)
1286 right_valid = _array(right, DIARIZATION_VALID_INTERVALS)
1287 left_embs = _array(left, DIARIZATION_INTERVAL_EMBEDDINGS)
1288 right_embs = _array(right, DIARIZATION_INTERVAL_EMBEDDINGS)
1289 if left_valid.shape != right_valid.shape:
1290 return _component_report(
1291 UNEXPECTED_DIFFERS,
1292 thresholds=thresholds,
1293 metrics={
1294 "valid_interval_shape_mismatch": [
1295 list(left_valid.shape),
1296 list(right_valid.shape),
1297 ]
1298 },
1299 )
1300 left_valid_set = _interval_set(left_valid)
1301 right_valid_set = _interval_set(right_valid)
1302 if left_valid_set != right_valid_set:
1303 return _component_report(
1304 UNEXPECTED_DIFFERS,
1305 thresholds=thresholds,
1306 metrics={
1307 "left_valid_interval_count": len(left_valid_set),
1308 "right_valid_interval_count": len(right_valid_set),
1309 },
1310 differences={
1311 "added_valid_intervals": sorted(right_valid_set - left_valid_set),
1312 "removed_valid_intervals": sorted(left_valid_set - right_valid_set),
1313 },
1314 )
1315 if left_embs.shape != right_embs.shape:
1316 return _component_report(
1317 UNEXPECTED_DIFFERS,
1318 thresholds=thresholds,
1319 metrics={"shape_mismatch": [list(left_embs.shape), list(right_embs.shape)]},
1320 )
1321 if left_embs.size == 0 and right_embs.size == 0:
1322 return _component_report(
1323 NOT_EVALUATED, thresholds=thresholds, metrics={"interval_count": 0}
1324 )
1325 metrics = _cosine_metrics(left_embs, right_embs)
1326 metrics["interval_count"] = int(left_embs.shape[0])
1327 classification = _embedding_classification(left_embs, right_embs, metrics)
1328 return _component_report(classification, thresholds=thresholds, metrics=metrics)
1329
1330
1331def _partition(labels: np.ndarray) -> set[frozenset[int]]:
1332 groups: dict[int, set[int]] = {}
1333 for index, label in enumerate(labels.tolist()):
1334 groups.setdefault(int(label), set()).add(index)
1335 return {frozenset(members) for members in groups.values()}
1336
1337
1338def _compare_clustering(left: Bundle, right: Bundle) -> dict[str, Any]:
1339 state_verdict = _component_state_verdict(
1340 left, right, COMPONENT_FIELDS["clustering"]
1341 )
1342 if state_verdict != COMPARE:
1343 return _component_report(state_verdict, thresholds="exact")
1344 left_labels = _array(left, DIARIZATION_CLUSTER_LABELS)
1345 right_labels = _array(right, DIARIZATION_CLUSTER_LABELS)
1346 if left_labels.shape != right_labels.shape:
1347 return _component_report(
1348 UNEXPECTED_DIFFERS,
1349 thresholds="exact",
1350 metrics={
1351 "shape_mismatch": [list(left_labels.shape), list(right_labels.shape)]
1352 },
1353 )
1354 if left_labels.size == 0 and right_labels.size == 0:
1355 return _component_report(
1356 NOT_EVALUATED, thresholds="exact", metrics={"cluster_label_count": 0}
1357 )
1358 scalar_diffs = []
1359 for key in (DIARIZATION_SILHOUETTE_K, DIARIZATION_EFFECTIVE_K):
1360 if _scalar(left, key) != _scalar(right, key):
1361 scalar_diffs.append(
1362 {"field": key, "left": _scalar(left, key), "right": _scalar(right, key)}
1363 )
1364 left_partition = _partition(left_labels)
1365 right_partition = _partition(right_labels)
1366 if scalar_diffs or left_partition != right_partition:
1367 return _component_report(
1368 UNEXPECTED_DIFFERS,
1369 thresholds="exact",
1370 metrics={
1371 "cluster_label_count": int(left_labels.size),
1372 "partition_equal_up_to_permutation": left_partition == right_partition,
1373 },
1374 differences={"scalars": scalar_diffs},
1375 )
1376 return _component_report(
1377 EQUAL,
1378 thresholds="exact",
1379 metrics={
1380 "cluster_label_count": int(left_labels.size),
1381 "partition_equal_up_to_permutation": True,
1382 },
1383 )
1384
1385
1386def _compare_statement_labels(left: Bundle, right: Bundle) -> dict[str, Any]:
1387 state_verdict = _component_state_verdict(
1388 left, right, COMPONENT_FIELDS["statement_labels"]
1389 )
1390 if state_verdict != COMPARE:
1391 return _component_report(state_verdict, thresholds="exact")
1392 left_ids = _array(left, INPUT_DIARIZATION_IDS)
1393 right_ids = _array(right, INPUT_DIARIZATION_IDS)
1394 left_labels = _array(left, DIARIZATION_STATEMENT_LABELS)
1395 right_labels = _array(right, DIARIZATION_STATEMENT_LABELS)
1396 _assert_unique_statement_ids(left_ids, INPUT_DIARIZATION_IDS)
1397 _assert_unique_statement_ids(right_ids, INPUT_DIARIZATION_IDS)
1398 _assert_parallel_length(left_ids, left_labels, DIARIZATION_STATEMENT_LABELS)
1399 _assert_parallel_length(right_ids, right_labels, DIARIZATION_STATEMENT_LABELS)
1400 if set(left_ids.tolist()) != set(right_ids.tolist()):
1401 return _component_report(
1402 UNEXPECTED_DIFFERS,
1403 thresholds="exact",
1404 differences={
1405 "left_ids": left_ids.tolist(),
1406 "right_ids": right_ids.tolist(),
1407 },
1408 )
1409 if left_labels.size == 0 and right_labels.size == 0:
1410 return _component_report(
1411 NOT_EVALUATED, thresholds="exact", metrics={"statement_count": 0}
1412 )
1413 if np.all(left_labels == LABEL_NULL_SENTINEL) and np.all(
1414 right_labels == LABEL_NULL_SENTINEL
1415 ):
1416 return _component_report(
1417 NOT_EVALUATED,
1418 thresholds="exact",
1419 metrics={"statement_count": int(left_labels.size), "all_null": True},
1420 )
1421 left_index = {int(statement_id): idx for idx, statement_id in enumerate(left_ids)}
1422 right_index = {int(statement_id): idx for idx, statement_id in enumerate(right_ids)}
1423 mismatches = []
1424 for statement_id in sorted(left_index):
1425 left_label = int(left_labels[left_index[statement_id]])
1426 right_label = int(right_labels[right_index[statement_id]])
1427 if left_label != right_label:
1428 mismatches.append(
1429 {
1430 "statement_id": statement_id,
1431 "left": None if left_label == LABEL_NULL_SENTINEL else left_label,
1432 "right": None
1433 if right_label == LABEL_NULL_SENTINEL
1434 else right_label,
1435 }
1436 )
1437 if mismatches:
1438 return _component_report(
1439 UNEXPECTED_DIFFERS,
1440 thresholds="exact",
1441 metrics={
1442 "statement_count": int(left_labels.size),
1443 "agreement_fraction": 1.0 - len(mismatches) / max(1, len(left_index)),
1444 },
1445 differences=mismatches,
1446 )
1447 return _component_report(
1448 EQUAL,
1449 thresholds="exact",
1450 metrics={"statement_count": int(left_labels.size), "agreement_fraction": 1.0},
1451 )
1452
1453
1454def _validate_input_alignment(left: Bundle, right: Bundle) -> None:
1455 for plane in ("statement_embedding", "diarization"):
1456 if left["manifest"]["inputs"].get(plane) != right["manifest"]["inputs"].get(
1457 plane
1458 ):
1459 raise HarnessError(f"input identity mismatch for {plane} plane")
1460 for key in (
1461 INPUT_STATEMENT_IDS,
1462 INPUT_STATEMENT_SPANS,
1463 INPUT_DIARIZATION_IDS,
1464 INPUT_DIARIZATION_SPANS,
1465 ):
1466 if not np.array_equal(_array(left, key), _array(right, key), equal_nan=True):
1467 raise HarnessError(f"input array mismatch for {key}")
1468
1469
1470def _both_gate_declined(left: Bundle, right: Bundle) -> bool:
1471 try:
1472 return (
1473 _field_state(left, EVIDENCE_SPEAKER) == PRESENT
1474 and _field_state(right, EVIDENCE_SPEAKER) == PRESENT
1475 and _scalar(left, EVIDENCE_SPEAKER) != "multi"
1476 and _scalar(right, EVIDENCE_SPEAKER) != "multi"
1477 )
1478 except HarnessError:
1479 return False
1480
1481
1482def _rollup(left: Bundle, right: Bundle, components: dict[str, dict[str, Any]]) -> str:
1483 values = [component["classification"] for component in components.values()]
1484 if any(value == UNEXPECTED_DIFFERS for value in values):
1485 return UNEXPECTED_DIFFERS
1486 if _both_gate_declined(left, right):
1487 return NOT_EVALUATED
1488 if any(value == FUNCTIONALLY_EQUAL for value in values):
1489 return FUNCTIONALLY_EQUAL
1490 if any(value == EQUAL for value in values):
1491 return EQUAL
1492 return NOT_EVALUATED
1493
1494
1495def _base_report(
1496 left: Bundle | None = None, right: Bundle | None = None
1497) -> dict[str, Any]:
1498 return {
1499 "schema": REPORT_SCHEMA,
1500 "schema_version": SCHEMA_VERSION,
1501 "classification": NOT_EVALUATED,
1502 "failure": None,
1503 "provenance": {
1504 "generated_at": datetime.now(UTC).isoformat().replace("+00:00", "Z"),
1505 "harness": {
1506 "name": "tests.verify_speaker_differential",
1507 "repo_commit": get_rev(),
1508 "version": None,
1509 },
1510 "host": {
1511 "platform": platform.platform(),
1512 "python": platform.python_version(),
1513 },
1514 "bundles": {
1515 "left": left["manifest"]["provenance"] if left is not None else None,
1516 "right": right["manifest"]["provenance"] if right is not None else None,
1517 },
1518 },
1519 "components": {},
1520 }
1521
1522
1523def compare_bundles(left: Bundle, right: Bundle) -> dict[str, Any]:
1524 report = _base_report(left, right)
1525 try:
1526 if left["manifest"].get("stage_errors") or right["manifest"].get(
1527 "stage_errors"
1528 ):
1529 raise HarnessError("one or both bundles contain stage_errors")
1530 _validate_input_alignment(left, right)
1531 components = {
1532 "pyannote_log_probs": _compare_pyannote(left, right),
1533 "evidence": _compare_evidence(left, right),
1534 "statement_embeddings": _compare_statement_embeddings(left, right),
1535 "intervals": _compare_intervals(left, right),
1536 "interval_embeddings": _compare_interval_embeddings(left, right),
1537 "clustering": _compare_clustering(left, right),
1538 "statement_labels": _compare_statement_labels(left, right),
1539 }
1540 report["components"] = {key: components[key] for key in COMPONENT_ORDER}
1541 report["classification"] = _rollup(left, right, components)
1542 except HarnessError as exc:
1543 report["classification"] = NOT_EVALUATED
1544 report["failure"] = {"class": HARNESS_ERROR, "message": str(exc)}
1545 return report
1546
1547
1548def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
1549 parser = argparse.ArgumentParser(description=__doc__)
1550 parser.add_argument("left_bundle", help="Left speaker differential .npz bundle")
1551 parser.add_argument("right_bundle", help="Right speaker differential .npz bundle")
1552 parser.add_argument(
1553 "--report", help="JSON report destination outside the repository"
1554 )
1555 return parser.parse_args(argv)
1556
1557
1558def main(argv: list[str] | None = None) -> int:
1559 args = parse_args(argv)
1560 requested_report_path = Path(args.report).resolve() if args.report else None
1561 report_path: Path | None = None
1562 try:
1563 if requested_report_path is not None:
1564 report_path = _refuse_repo_destination(requested_report_path)
1565 left = load_bundle(Path(args.left_bundle))
1566 right = load_bundle(Path(args.right_bundle))
1567 report = compare_bundles(left, right)
1568 except Exception as exc:
1569 report = _base_report()
1570 report["failure"] = {"class": HARNESS_ERROR, "message": str(exc)}
1571 report["classification"] = NOT_EVALUATED
1572
1573 rendered = _render_report(report)
1574 if report_path is not None:
1575 report_path.parent.mkdir(parents=True, exist_ok=True)
1576 report_path.write_text(rendered, encoding="utf-8")
1577 sys.stdout.write(rendered)
1578 return 0 if report.get("classification") in {EQUAL, FUNCTIONALLY_EQUAL} else 1
1579
1580
1581if __name__ == "__main__":
1582 raise SystemExit(main())