From d45fe41fe9457708f2d7654ea2a6f43559322eb1 Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Wed, 12 Aug 2026 09:22:01 +0200 Subject: [PATCH] refactor(soniox): make cue grouping functional to satisfy type discipline gate --- litellm/llms/soniox/common_utils.py | 237 ++++++++++-------- ...niox_audio_transcription_transformation.py | 2 +- type-discipline-budget.json | 6 +- 3 files changed, 136 insertions(+), 109 deletions(-) diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index ee25465e128..fc85e2b719c 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -3,6 +3,9 @@ Shared utilities for the Soniox provider (https://soniox.com). """ import unicodedata +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import reduce from typing import Any, Final from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -120,9 +123,9 @@ _CUE_MAX_DURATION_MS: Final[int] = 7000 _CUE_GAP_MS: Final[int] = 700 -_SENTENCE_END_CHARS = (".", "!", "?", "。", "!", "?", "؟", "۔", "।", "॥", "։", "።") +_SENTENCE_END_CHARS: Final = (".", "!", "?", "。", "!", "?", "؟", "۔", "।", "॥", "։", "።") -_CJK_RANGES = ( +_CJK_RANGES: Final = ( (0x3400, 0x4DBF), (0x4E00, 0x9FFF), (0xF900, 0xFAFF), @@ -131,13 +134,13 @@ _CJK_RANGES = ( (0x31F0, 0x31FF), ) -_CJK_NO_BREAK_BEFORE = "、。,.!?:;・ー…」』)〉》】〕" +_CJK_NO_BREAK_BEFORE: Final = "、。,.!?:;・ー…」』)〉》】〕" -_CJK_NO_BREAK_AFTER = "「『(〈《【〔" +_CJK_NO_BREAK_AFTER: Final = "「『(〈《【〔" def _is_cjk(ch: str) -> bool: - cp = ord(ch) + cp: Final = ord(ch) return any(lo <= cp <= hi for lo, hi in _CJK_RANGES) @@ -175,7 +178,47 @@ def _format_timestamp_vtt(ms: int) -> str: return f"{hours:02d}:{minutes:02d}:{seconds:02d}.{millis:03d}" -def _merge_tokens_into_words(tokens: list[dict[str, Any]]) -> list[dict[str, Any]]: +@dataclass(frozen=True, slots=True) +class _Word: + text: str + start_ms: int | None + end_ms: int | None + speaker: str | int | None + + +@dataclass(frozen=True, slots=True) +class _Cue: + start_ms: int + end_ms: int + text: str + + +def _keeps_token(token: Mapping[str, Any]) -> bool: + text: Final = token.get("text", "") + return isinstance(text, str) and text != "" and token.get("translation_status") != "translation" + + +def _starts_new_word(prev: Mapping[str, Any], token: Mapping[str, Any]) -> bool: + prev_last: Final = prev["text"][-1:] + first: Final = token["text"][0] + return ( + first.isspace() + or prev_last.isspace() + or token.get("speaker") != prev.get("speaker") + or _is_cjk_word_boundary(prev_last, first) + ) + + +def _build_word(group: Sequence[Mapping[str, Any]]) -> _Word: + return _Word( + text="".join(t["text"] for t in group), + start_ms=next((t.get("start_ms") for t in group if t.get("start_ms") is not None), None), + end_ms=next((t.get("end_ms") for t in reversed(group) if t.get("end_ms") is not None), None), + speaker=group[0].get("speaker"), + ) + + +def _merge_tokens_into_words(tokens: Sequence[Mapping[str, Any]]) -> tuple[_Word, ...]: """ Merge Soniox subword tokens (e.g. ``"Hel"``, ``"lo"``) into whole words. @@ -190,50 +233,62 @@ def _merge_tokens_into_words(tokens: list[dict[str, Any]]) -> list[dict[str, Any Soniox does not timestamp them, so they cannot be aligned to the audio and would otherwise mix translated text into original-language cues. """ - words: list[dict[str, Any]] = [] - for token in tokens: - text = token.get("text", "") - if not isinstance(text, str) or text == "": - continue - if token.get("translation_status") == "translation": - continue - is_continuation = ( - bool(words) - and not text[0].isspace() - and not words[-1]["text"][-1:].isspace() - and token.get("speaker") == words[-1]["speaker"] - and not _is_cjk_word_boundary(words[-1]["text"][-1:], text[0]) - ) - if is_continuation: - last = words[-1] - last["text"] += text - if last["start_ms"] is None: - last["start_ms"] = token.get("start_ms") - if token.get("end_ms") is not None: - last["end_ms"] = token.get("end_ms") - else: - words.append( - { - "text": text, - "start_ms": token.get("start_ms"), - "end_ms": token.get("end_ms"), - "speaker": token.get("speaker"), - } - ) - return words + kept: Final = tuple(t for t in tokens if _keeps_token(t)) + starts: Final = tuple(i for i, t in enumerate(kept) if i == 0 or _starts_new_word(kept[i - 1], t)) + return tuple(_build_word(kept[begin:end]) for begin, end in zip(starts, (*starts[1:], len(kept)))) -def _group_tokens_into_cues( - tokens: list[dict[str, Any]], -) -> list[dict[str, Any]]: +def _cue_start(ws: Sequence[_Word]) -> int | None: + return next((w.start_ms for w in ws if w.start_ms is not None), None) + + +def _cue_end(ws: Sequence[_Word]) -> int | None: + return next((w.end_ms for w in reversed(ws) if w.end_ms is not None), _cue_start(ws)) + + +def _cue_text(ws: Sequence[_Word]) -> str: + return "".join(w.text for w in ws).strip() + + +def _should_break(cue: Sequence[_Word], word: _Word) -> bool: + speaker_changed: Final = word.speaker is not None and any( + w.speaker is not None and w.speaker != word.speaker for w in cue + ) + cue_start: Final = _cue_start(cue) + cue_end: Final = _cue_end(cue) + gap_exceeded: Final = word.start_ms is not None and cue_end is not None and (word.start_ms - cue_end) >= _CUE_GAP_MS + chars_exceeded: Final = _text_width(_cue_text(cue)) + _text_width(word.text) > _CUE_MAX_CHARS + word_end: Final = word.end_ms if word.end_ms is not None else word.start_ms + duration_exceeded: Final = ( + word_end is not None and cue_start is not None and (word_end - cue_start) > _CUE_MAX_DURATION_MS + ) + return speaker_changed or gap_exceeded or chars_exceeded or duration_exceeded + + +def _cue_start_indices(words: Sequence[_Word]) -> tuple[int, ...]: + def step(starts: tuple[int, ...], index: int) -> tuple[int, ...]: + if words[index - 1].text.rstrip().endswith(_SENTENCE_END_CHARS): + return (*starts, index) + if _should_break(words[starts[-1] : index], words[index]): + return (*starts, index) + return starts + + return reduce(step, range(1, len(words)), (0,)) if words else () + + +def _build_cue(ws: Sequence[_Word]) -> _Cue | None: + text: Final = _cue_text(ws) + start: Final = _cue_start(ws) + if not text or start is None: + return None + end: Final = _cue_end(ws) + return _Cue(start_ms=start, end_ms=end if end is not None else start, text=text) + + +def _group_tokens_into_cues(tokens: Sequence[Mapping[str, Any]]) -> tuple[_Cue, ...]: """ Group Soniox tokens into subtitle cues aligned to the actual speech. - Each cue has: - - start_ms: int - - end_ms: int - - text: str - Cues only ever break at word boundaries (Soniox tokens are subwords, so tokens are first merged into words). A new cue starts when: - the speaker changes (if diarization is on), @@ -245,50 +300,16 @@ def _group_tokens_into_cues( _CUE_MAX_DURATION_MS. A cue also ends after sentence-final punctuation, which keeps cue breaks at natural seams. Cue timestamps come straight from token timestamps; - words without timestamps stay attached to the surrounding cue. + words without timestamps stay attached to the surrounding cue, and a cue + whose words carry no timestamps at all is dropped. """ - words = _merge_tokens_into_words(tokens) - cues: list[dict[str, Any]] = [] - current: list[dict[str, Any]] = [] - - def _cue_start(ws: list[dict[str, Any]]) -> int | None: - return next((w["start_ms"] for w in ws if w["start_ms"] is not None), None) - - def _cue_end(ws: list[dict[str, Any]]) -> int | None: - return next((w["end_ms"] for w in reversed(ws) if w["end_ms"] is not None), _cue_start(ws)) - - def _cue_text(ws: list[dict[str, Any]]) -> str: - return "".join(w["text"] for w in ws).strip() - - def _flush() -> None: - text = _cue_text(current) - start = _cue_start(current) - if text and start is not None: - cues.append({"start_ms": start, "end_ms": _cue_end(current), "text": text}) - current.clear() - - for word in words: - if current: - start_ms = word["start_ms"] - cue_start = _cue_start(current) - cue_end = _cue_end(current) - speaker_changed = word["speaker"] is not None and any( - w["speaker"] is not None and w["speaker"] != word["speaker"] for w in current - ) - gap_exceeded = start_ms is not None and cue_end is not None and (start_ms - cue_end) >= _CUE_GAP_MS - chars_exceeded = _text_width(_cue_text(current)) + _text_width(word["text"]) > _CUE_MAX_CHARS - word_end = word["end_ms"] if word["end_ms"] is not None else start_ms - duration_exceeded = ( - word_end is not None and cue_start is not None and (word_end - cue_start) > _CUE_MAX_DURATION_MS - ) - if speaker_changed or gap_exceeded or chars_exceeded or duration_exceeded: - _flush() - current.append(word) - if word["text"].rstrip().endswith(_SENTENCE_END_CHARS): - _flush() - - _flush() - return cues + words: Final = _merge_tokens_into_words(tokens) + starts: Final = _cue_start_indices(words) + return tuple( + cue + for begin, end in zip(starts, (*starts[1:], len(words))) + if (cue := _build_cue(words[begin:end])) is not None + ) def render_soniox_tokens_as_srt(tokens: list[dict[str, Any]]) -> str: @@ -301,16 +322,16 @@ def render_soniox_tokens_as_srt(tokens: list[dict[str, Any]]) -> str: if not cues: return "" - lines: Final[list[str]] = [] - for idx, cue in enumerate(cues, start=1): - start = _format_timestamp_srt(cue["start_ms"]) - end = _format_timestamp_srt(cue["end_ms"]) - lines.append(str(idx)) - lines.append(f"{start} --> {end}") - lines.append(cue["text"]) - lines.append("") # blank line between cues - - return "\n".join(lines) + return "\n".join( + line + for idx, cue in enumerate(cues, start=1) + for line in ( + str(idx), + f"{_format_timestamp_srt(cue.start_ms)} --> {_format_timestamp_srt(cue.end_ms)}", + cue.text, + "", + ) + ) def render_soniox_tokens_as_vtt(tokens: list[dict[str, Any]]) -> str: @@ -321,12 +342,18 @@ def render_soniox_tokens_as_vtt(tokens: list[dict[str, Any]]) -> str: """ cues: Final = _group_tokens_into_cues(tokens) - lines: Final[list[str]] = ["WEBVTT", ""] - for cue in cues: - start = _format_timestamp_vtt(cue["start_ms"]) - end = _format_timestamp_vtt(cue["end_ms"]) - lines.append(f"{start} --> {end}") - lines.append(cue["text"]) - lines.append("") # blank line between cues + lines: Final = ( + "WEBVTT", + "", + *( + line + for cue in cues + for line in ( + f"{_format_timestamp_vtt(cue.start_ms)} --> {_format_timestamp_vtt(cue.end_ms)}", + cue.text, + "", + ) + ), + ) return "\n".join(lines) diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index cf7b5679dfe..cb053ac5f1e 100644 --- a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -562,7 +562,7 @@ class TestMultilingualCueGrouping: {"text": "保存", "start_ms": 250, "end_ms": 400}, ] words = _merge_tokens_into_words(tokens) - assert [w["text"] for w in words] == ["編", "集、", "保存"] + assert [w.text for w in words] == ["編", "集、", "保存"] class TestRenderSonioxTokensAsVtt: diff --git a/type-discipline-budget.json b/type-discipline-budget.json index fdacf375844..daa3e89fa0a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23003 + "limit": 22994 }, "LIT002": { - "limit": 27146 + "limit": 27139 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16731 + "limit": 16727 }, "LIT011": { "limit": 5596