ddidderr
3042b7743b
Sampled Gemini voice matching could leave one person under separate IDs across transcription chunks. Use hosted Precision-2 on the complete recording and align its exclusive speaker timeline with MAI word timing. Preserve transcription text, order, and timestamps; leave ambiguous word assignments unknown instead of merging identities to satisfy a count. Upload lossless FLAC, bound polling, retry transient status reads against the same job, and validate results before emitting stdout. Keep optional speaker-count and hosted model overrides. Disable MAI diarization and use 30-minute transcription uploads in both modes. Respect environment keys, keep authentication out of process arguments, and save optional diagnostics without audio, credentials, or signed upload URLs. Test Plan: - python3 -B -m unittest -v: all 101 regression tests passed, including mocked provider upload, polling, alignment, failure, and cleanup paths. - bash -n transcribe and shellcheck transcribe: passed. - git diff --cached --check: passed. - Prepared the 74-minute input locally and verified identical decoded PCM hashes for the normalized WAV and full-recording lossless FLAC. - Live API testing was explicitly deferred by the user; mocked results do not establish speaker accuracy on the recording. BREAKING CHANGE: Default diarization requires PYANNOTE_API_KEY or its 1Password reference. --speaker-model accepts hosted precision-2 or community-1 instead of OpenRouter audio-chat model IDs.
381 lines
21 KiB
Python
381 lines
21 KiB
Python
import contextlib
|
|
import copy
|
|
import io
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
import wave
|
|
|
|
import diarization
|
|
|
|
|
|
def word(start, end, speaker=9, text="speech"):
|
|
return {"word": text, "start": start, "end": end, "speaker": speaker}
|
|
|
|
|
|
def transcript(*words):
|
|
return {"text": " ".join(item["word"] for item in words), "words": list(words)}
|
|
|
|
|
|
def turn(start, end, speaker="voice-A", **extra):
|
|
return dict(start=start, end=end, speaker=speaker, **extra)
|
|
|
|
|
|
def job(*segments, **extra):
|
|
return dict(jobId="test-job", status="succeeded",
|
|
output={"exclusiveDiarization": list(segments)}, **extra)
|
|
|
|
|
|
def align(source, result, speakers=None):
|
|
with contextlib.redirect_stderr(io.StringIO()):
|
|
return diarization.align(source, result, speakers)
|
|
|
|
|
|
def write_audio(directory, duration=30):
|
|
path = Path(directory) / "audio.wav"
|
|
with wave.open(str(path), "wb") as audio:
|
|
audio.setparams((1, 2, 16000, 0, "NONE", "not compressed"))
|
|
audio.writeframes(b"\0\0" * int(duration * 16000))
|
|
return path
|
|
|
|
|
|
class AlignmentTests(unittest.TestCase):
|
|
def test_global_timeline_reunites_host_in_ads_and_across_request_local_ids(self):
|
|
source = transcript(word(1, 3, 7, "Host interview."), word(4, 6, 8, "Guest answer."),
|
|
word(601, 603, 8, "Host ad with a different microphone."),
|
|
word(1001, 1003, 9, "Host interview again."),
|
|
word(1004, 1006, 7, "Guest answer again."))
|
|
timeline = job(turn(0, 3.5, "HOST"), turn(4, 8, "GUEST"), turn(600, 650, "HOST"),
|
|
turn(1000, 1003.5, "HOST"), turn(1004, 1020, "GUEST"))
|
|
result = align(source, timeline, speakers=2)
|
|
self.assertEqual([w["speaker"] for w in result["words"]], [0, 1, 0, 0, 1])
|
|
|
|
def test_reused_provider_local_id_never_merges_distinct_global_voices(self):
|
|
source = transcript(word(1, 2, 0), word(4, 5, 0))
|
|
result = align(source, job(turn(0, 3, "a"), turn(4, 6, "b")))
|
|
self.assertEqual([w["speaker"] for w in result["words"]], [0, 1])
|
|
|
|
def test_all_original_text_order_timestamps_and_other_fields_survive(self):
|
|
source = transcript(word(8, 12, None, " Original\t text "), word(7, 9, 3, "Earlier timing"))
|
|
source["words"][0]["confidence"] = 0.7
|
|
source["words"][1].pop("speaker")
|
|
source["metadata"] = {"nested": [1, 2]}
|
|
before = copy.deepcopy(source)
|
|
result = align(source, job(turn(0, 20)))
|
|
self.assertEqual(source, before)
|
|
self.assertEqual({k: v for k, v in result.items() if k != "words"},
|
|
{k: v for k, v in before.items() if k != "words"})
|
|
for original, changed in zip(before["words"], result["words"]):
|
|
self.assertEqual({k: v for k, v in original.items() if k != "speaker"},
|
|
{k: v for k, v in changed.items() if k != "speaker"})
|
|
self.assertEqual([w["speaker"] for w in result["words"]], [0, 0])
|
|
|
|
def test_ids_follow_first_assigned_nonempty_word_not_untranscribed_timeline_turns(self):
|
|
source = transcript(word(5, 6, text=" "), word(8, 9, text="First"), word(4, 4.5, text="Second"))
|
|
result = align(source, job(turn(0, 2, "ignored"), turn(3, 6, "later"), turn(7, 10, "first")))
|
|
self.assertEqual([w["speaker"] for w in result["words"]], [None, 0, 1])
|
|
|
|
def test_overlap_aggregates_same_speaker_across_separate_turns(self):
|
|
result = align(transcript(word(0, 10)), job(turn(0, 4, "a"), turn(4, 6, "b"), turn(6, 10, "a")))
|
|
self.assertEqual(result["words"][0]["speaker"], 0)
|
|
|
|
def test_dominant_overlap_wins_but_equal_voices_or_low_coverage_stay_unknown(self):
|
|
for segments, interval, expected in [
|
|
([turn(0, 6, "a"), turn(6, 10, "b")], (0, 10), 0),
|
|
([turn(0, 5, "a"), turn(5, 10, "b")], (0, 10), None),
|
|
([turn(0, 1, "a")], (0, 10), None),
|
|
]:
|
|
with self.subTest(segments=segments):
|
|
self.assertEqual(align(transcript(word(*interval)), job(*segments))["words"][0]["speaker"], expected)
|
|
|
|
def test_nearest_is_bounded_for_whole_word_and_not_only_closest_edge(self):
|
|
source = transcript(word(1.01, 1.2), word(1.01, 2), word(0, 0.1), word(8, 8.1))
|
|
log = io.StringIO()
|
|
with contextlib.redirect_stderr(log):
|
|
result = diarization.align(source, job(turn(0.2, 1)))
|
|
self.assertEqual([w["speaker"] for w in result["words"]], [0, None, 0, None])
|
|
self.assertIn("2 word(s) used a nearest", log.getvalue())
|
|
self.assertIn("2 word(s) have ambiguous", log.getvalue())
|
|
|
|
def test_zero_duration_and_nearest_ties_do_not_force_a_speaker(self):
|
|
cases = [
|
|
(job(turn(0, 1, "a"), turn(1, 2, "b")), 1, None),
|
|
(job(turn(0, 1, "a"), turn(1.2, 2, "b")), 1.1, None),
|
|
(job(turn(0, 1, "a"), turn(1.2, 2, "a")), 1.1, 0),
|
|
(job(turn(0, 1, "a")), 0.5, 0),
|
|
(job(turn(0, 1, "a")), 1.3, None),
|
|
]
|
|
for timeline, timestamp, expected in cases:
|
|
with self.subTest(timestamp=timestamp, timeline=timeline):
|
|
self.assertEqual(align(transcript(word(timestamp, timestamp)), timeline)["words"][0]["speaker"], expected)
|
|
|
|
def test_empty_timeline_is_valid_only_without_transcribed_speech(self):
|
|
self.assertEqual(align(transcript(), job()), transcript())
|
|
with self.assertRaisesRegex(ValueError, "empty speaker timeline"):
|
|
align(transcript(word(0, 1, 42)), job())
|
|
|
|
def test_known_count_warns_without_merging_or_fabricating_speakers(self):
|
|
log = io.StringIO()
|
|
with contextlib.redirect_stderr(log):
|
|
result = diarization.align(transcript(word(0, 1), word(2, 3), word(4, 5)),
|
|
job(turn(0, 1, "a"), turn(2, 3, "b"), turn(4, 5, "c")), 2)
|
|
self.assertEqual([w["speaker"] for w in result["words"]], [0, 1, 2])
|
|
self.assertIn("expected 2 speakers", log.getvalue())
|
|
|
|
def test_malformed_incomplete_unsorted_or_overlapping_timelines_fail(self):
|
|
invalid = [None, {}, {"status": "running"}, dict(job(), status="failed"),
|
|
dict(job(), output={}), dict(job(), error="failure"),
|
|
dict(job(), output={"error": "failure", "exclusiveDiarization": []}),
|
|
job(None), job(turn(-1, 1)), job(turn(2, 1)), job(turn(1, 1)),
|
|
job(turn(True, 2)), job(turn(0, float("nan"))), job(turn(0, float("inf"))),
|
|
job(turn(0, 1, None)), job(turn(0, 1, " ")), job(turn(0, 1, 0)),
|
|
job(turn(0, 3), turn(2, 4)), job(turn(5, 6), turn(1, 2)),
|
|
job(turn(0, 1, confidence=0.7)), job(turn(0, 1, confidence={"a": 101})),
|
|
job(turn(0, 1, confidence={"a": True})),
|
|
job(turn(0, 1, confidence={"a": float("nan")}))]
|
|
for timeline in invalid:
|
|
with self.subTest(timeline=timeline), self.assertRaises(ValueError):
|
|
align(transcript(word(0, 1)), timeline)
|
|
align(transcript(word(0, 1)), job(turn(0, 1, confidence={"voice-A": 97.2, "voice-B": 2.8})))
|
|
|
|
def test_invalid_words_or_known_speaker_counts_fail(self):
|
|
for item in (word(-1, 1), word(1, 0), word(False, 1), word(0, float("nan"))):
|
|
with self.subTest(item=item), self.assertRaises(ValueError):
|
|
align(transcript(item), job(turn(0, 2)))
|
|
for count in (0, -1, True, 1.5):
|
|
with self.subTest(count=count), self.assertRaises(ValueError):
|
|
align(transcript(word(0, 1)), job(turn(0, 2)), count)
|
|
|
|
def test_cli_output_only_after_complete_validation(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
source = Path(directory) / "transcript.json"
|
|
result = Path(directory) / "result.json"
|
|
source.write_text(json.dumps(transcript(word(0, 1, text="Must survive."))))
|
|
result.write_text(json.dumps(job(turn(0, 2))))
|
|
command = [sys.executable, "-B", str(Path(diarization.__file__)), "apply", str(source), str(result)]
|
|
completed = subprocess.run([*command, "--json"], capture_output=True, text=True)
|
|
self.assertEqual(completed.returncode, 0, completed.stderr)
|
|
self.assertEqual(json.loads(completed.stdout)["words"][0]["speaker"], 0)
|
|
completed = subprocess.run([*command, "--srt"], capture_output=True, text=True)
|
|
self.assertEqual(completed.returncode, 0, completed.stderr)
|
|
self.assertIn("00:00:00,000 --> 00:00:01,000", completed.stdout)
|
|
result.write_text('{"status":"succeeded","status":"running"}')
|
|
completed = subprocess.run(command, capture_output=True, text=True)
|
|
self.assertNotEqual(completed.returncode, 0)
|
|
self.assertEqual(completed.stdout, "")
|
|
self.assertIn("Duplicate JSON key", completed.stderr)
|
|
|
|
|
|
class FakeTransport:
|
|
"""Exercise real Client command construction without executing curl."""
|
|
def __init__(self, polls=None, failure=None):
|
|
self.polls = list(polls or [job(turn(0, 10, confidence={"voice-A": 99}))])
|
|
self.failure = failure
|
|
self.calls = []
|
|
self.submissions = []
|
|
self.uploads = []
|
|
self.authenticated = []
|
|
|
|
def __call__(self, command, **kwargs):
|
|
self.calls.append(command)
|
|
if command[0] == "ffmpeg":
|
|
Path(command[-1]).write_bytes(b"fLaC fake whole recording")
|
|
return subprocess.CompletedProcess(command, 0, stderr=b"")
|
|
output = Path(command[command.index("-o") + 1])
|
|
method = command[command.index("--request") + 1]
|
|
url = command[-1]
|
|
if method != "PUT":
|
|
header = Path(command[command.index("--header") + 1][1:])
|
|
assert header.stat().st_mode & 0o777 == 0o600
|
|
self.authenticated.append(header.read_text())
|
|
if self.failure and self.failure[0] in url:
|
|
status, body, code = self.failure[1:]
|
|
elif url.endswith("/media/input"):
|
|
self.media_request = json.loads(Path(command[command.index("--data-binary") + 1][1:]).read_text())
|
|
status, body, code = 200, {"url": "https://storage.example/file?secret=signature"}, 0
|
|
elif method == "PUT":
|
|
self.uploads.append(Path(command[command.index("--upload-file") + 1]).read_bytes())
|
|
status, body, code = 200, "", 0
|
|
elif url.endswith("/diarize"):
|
|
self.submissions.append(json.loads(Path(command[command.index("--data-binary") + 1][1:]).read_text()))
|
|
status, body, code = 200, {"jobId": "test-job", "status": "created"}, 0
|
|
else:
|
|
item = self.polls.pop(0)
|
|
status, body, code = item if isinstance(item, tuple) else (200, item, 0)
|
|
output.write_text(body if isinstance(body, str) else json.dumps(body))
|
|
return subprocess.CompletedProcess(command, code, stdout=str(status).encode(),
|
|
stderr=b"sensitive https://storage.example/file?secret=signature")
|
|
|
|
|
|
class RunTests(unittest.TestCase):
|
|
def execute(self, fake, directory, **kwargs):
|
|
audio = write_audio(directory)
|
|
with patch.dict(os.environ, {"PYANNOTE_API_KEY": "test-secret-key"}), \
|
|
patch.object(diarization.subprocess, "run", side_effect=fake), \
|
|
patch.object(diarization.time, "sleep"), contextlib.redirect_stderr(io.StringIO()) as log:
|
|
result = diarization.run(audio, Path(directory) / "work", **kwargs)
|
|
return result, log.getvalue()
|
|
|
|
def test_complete_flac_upload_exact_model_flags_and_private_request_cleanup(self):
|
|
fake = FakeTransport()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
result, _ = self.execute(fake, directory, speakers=2)
|
|
work = Path(directory) / "work"
|
|
self.assertEqual(json.loads((work / "result.json").read_text()), result)
|
|
self.assertEqual([p.name for p in work.iterdir()], ["result.json"])
|
|
self.assertEqual(fake.uploads, [b"fLaC fake whole recording"])
|
|
self.assertEqual(fake.submissions, [dict(url=fake.media_request["url"], model="precision-2",
|
|
exclusive=True, turnLevelConfidence=True, transcription=False, numSpeakers=2)])
|
|
self.assertRegex(fake.media_request["url"], r"^media://[a-f0-9-]+\.flac$")
|
|
encoding = fake.calls[0]
|
|
self.assertNotIn("-ss", encoding)
|
|
self.assertNotIn("-t", encoding)
|
|
self.assertIn("flac", encoding)
|
|
for command in fake.calls[1:]:
|
|
self.assertIn("--max-time", command)
|
|
self.assertIn("--connect-timeout", command)
|
|
self.assertIn("--proto", command)
|
|
self.assertNotIn("--location", command)
|
|
self.assertNotIn("test-secret-key", " ".join(command))
|
|
self.assertEqual(any(value.startswith("@") for value in command
|
|
if "authorization.header" in value),
|
|
command[command.index("--request") + 1] != "PUT")
|
|
self.assertEqual(fake.authenticated, ["Authorization: Bearer test-secret-key\n"] * 3)
|
|
|
|
def test_hosted_community_override_is_sent_unchanged(self):
|
|
fake = FakeTransport()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
self.execute(fake, directory, model="community-1")
|
|
self.assertEqual(fake.submissions[0]["model"], "community-1")
|
|
|
|
def test_saved_result_keeps_timelines_and_confidence_without_request_or_url_metadata(self):
|
|
original = job(turn(0, 10, confidence={"voice-A": 91}, url="https://secret/segment"),
|
|
request={"key": "test-secret-key"}, url="https://secret/job")
|
|
original["output"].update(diarization=[turn(2, 9, "b"), turn(0, 10, "a")],
|
|
warning="constraints adjusted https://secret/warning", url="https://secret/output")
|
|
fake = FakeTransport(polls=[original])
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
result, log = self.execute(fake, directory)
|
|
saved = (Path(directory) / "work" / "result.json").read_text()
|
|
self.assertEqual(set(result), {"jobId", "status", "output"})
|
|
self.assertEqual(set(result["output"]), {"diarization", "exclusiveDiarization"})
|
|
self.assertEqual(result["output"]["diarization"], original["output"]["diarization"])
|
|
self.assertEqual(result["output"]["exclusiveDiarization"][0]["confidence"], {"voice-A": 91})
|
|
for value in ("https://", "test-secret-key", "request", "warning"):
|
|
self.assertNotIn(value, saved)
|
|
self.assertIn("constraints adjusted", log)
|
|
self.assertNotIn("https://", log)
|
|
|
|
def test_automatic_count_is_omitted_and_created_pending_running_are_polled(self):
|
|
fake = FakeTransport(polls=[{"status": status} for status in ("created", "pending", "running")]
|
|
+ [job(turn(0, 10))])
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
_, log = self.execute(fake, directory)
|
|
self.assertNotIn("numSpeakers", fake.submissions[0])
|
|
self.assertEqual(len(fake.submissions), 1)
|
|
self.assertEqual(sum("GET" in call for call in fake.calls), 4)
|
|
self.assertIn("status: created", log)
|
|
self.assertIn("status: pending", log)
|
|
self.assertIn("status: running", log)
|
|
|
|
def test_transient_get_http_and_transport_errors_repoll_same_job(self):
|
|
fake = FakeTransport(polls=[(503, {"error": "try again"}, 0), (0, "", 28), job(turn(0, 10))])
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
_, log = self.execute(fake, directory)
|
|
urls = [call[-1] for call in fake.calls if "GET" in call]
|
|
self.assertEqual(urls, ["https://api.pyannote.ai/v1/jobs/test-job"] * 3)
|
|
self.assertEqual(len(fake.submissions), 1)
|
|
self.assertEqual(len(fake.uploads), 1)
|
|
self.assertIn("retrying the same job", log)
|
|
self.assertNotIn("secret=signature", log)
|
|
|
|
def test_post_and_upload_failures_are_never_retried_and_leave_no_result(self):
|
|
for endpoint in ("/media/input", "storage.example", "/diarize"):
|
|
fake = FakeTransport(failure=(endpoint, 503, {"error": "provider unavailable"}, 0))
|
|
with self.subTest(endpoint=endpoint), tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaisesRegex(ValueError, "HTTP 503"):
|
|
self.execute(fake, directory)
|
|
self.assertEqual(list((Path(directory) / "work").iterdir()), [])
|
|
self.assertEqual(sum(endpoint in call[-1] for call in fake.calls), 1)
|
|
|
|
def test_terminal_and_malformed_poll_failures_never_yield_partial_result(self):
|
|
failures = [dict(job(), status="failed"), dict(job(), status="canceled"),
|
|
{"status": "unexpected"}, {"status": "succeeded", "output": {}},
|
|
dict(job(turn(0, 10)), jobId="wrong-job"), job(turn(0, 31)),
|
|
(401, {"error": "bad key"}, 0), (200, "not JSON", 0)]
|
|
for response in failures:
|
|
with self.subTest(response=response), tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaises(ValueError):
|
|
self.execute(FakeTransport(polls=[response]), directory)
|
|
self.assertEqual(list((Path(directory) / "work").iterdir()), [])
|
|
|
|
def test_error_details_redact_key_and_urls(self):
|
|
fake = FakeTransport(failure=("/media/input", 401, {"error": {
|
|
"message": "invalid test-secret-key Bearer abc https://storage.example/?token=secret"}}, 0))
|
|
with tempfile.TemporaryDirectory() as directory, self.assertRaises(ValueError) as caught:
|
|
self.execute(fake, directory)
|
|
self.assertIn("invalid", str(caught.exception))
|
|
for secret in ("test-secret-key", "abc", "https://", "token=secret"):
|
|
self.assertNotIn(secret, str(caught.exception))
|
|
|
|
def test_missing_key_or_invalid_model_timeout_and_count_prevent_network(self):
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ, {}, clear=True), \
|
|
patch.object(diarization.subprocess, "run") as provider:
|
|
with self.assertRaisesRegex(ValueError, "PYANNOTE_API_KEY"):
|
|
diarization.run(write_audio(directory), directory)
|
|
provider.assert_not_called()
|
|
for kwargs in ({"model": "google/gemini"}, {"timeout": 0}, {"timeout": float("inf")},
|
|
{"timeout": True}, {"speakers": 0}):
|
|
with self.subTest(kwargs=kwargs), tempfile.TemporaryDirectory() as directory:
|
|
fake = FakeTransport()
|
|
with self.assertRaises(ValueError):
|
|
self.execute(fake, directory, **kwargs)
|
|
self.assertEqual(fake.calls, [])
|
|
|
|
def test_overall_deadline_is_bounded_even_while_waiting(self):
|
|
with tempfile.TemporaryDirectory() as directory, patch.object(diarization.time, "monotonic", return_value=10):
|
|
client = diarization.Client(directory, "key", 9)
|
|
with self.assertRaisesRegex(ValueError, "timed out"):
|
|
client.request("GET", "https://api.pyannote.ai/v1/jobs/one")
|
|
with tempfile.TemporaryDirectory() as directory, patch.object(diarization.time, "monotonic", return_value=10), \
|
|
patch.object(diarization.time, "sleep") as sleep:
|
|
diarization.Client(directory, "key", 10.5).pause()
|
|
sleep.assert_called_once_with(0.5)
|
|
|
|
def test_pending_job_times_out_without_restarting_upload_or_submission(self):
|
|
fake = FakeTransport(polls=[{"status": "running"}] * 10)
|
|
clock = [0]
|
|
with tempfile.TemporaryDirectory() as directory, \
|
|
patch.dict(os.environ, {"PYANNOTE_API_KEY": "test-secret-key"}), \
|
|
patch.object(diarization.subprocess, "run", side_effect=fake), \
|
|
patch.object(diarization.time, "monotonic", side_effect=lambda: clock[0]), \
|
|
patch.object(diarization.time, "sleep", side_effect=lambda seconds: clock.__setitem__(0, clock[0] + seconds)), \
|
|
contextlib.redirect_stderr(io.StringIO()), self.assertRaisesRegex(ValueError, "timed out"):
|
|
diarization.run(write_audio(directory), Path(directory) / "work", timeout=5)
|
|
self.assertEqual(len(fake.submissions), 1)
|
|
self.assertEqual(len(fake.uploads), 1)
|
|
self.assertEqual(sum("GET" in call for call in fake.calls), 3)
|
|
|
|
def test_upload_size_limit_prevents_any_http_request(self):
|
|
fake = FakeTransport()
|
|
with tempfile.TemporaryDirectory() as directory, patch.object(diarization, "MAX_UPLOAD_BYTES", 2), \
|
|
self.assertRaisesRegex(ValueError, "1 GiB"):
|
|
self.execute(fake, directory)
|
|
self.assertEqual(len(fake.calls), 1)
|
|
|
|
def test_invalid_upload_url_never_receives_audio_or_bearer(self):
|
|
for url in ("http://storage.example/upload", "file:///tmp/upload", "https://key@storage.example/upload"):
|
|
with self.subTest(url=url), tempfile.TemporaryDirectory() as directory, \
|
|
patch.object(diarization.subprocess, "run") as transport, self.assertRaises(ValueError):
|
|
diarization.Client(directory, "key", float("inf")).request("PUT", url, authenticated=False)
|
|
transport.assert_not_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|