Files
transcribe/diarization.py
T
ddidderr 3042b7743b feat!: use whole-recording Pyannote speaker diarization
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.
2026-09-12 11:18:30 +02:00

413 lines
20 KiB
Python

"""Run whole-recording Pyannote diarization and align its speakers to MAI words.
Only speaker labels change. A word needs a unique speaker covering at least
half of its duration. If there is no overlap, a nearest turn is usable only
when every point of the word lies within 200 ms of that turn. Equal candidates
remain unknown, including zero-duration words exactly at a speaker boundary.
"""
import argparse
from bisect import bisect_left, bisect_right
import json
import math
import os
from pathlib import Path
import re
import subprocess
import sys
import tempfile
import time
from urllib.parse import quote, urlsplit
import uuid
import wave
from subtitles import convert, normalized_words
from transcript import render_text
API_URL = "https://api.pyannote.ai/v1"
DEFAULT_TIMEOUT = 3600
POLL_INTERVAL = 2.0
MAX_NEAREST_GAP = 0.2
MAX_AUDIO_SECONDS = 24 * 60 * 60
MAX_UPLOAD_BYTES = 1024 ** 3
TRANSIENT_HTTP = {408, 425, 429, 500, 502, 503, 504}
EPSILON = 1e-9
def check_speakers(speakers):
if speakers is not None and (type(speakers) is not int or speakers < 1):
raise ValueError("The known speaker count must be a positive integer")
def unique_object(pairs):
result = {}
for key, value in pairs:
if key in result:
raise ValueError("Duplicate JSON key")
result[key] = value
return result
def read_json(path):
with open(path, encoding="utf-8") as source:
return json.load(source, object_pairs_hook=unique_object)
def finite_number(value):
return type(value) in (int, float) and math.isfinite(value)
def validate_segments(segments, exclusive=True):
if not isinstance(segments, list):
raise ValueError("Pyannote result is missing the exclusive diarization timeline"
if exclusive else "Invalid Pyannote standard diarization timeline")
previous_end = 0
for segment in segments:
if not isinstance(segment, dict):
raise ValueError("Invalid Pyannote diarization turn")
start, end, speaker = (segment.get(key) for key in ("start", "end", "speaker"))
if (not finite_number(start) or not finite_number(end)
or not 0 <= start < end or (exclusive and start < previous_end)
or not isinstance(speaker, str) or not speaker.strip()):
raise ValueError("Pyannote turns must have finite valid times and speaker IDs; "
"exclusive turns must be ordered and nonoverlapping")
confidence = segment.get("confidence")
if confidence is not None and (
not isinstance(confidence, dict)
or any(not isinstance(key, str) or not key.strip()
or not finite_number(value) or not 0 <= value <= 100
for key, value in confidence.items())):
raise ValueError("Invalid Pyannote turn confidence scores")
previous_end = end
return segments
def validated_timeline(job):
"""Require a completed exclusive timeline; never sort or repair its turns."""
if (not isinstance(job, dict) or job.get("status") != "succeeded"
or job.get("error") or not isinstance(job.get("output"), dict)
or job["output"].get("error")):
raise ValueError("Pyannote result must be a succeeded job with complete output")
return validate_segments(job["output"].get("exclusiveDiarization"))
def retained_result(job, job_id):
"""Keep only validated diarization data, never input URLs or request metadata."""
validated_timeline(job)
output = {}
for name in ("exclusiveDiarization", "diarization"):
if name in job["output"]:
segments = validate_segments(job["output"][name], exclusive=name == "exclusiveDiarization")
output[name] = [{key: value for key, value in segment.items()
if key in ("start", "end", "speaker", "confidence")}
for segment in segments]
return {"jobId": job_id, "status": "succeeded", "output": output}
def align(response, job, speakers=None):
"""Replace local labels using the complete audio timeline, without editing words."""
check_speakers(speakers)
# Validate the originals without taking normalized copies: those normalize
# whitespace and clamp regressing STT timestamps, which must survive here.
normalized_words(response, "Speaker alignment")
segments = validated_timeline(job)
if not segments and any(word["word"].strip() for word in response["words"]):
raise ValueError("Pyannote returned an empty speaker timeline for a nonempty transcript")
recording_speakers = {segment["speaker"] for segment in segments}
labels = {}
starts = [segment["start"] for segment in segments]
ends = [segment["end"] for segment in segments]
counts = {"unknown": 0, "nearest": 0}
result_words = []
for original in response["words"]:
start, end = original["start"], original["end"]
support = {}
if end > start:
index = bisect_right(ends, start)
while index < len(segments) and starts[index] < end:
segment = segments[index]
overlap = min(end, segment["end"]) - max(start, segment["start"])
if overlap > 0:
key = segment["speaker"]
support[key] = support.get(key, 0) + overlap
index += 1
else:
index = bisect_left(ends, start)
while index < len(segments) and starts[index] <= start:
support[segments[index]["speaker"]] = 1
index += 1
selected = None
used_nearest = False
if support:
best = max(support.values())
winners = [key for key, amount in support.items() if abs(amount - best) <= EPSILON]
if len(winners) == 1 and best + EPSILON >= (end - start) / 2:
selected = winners[0]
elif segments:
# Only adjacent turns can be nearest in an ordered exclusive
# timeline. The bound applies to the whole word, not just its edge.
index = bisect_right(starts, start)
distances = {}
for candidate in (index - 1, index):
if 0 <= candidate < len(segments):
segment = segments[candidate]
distance = max(segment["start"] - start, end - segment["end"], 0)
key = segment["speaker"]
distances[key] = min(distances.get(key, math.inf), distance)
best = min(distances.values())
winners = [key for key, distance in distances.items() if abs(distance - best) <= EPSILON]
if len(winners) == 1 and best <= MAX_NEAREST_GAP + EPSILON:
selected = winners[0]
used_nearest = True
if not original["word"].strip():
selected = None
if selected is not None:
labels.setdefault(selected, len(labels))
counts["nearest"] += used_nearest
counts["unknown"] += selected is None
result_words.append(dict(original, speaker=labels[selected] if selected is not None else None))
print(f"Speaker alignment: {len(recording_speakers)} recording speaker(s), "
f"{len(result_words) - counts['unknown']}/{len(result_words)} words assigned.", file=sys.stderr)
if counts["nearest"]:
print(f"Warning: {counts['nearest']} word(s) used a nearest turn within "
f"{MAX_NEAREST_GAP:.1f} seconds of the entire word.", file=sys.stderr)
if counts["unknown"]:
print(f"Warning: {counts['unknown']} word(s) have ambiguous or missing speaker timing; "
"their speaker is unknown.", file=sys.stderr)
if speakers is not None and len(recording_speakers) != speakers:
print(f"Warning: expected {speakers} speakers, but Pyannote returned {len(recording_speakers)}; "
"the count is not enforced by merging identities.", file=sys.stderr)
return dict(response, words=result_words)
class RequestError(ValueError):
def __init__(self, message, transient=False):
super().__init__(message)
self.transient = transient
def safe_detail(payload, api_key):
"""Keep useful provider errors without echoing credentials or signed URLs."""
if not isinstance(payload, dict):
return ""
detail = payload.get("error", payload.get("detail", payload.get("message", "")))
if isinstance(detail, dict):
detail = detail.get("message", detail.get("detail", ""))
if not isinstance(detail, str):
return ""
detail = detail.replace(api_key, "[redacted]") if api_key else detail
detail = re.sub(r"(?:https?://|media://)\S+", "[redacted URL]", detail)
detail = re.sub(r"(?i)bearer\s+\S+", "Bearer [redacted]", detail)
return " ".join(detail.split())[:400]
class Client:
def __init__(self, directory, api_key, deadline):
self.directory = Path(directory)
self.api_key = api_key
self.deadline = deadline
self.auth_path = self.directory / "authorization.header"
def remaining(self):
remaining = self.deadline - time.monotonic()
if remaining <= 0:
raise ValueError("Pyannote operation timed out; a submitted remote job may continue")
return remaining
def request(self, method, url, *, payload=None, upload=None, authenticated=True):
parsed = urlsplit(url)
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password:
raise ValueError("Pyannote returned an invalid HTTPS upload URL")
response_path = self.directory / "http-response.json"
response_path.unlink(missing_ok=True)
timeout = min(self.remaining(), 300 if upload else 120)
command = ["curl", "--silent", "--show-error", "--request", method,
"--connect-timeout", str(min(20, timeout)), "--max-time", str(timeout),
"--proto", "=https", "-o", str(response_path), "--write-out", "%{http_code}"]
if authenticated:
# Keep credentials out of process arguments. The parent directory
# is private, and this file is removed with it after the operation.
descriptor = os.open(self.auth_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
with os.fdopen(descriptor, "w", encoding="utf-8") as header:
header.write(f"Authorization: Bearer {self.api_key}\n")
command.extend(["--header", f"@{self.auth_path}"])
if payload is not None:
request_path = self.directory / "http-request.json"
request_path.write_text(json.dumps(payload), encoding="utf-8")
command.extend(["--header", "Content-Type: application/json",
"--data-binary", f"@{request_path}"])
if upload is not None:
command.extend(["--header", "Content-Type: application/octet-stream",
"--upload-file", str(upload)])
command.append(url)
try:
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
timeout=timeout + 1, check=False)
except subprocess.TimeoutExpired as error:
raise RequestError("Pyannote HTTP request timed out", transient=True) from error
if result.returncode:
# curl diagnostics can contain a presigned storage URL. Report the
# exit status without copying their potentially sensitive contents.
raise RequestError(f"Pyannote HTTP transport failed (curl exit {result.returncode})",
transient=True)
status_text = result.stdout.decode("ascii", errors="replace").strip()
if not re.fullmatch(r"[1-5][0-9]{2}", status_text):
raise RequestError("Pyannote HTTP request returned no valid status code")
status = int(status_text)
if not 200 <= status < 300:
try:
body = read_json(response_path)
except (ValueError, OSError):
body = None
detail = safe_detail(body, self.api_key)
raise RequestError(f"Pyannote HTTP {status}" + (f": {detail}" if detail else ""),
transient=status in TRANSIENT_HTTP)
if upload is not None:
return None
try:
body = read_json(response_path)
except (ValueError, OSError) as error:
raise RequestError("Pyannote returned invalid JSON") from error
if not isinstance(body, dict):
raise RequestError("Pyannote returned a JSON value instead of an object")
if body.get("error") and body.get("status") not in ("failed", "canceled"):
detail = safe_detail(body, self.api_key)
raise RequestError("Pyannote returned an error" + (f": {detail}" if detail else ""))
return body
def pause(self):
time.sleep(min(POLL_INTERVAL, self.remaining()))
def run(audio, work, model="precision-2", speakers=None, timeout=DEFAULT_TIMEOUT):
"""Upload the complete lossless recording once, submit once, and poll its job."""
check_speakers(speakers)
if model not in ("precision-2", "community-1"):
raise ValueError("The speaker model must be precision-2 or community-1")
if not finite_number(timeout) or timeout <= 0:
raise ValueError("The Pyannote timeout must be a positive finite number")
api_key = os.environ.get("PYANNOTE_API_KEY", "")
if not api_key.strip() or "\r" in api_key or "\n" in api_key:
raise ValueError("PYANNOTE_API_KEY must contain a Pyannote API key")
work = Path(work)
work.mkdir(parents=True, exist_ok=True)
result_path = work / "result.json"
result_path.unlink(missing_ok=True)
with wave.open(str(audio), "rb") as recording:
if (recording.getnchannels(), recording.getsampwidth(), recording.getframerate()) != (1, 2, 16000):
raise ValueError("Pyannote requires the prepared mono 16 kHz 16-bit PCM WAV")
duration = recording.getnframes() / recording.getframerate()
if not 0 < duration <= MAX_AUDIO_SECONDS:
raise ValueError("Pyannote requires nonempty audio no longer than 24 hours")
deadline = time.monotonic() + timeout
# API requests/responses can contain temporary signed storage URLs. Only a
# validated final job survives this private upload directory's lifetime.
with tempfile.TemporaryDirectory(prefix=".upload-", dir=work) as temporary:
client = Client(temporary, api_key, deadline)
flac = Path(temporary) / "recording.flac"
print("Preparing the complete recording as lossless FLAC for Pyannote...", file=sys.stderr)
command = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-nostdin", "-y",
"-i", str(audio), "-map", "0:a:0", "-vn", "-c:a", "flac",
"-compression_level", "5", str(flac)]
try:
encoded = subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
timeout=client.remaining(), check=False)
except subprocess.TimeoutExpired as error:
raise ValueError("Pyannote audio encoding timed out") from error
if encoded.returncode or not flac.is_file() or flac.stat().st_size == 0:
raise ValueError("Could not encode the complete Pyannote recording as FLAC")
if flac.stat().st_size > MAX_UPLOAD_BYTES:
raise ValueError("Pyannote FLAC upload exceeds the 1 GiB limit")
media_url = f"media://{uuid.uuid4()}.flac"
media = client.request("POST", f"{API_URL}/media/input", payload={"url": media_url})
signed_url = media.get("url")
if not isinstance(signed_url, str):
raise ValueError("Pyannote did not return an upload URL")
print(f"Uploading the complete {duration / 60:.1f}-minute recording to Pyannote...", file=sys.stderr)
client.request("PUT", signed_url, upload=flac, authenticated=False)
payload = {"url": media_url, "model": model, "exclusive": True,
"turnLevelConfidence": True, "transcription": False}
if speakers is not None:
payload["numSpeakers"] = speakers
created = client.request("POST", f"{API_URL}/diarize", payload=payload)
warning = safe_detail({"message": created.get("warning")}, api_key)
if warning:
print(f"Warning: Pyannote: {warning}", file=sys.stderr)
job_id = created.get("jobId")
if not isinstance(job_id, str) or not job_id.strip() or any(ord(c) < 32 for c in job_id):
raise ValueError("Pyannote did not return a valid job ID")
print("Pyannote job submitted; waiting for whole-recording speaker detection...", file=sys.stderr)
previous_status = None
retries = 0
while True:
try:
job = client.request("GET", f"{API_URL}/jobs/{quote(job_id, safe='')}")
except RequestError as error:
if not error.transient:
raise
retries += 1
print(f"Warning: {error}; retrying the same job (attempt {retries}).", file=sys.stderr)
client.pause()
continue
if "jobId" in job and job["jobId"] != job_id:
raise ValueError("Pyannote returned a different job ID while polling")
status = job.get("status")
if status == "succeeded":
segments = validated_timeline(job)
if any(segment["end"] > duration + MAX_NEAREST_GAP for segment in segments):
raise ValueError("Pyannote returned speaker times outside the uploaded recording")
warning = safe_detail({"message": job["output"].get("warning")}, api_key)
if warning:
print(f"Warning: Pyannote: {warning}", file=sys.stderr)
retained = retained_result(job, job_id)
completed = Path(temporary) / "completed.json"
completed.write_text(json.dumps(retained, ensure_ascii=False) + "\n", encoding="utf-8")
completed.replace(result_path)
print("Pyannote whole-recording speaker detection succeeded.", file=sys.stderr)
return retained
if status in ("failed", "canceled"):
detail = safe_detail(job, api_key) or safe_detail(job.get("output"), api_key)
raise ValueError(f"Pyannote job {status}" + (f": {detail}" if detail else ""))
if status not in ("created", "pending", "running"):
raise ValueError("Pyannote job returned an unknown or missing status")
if status != previous_status:
print(f"Pyannote job status: {status}.", file=sys.stderr)
previous_status = status
client.pause()
def main():
parser = argparse.ArgumentParser(description=__doc__)
subparsers = parser.add_subparsers(dest="command", required=True)
runner = subparsers.add_parser("run")
runner.add_argument("audio")
runner.add_argument("work")
runner.add_argument("--model", default="precision-2")
runner.add_argument("--speakers", type=int)
runner.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT)
apply = subparsers.add_parser("apply")
apply.add_argument("transcript")
apply.add_argument("result")
apply.add_argument("--speakers", type=int)
output = apply.add_mutually_exclusive_group()
output.add_argument("--srt", action="store_true")
output.add_argument("--json", action="store_true")
args = parser.parse_args()
try:
if args.command == "run":
run(args.audio, args.work, args.model, args.speakers, args.timeout)
return
response = align(read_json(args.transcript), read_json(args.result), args.speakers)
if args.json:
rendered = json.dumps(response, ensure_ascii=False) + "\n"
else:
rendered = convert(response) if args.srt else render_text(response)
except (ValueError, OSError, wave.Error, EOFError) as error:
print(f"Speaker detection failed: {error}", file=sys.stderr)
sys.exit(1)
sys.stdout.write(rendered)
if __name__ == "__main__":
main()