updates

modules/capabilities/transcription/diarize-transcript.py

@@ -0,0 +1,118 @@
+import argparse
+import json
+import re
+from dataclasses import dataclass
+from itertools import groupby
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class Turn:
+    start: float
+    end: float
+    speaker: str
+
+
+@dataclass
+class Token:
+    start: float
+    end: float
+    text: str
+
+
+def load_tokens(path: Path) -> list[Token]:
+    with path.open(encoding="utf-8") as file:
+        transcript = json.load(file)
+
+    tokens: list[Token] = []
+    pending = ""
+    # Whisper reserves all IDs from end-of-text onward for control tokens.
+    end_of_text = 50256 + int(transcript["model"]["multilingual"])
+
+    for segment in transcript["transcription"]:
+        for raw_token in segment["tokens"]:
+            if raw_token["id"] >= end_of_text:
+                continue
+            text = raw_token["text"]
+            offsets = raw_token.get("offsets")
+
+            if not text:
+                continue
+            if offsets is None:
+                if tokens:
+                    tokens[-1].text += text
+                else:
+                    pending += text
+                continue
+
+            tokens.append(
+                Token(offsets["from"] / 1000, offsets["to"] / 1000, pending + text)
+            )
+            pending = ""
+
+    if not tokens:
+        raise ValueError("Whisper produced no timed text")
+
+    return tokens
+
+
+def load_turns(audio: Path, model: Path) -> list[Turn]:
+    from pyannote.audio import Pipeline
+
+    pipeline = Pipeline.from_pretrained(model)
+    result = pipeline(str(audio))
+    turns = [
+        Turn(segment.start, segment.end, speaker)
+        for segment, _, speaker in result.exclusive_speaker_diarization.itertracks(
+            yield_label=True
+        )
+    ]
+    if not turns:
+        raise RuntimeError("Pyannote detected no speech")
+
+    return turns
+
+
+def token_speaker(token: Token, turns: list[Turn]) -> str:
+    midpoint = (token.start + token.end) / 2
+
+    def rank(turn: Turn) -> tuple[float, float]:
+        overlap = max(0.0, min(token.end, turn.end) - max(token.start, turn.start))
+        distance = max(turn.start - midpoint, midpoint - turn.end, 0.0)
+        return overlap, -distance
+
+    return max(turns, key=rank).speaker
+
+
+def render(tokens: list[Token], turns: list[Turn]) -> str:
+    speaker_numbers: dict[str, int] = {}
+    rendered = []
+    for speaker, group in groupby(
+        tokens, key=lambda token: token_speaker(token, turns)
+    ):
+        number = speaker_numbers.setdefault(speaker, len(speaker_numbers) + 1)
+        text = "".join(token.text for token in group)
+        text = re.sub(r"\s+", " ", text).strip()
+        if text:
+            rendered.append(f"**Speaker {number}:**\n\n{text}")
+
+    if not rendered:
+        raise ValueError("Whisper produced no transcript text")
+    return "\n\n".join(rendered) + "\n"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--audio", required=True, type=Path)
+    parser.add_argument("--model", required=True, type=Path)
+    parser.add_argument("--transcript", required=True, type=Path)
+    parser.add_argument("--output", required=True, type=Path)
+    args = parser.parse_args()
+
+    tokens = load_tokens(args.transcript)
+    turns = load_turns(args.audio, args.model)
+    args.output.write_text(render(tokens, turns), encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()