updates

modules/capabilities/_diarize-transcript.py

@@ -2,6 +2,7 @@ import argparse
 import json
 import re
 from dataclasses import dataclass
+from itertools import groupby
 from pathlib import Path
 
 
@@ -25,9 +26,13 @@ def load_tokens(path: Path) -> list[Token]:
 
     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")
 
@@ -81,19 +86,12 @@ def token_speaker(token: Token, turns: list[Turn]) -> str:
 
 def render(tokens: list[Token], turns: list[Turn]) -> str:
     speaker_numbers: dict[str, int] = {}
-    blocks: list[tuple[int, str]] = []
-
-    for token in tokens:
-        speaker = token_speaker(token, turns)
-        number = speaker_numbers.setdefault(speaker, len(speaker_numbers) + 1)
-
-        if blocks and blocks[-1][0] == number:
-            blocks[-1] = number, blocks[-1][1] + token.text
-        else:
-            blocks.append((number, token.text))
-
     rendered = []
-    for number, text in blocks:
+    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}")