feat: reconcile speaker identities through OpenRouter audio
Short chunk overlaps often contain only one participant, so returning speakers receive new transcript IDs. MAI can also split one voice into several IDs within a single request, which overlap stitching cannot fix. Add an opt-in audio comparison pass: --speakers N selects Gemini 3.8 Flash with a known count, and --speaker-model allows another compatible model. Keep request-local identities separate, select up to three clean excerpts per ID, and compare all excerpts in one OpenRouter request. Gemini uses Google AI Studio because the Vertex backend limits prompts to one audio file. No local model or extra Python dependency is required. Only reconcile identities whose sample assignments agree. Preserve uncertain identities separately, reject incomplete or malformed replies, and retain the original words and timestamps. Bound sample count and request size, clip sampling to the real audio duration, and extend worker cleanup and buffered output to the comparison step. Document the sampled method's limitations and additional provider request. Test Plan: - python3 -B -m unittest -v: all 82 tests passed, including real MP3 encoding with fake transcription/comparison responses and cancellation. - bash -n transcribe: passed. - git diff --cached --check: passed. - Live API and voice accuracy unverified: this workspace could not access the configured 1Password credential. Mocked responses prove pipeline behavior, not model accuracy on the supplied recording.
This commit is contained in:
ddidderr committed
1 parent
1d03dcd587
commit
138c796399
9 files changed
+988
-19
No files matched your search
@@ -0,0 +1,276 @@
|
||||
import base64
|
||||
import contextlib
|
||||
import copy
|
||||
import io
|
||||
import json
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
import wave
|
||||
|
||||
import diarization
|
||||
from diarization import prepare, reconcile, select_samples
|
||||
|
||||
|
||||
def word(start, end, speaker, 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 answer(assignments, finish="stop"):
|
||||
return {"choices": [{"finish_reason": finish, "message": {
|
||||
"content": json.dumps({"assignments": [
|
||||
{"sample": index, "speaker": speaker} for index, speaker in assignments
|
||||
]}),
|
||||
}}]}
|
||||
|
||||
|
||||
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 SampleTests(unittest.TestCase):
|
||||
def test_separate_runs_preferred_and_source_is_unchanged(self):
|
||||
source = transcript(word(0, 24, 7), word(25, 28, 4), word(30, 34, 7),
|
||||
word(35, 38, None), word(40, 45, 7))
|
||||
before = copy.deepcopy(source)
|
||||
with contextlib.redirect_stderr(io.StringIO()):
|
||||
samples = select_samples(source)
|
||||
selected = [sample for sample in samples if sample["speaker"] == 7]
|
||||
self.assertEqual(len(selected), 3)
|
||||
self.assertTrue(any(0 <= sample["start"] < sample["end"] <= 24 for sample in selected))
|
||||
self.assertTrue(any(30 <= sample["start"] < sample["end"] <= 34 for sample in selected))
|
||||
self.assertTrue(any(40 <= sample["start"] < sample["end"] <= 45 for sample in selected))
|
||||
self.assertEqual(source, before)
|
||||
self.assertEqual([sample["id"] for sample in samples], list(range(len(samples))))
|
||||
|
||||
def test_long_run_gives_three_separated_bounded_excerpts(self):
|
||||
samples = select_samples(transcript(word(0, 60, 0)))
|
||||
self.assertEqual(len(samples), 3)
|
||||
for sample in samples:
|
||||
self.assertGreaterEqual(sample["end"] - sample["start"], 1.5)
|
||||
self.assertLessEqual(sample["end"] - sample["start"], 8)
|
||||
for previous, current in zip(samples, samples[1:]):
|
||||
self.assertLess(previous["end"], current["start"])
|
||||
|
||||
def test_two_runs_give_two_samples_and_pause_or_unknown_never_bridged(self):
|
||||
source = transcript(word(0, 2, 0), word(4, 6, 0), word(6, 7, None),
|
||||
word(7, 8, 0), word(9, 10, 1))
|
||||
log = io.StringIO()
|
||||
with contextlib.redirect_stderr(log):
|
||||
samples = select_samples(source)
|
||||
self.assertEqual([(s["start"], s["end"]) for s in samples], [(0, 2), (4, 6)])
|
||||
self.assertIn("1 have none", log.getvalue())
|
||||
|
||||
def test_overlapping_foreign_or_unknown_words_are_excluded(self):
|
||||
for interfering in (None, 1):
|
||||
with self.subTest(interfering=interfering), contextlib.redirect_stderr(io.StringIO()):
|
||||
samples = select_samples(transcript(word(0, 20, 0), word(5, 7, interfering)))
|
||||
self.assertTrue(samples)
|
||||
self.assertTrue(all(s["end"] <= 5 or s["start"] >= 7 for s in samples))
|
||||
self.assertTrue(all(s["speaker"] == 0 for s in samples))
|
||||
|
||||
def test_same_voice_overlapping_runs_do_not_duplicate_audio(self):
|
||||
with contextlib.redirect_stderr(io.StringIO()):
|
||||
samples = select_samples(transcript(word(0, 30, 0), word(5, 7, 1), word(9, 30, 0)))
|
||||
samples.sort(key=lambda sample: sample["start"])
|
||||
self.assertTrue(all(left["end"] <= right["start"] for left, right in zip(samples, samples[1:])))
|
||||
|
||||
def test_single_excerpt_warns_and_short_or_unknown_speech_has_none(self):
|
||||
log = io.StringIO()
|
||||
with contextlib.redirect_stderr(log):
|
||||
samples = select_samples(transcript(word(0, 2, 0), word(3, 4, 1), word(5, 9, None)))
|
||||
self.assertEqual(len(samples), 1)
|
||||
self.assertIn("1 provisional speaker(s) have only one", log.getvalue())
|
||||
self.assertIn("1 have none", log.getvalue())
|
||||
|
||||
def test_excess_samples_fail_instead_of_silently_dropping_speakers(self):
|
||||
with patch.object(diarization, "MAX_SAMPLES", 2), self.assertRaisesRegex(ValueError, "limit is 2"):
|
||||
select_samples(transcript(word(0, 60, 0)))
|
||||
|
||||
def test_invalid_speaker_types_raise_value_error(self):
|
||||
for speaker in ([], {}, True, 1.5):
|
||||
with self.subTest(speaker=speaker), self.assertRaises(ValueError):
|
||||
select_samples(transcript(word(0, 8, speaker)))
|
||||
|
||||
|
||||
class ReconciliationTests(unittest.TestCase):
|
||||
def test_corrects_within_and_cross_request_fragmentation_from_shared_audio(self):
|
||||
source = transcript(word(0, 12, 0, "First voice."), word(13, 25, 1, "Second voice."),
|
||||
word(26, 38, 2, "First voice split in same request."),
|
||||
word(601, 613, 3, "First voice in next request."),
|
||||
word(614, 626, 4, "Second voice in next request."))
|
||||
samples = select_samples(source)
|
||||
before = copy.deepcopy(source)
|
||||
response = answer([(s["id"], 1 if s["speaker"] in (0, 2, 3) else 0) for s in samples])
|
||||
with contextlib.redirect_stderr(io.StringIO()):
|
||||
result = reconcile(source, samples, response, speakers=2)
|
||||
self.assertEqual([w["speaker"] for w in result["words"]], [0, 1, 0, 0, 1])
|
||||
self.assertEqual(source, before)
|
||||
self.assertEqual(result["text"], before["text"])
|
||||
self.assertEqual([{k: v for k, v in w.items() if k != "speaker"} for w in result["words"]],
|
||||
[{k: v for k, v in w.items() if k != "speaker"} for w in before["words"]])
|
||||
|
||||
def test_disagreement_null_and_unsampled_ids_stay_separate_without_namespace_collision(self):
|
||||
source = transcript(word(0, 12, 0), word(13, 25, 1), word(26, 38, "1"),
|
||||
word(40, 41, 2), word(42, 43, None))
|
||||
with contextlib.redirect_stderr(io.StringIO()):
|
||||
samples = select_samples(source)
|
||||
response = answer([(s["id"], 0 if s["speaker"] == 0 else
|
||||
(None if s["speaker"] == "1" else s["id"] % 2)) for s in samples])
|
||||
log = io.StringIO()
|
||||
with contextlib.redirect_stderr(log):
|
||||
result = reconcile(source, samples, response, speakers=2)
|
||||
self.assertEqual([w["speaker"] for w in result["words"]], [0, 1, 2, 3, None])
|
||||
self.assertIn("3 provisional", log.getvalue())
|
||||
self.assertIn("4 provisional IDs -> 4 transcript speakers", log.getvalue())
|
||||
self.assertIn("expected 2 speakers", log.getvalue())
|
||||
|
||||
def test_reconciliation_preserves_regressed_timestamps_and_unknown_key_absence(self):
|
||||
source = transcript(word(8, 12, 2), word(7, 9, 2), word(14, 16, None))
|
||||
del source["words"][-1]["speaker"]
|
||||
samples = [{"id": 0, "speaker": 2, "start": 8, "end": 12}]
|
||||
with contextlib.redirect_stderr(io.StringIO()):
|
||||
result = reconcile(source, samples, answer([(0, 0)]))
|
||||
self.assertEqual([w["start"] for w in result["words"]], [8, 7, 14])
|
||||
self.assertNotIn("speaker", result["words"][-1])
|
||||
|
||||
def test_empty_transcript_and_assignments(self):
|
||||
with contextlib.redirect_stderr(io.StringIO()):
|
||||
self.assertEqual(reconcile(transcript(), [], answer([])), transcript())
|
||||
|
||||
def test_single_sample_is_usable_without_assuming_provisional_numeric_id(self):
|
||||
source = transcript(word(0, 2, 8), word(3, 5, 9))
|
||||
with contextlib.redirect_stderr(io.StringIO()):
|
||||
samples = select_samples(source)
|
||||
result = reconcile(source, samples, answer([(0, 7), (1, 7)]))
|
||||
self.assertEqual([w["speaker"] for w in result["words"]], [0, 0])
|
||||
|
||||
def test_malformed_incomplete_duplicate_or_out_of_range_assignments_fail(self):
|
||||
source = transcript(word(0, 12, 0))
|
||||
samples = select_samples(source)
|
||||
valid = [(s["id"], 0) for s in samples]
|
||||
cases = [answer(valid[:-1]), answer(valid + [valid[0]]), answer(valid + [(99, 0)]),
|
||||
answer([(0, True), *valid[1:]]), answer([(0, -1), *valid[1:]]),
|
||||
answer([(0, 2), *valid[1:]]), answer([(False, 0), *valid[1:]]),
|
||||
answer([(0, "0"), *valid[1:]]), answer([(0, []), *valid[1:]]),
|
||||
answer(valid, "length"), answer(valid, None), {"error": {"message": "failure"}},
|
||||
{"choices": []}, {"choices": [None]}, {"choices": [{"finish_reason": "stop"}]},
|
||||
{"choices": [{"finish_reason": "stop", "message": {"content": "not JSON"}}]},
|
||||
{"choices": [{"finish_reason": "stop", "message": {"content": "[]"}}]},
|
||||
{"choices": [{"finish_reason": "stop", "message": {"content":
|
||||
'{"assignments":[{"sample":0,"speaker":0,"speaker":1}]}'}}]},
|
||||
{"choices": [{"finish_reason": "stop", "message": {"content": '{"assignments": {}}'}}]}]
|
||||
for response in cases:
|
||||
with self.subTest(response=response), self.assertRaises(ValueError):
|
||||
reconcile(source, samples, response, speakers=2)
|
||||
|
||||
def test_invalid_samples_or_known_count_fail(self):
|
||||
source = transcript(word(0, 12, 0))
|
||||
samples = select_samples(source)
|
||||
for invalid in (samples + [samples[0]], [dict(samples[0], speaker=[])],
|
||||
[dict(samples[0], speaker=99)], [dict(samples[0], start=float("nan"))],
|
||||
[dict(samples[0], end=True)], [dict(samples[0], id=True)], {}):
|
||||
with self.subTest(samples=invalid), self.assertRaises(ValueError):
|
||||
reconcile(source, invalid, answer([]))
|
||||
for count in (0, -1, True, 1.5):
|
||||
with self.subTest(count=count), self.assertRaises(ValueError):
|
||||
reconcile(source, samples, answer([]), count)
|
||||
|
||||
def test_cli_failure_emits_no_partial_transcript(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
paths = [Path(directory) / name for name in ("transcript.json", "samples.json", "response.json")]
|
||||
source = transcript(word(0, 12, 0, "Must not escape."))
|
||||
for path, value in zip(paths, (source, select_samples(source), answer([], "length"))):
|
||||
path.write_text(json.dumps(value), encoding="utf-8")
|
||||
result = subprocess.run([sys.executable, "-B", str(Path(diarization.__file__)), "apply",
|
||||
*(str(path) for path in paths)], capture_output=True, text=True)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertEqual(result.stdout, "")
|
||||
self.assertIn("completed response", result.stderr)
|
||||
|
||||
|
||||
class PrepareTests(unittest.TestCase):
|
||||
@staticmethod
|
||||
def fake_ffmpeg(command, **kwargs):
|
||||
Path(command[-1]).write_bytes(b"fake mp3")
|
||||
return subprocess.CompletedProcess(command, 0, stderr=b"")
|
||||
|
||||
def test_request_contains_only_labelled_audio_and_strict_schema(self):
|
||||
source = transcript(word(0, 12, 0, "Private transcript text."), word(13, 25, 1))
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
diarization.subprocess, "run", side_effect=self.fake_ffmpeg) as encoder:
|
||||
prepare(source, write_audio(directory), directory, "example/audio-model", 2)
|
||||
request_text = (Path(directory) / "request.json").read_text()
|
||||
request = json.loads(request_text)
|
||||
samples = json.loads((Path(directory) / "samples.json").read_text())
|
||||
self.assertNotIn("Private transcript text", request_text)
|
||||
self.assertEqual(request["model"], "example/audio-model")
|
||||
self.assertEqual(request["provider"], {"require_parameters": True})
|
||||
self.assertEqual(len(encoder.call_args_list), len(samples))
|
||||
self.assertIn("exactly 2 speakers", request["messages"][0]["content"])
|
||||
self.assertTrue(request["response_format"]["json_schema"]["strict"])
|
||||
content = request["messages"][1]["content"]
|
||||
for index, sample in enumerate(samples):
|
||||
self.assertEqual(content[index * 2]["text"], f"Sample ID {sample['id']}")
|
||||
self.assertEqual(base64.b64decode(content[index * 2 + 1]["input_audio"]["data"]), b"fake mp3")
|
||||
command = encoder.call_args_list[index].args[0]
|
||||
self.assertEqual(float(command[command.index("-ss") + 1]), sample["start"])
|
||||
self.assertEqual(float(command[command.index("-t") + 1]), sample["end"] - sample["start"])
|
||||
|
||||
def test_no_samples_removes_stale_request_and_does_not_encode(self):
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(diarization.subprocess, "run") as encoder:
|
||||
request = Path(directory) / "request.json"
|
||||
request.write_text("stale")
|
||||
with contextlib.redirect_stderr(io.StringIO()) as log:
|
||||
prepare(transcript(word(0, 1, 0)), write_audio(directory), directory, "audio-model")
|
||||
self.assertFalse(request.exists())
|
||||
self.assertEqual(json.loads((Path(directory) / "samples.json").read_text()), [])
|
||||
self.assertIn("no clean speaker samples", log.getvalue())
|
||||
encoder.assert_not_called()
|
||||
|
||||
def test_encoding_failure_and_request_size_limit_leave_no_request(self):
|
||||
source = transcript(word(0, 12, 0))
|
||||
for failure in (True, False):
|
||||
with self.subTest(failure=failure), tempfile.TemporaryDirectory() as directory:
|
||||
result = subprocess.CompletedProcess([], 1, stderr=b"bad audio")
|
||||
with patch.object(diarization.subprocess, "run", return_value=result,
|
||||
side_effect=None if failure else self.fake_ffmpeg), \
|
||||
patch.object(diarization, "MAX_REQUEST_BYTES", 10), self.assertRaises(ValueError):
|
||||
prepare(source, write_audio(directory), directory, "audio-model")
|
||||
self.assertFalse((Path(directory) / "request.json").exists())
|
||||
|
||||
def test_gemini_aliases_require_ai_studio_without_provider_fallback(self):
|
||||
for model in ("google/gemini-3.8-flash", "~google/gemini-3.8-flash"):
|
||||
with self.subTest(model=model), tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
diarization.subprocess, "run", side_effect=self.fake_ffmpeg):
|
||||
prepare(transcript(word(0, 12, 0)), write_audio(directory), directory, model)
|
||||
request = json.loads((Path(directory) / "request.json").read_text())
|
||||
self.assertEqual(request["provider"], {"require_parameters": True,
|
||||
"only": ["google-ai-studio"], "allow_fallbacks": False})
|
||||
|
||||
def test_samples_are_clipped_and_reselected_at_real_audio_end(self):
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
diarization.subprocess, "run", side_effect=self.fake_ffmpeg), \
|
||||
contextlib.redirect_stderr(io.StringIO()) as log:
|
||||
prepare(transcript(word(0, 12, 0), word(13, 15, 1)),
|
||||
write_audio(directory, 13.2), directory, "audio-model")
|
||||
samples = json.loads((Path(directory) / "samples.json").read_text())
|
||||
self.assertTrue(samples)
|
||||
self.assertTrue(all(sample["end"] <= 13.2 for sample in samples))
|
||||
self.assertTrue(all(sample["speaker"] == 0 for sample in samples))
|
||||
self.assertIn("clipped to the audio duration", log.getvalue())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user