mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
Merge pull request #34440 from dan2k3k4/litellm_soniox_srt_cue_grouping
fix(soniox): align synthesized SRT/VTT cues to real speech timing
This commit is contained in:
Mateo Wang
GitHub
commit
817bbe1dc6
4 changed files with 381 additions and 63 deletions
|
|
@ -1,19 +1,36 @@
|
|||
"""Provider-agnostic SRT/WebVTT subtitle synthesis from timestamped transcription tokens."""
|
||||
|
||||
import unicodedata
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate, chain
|
||||
from itertools import accumulate, groupby
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
CUE_MAX_TOKENS: Final = 15
|
||||
CUE_MAX_DURATION_MS: Final = 5000
|
||||
CUE_MAX_CHARS: Final = 84
|
||||
CUE_MAX_DURATION_MS: Final = 7000
|
||||
CUE_GAP_MS: Final = 700
|
||||
|
||||
SRT_RESPONSE_FORMAT: Final = "srt"
|
||||
VTT_RESPONSE_FORMAT: Final = "vtt"
|
||||
SUBTITLE_RESPONSE_FORMATS: Final = frozenset((SRT_RESPONSE_FORMAT, VTT_RESPONSE_FORMAT))
|
||||
|
||||
_SENTENCE_END_CHARS: Final = (".", "!", "?", "。", "！", "？", "؟", "۔", "।", "॥", "։", "።")
|
||||
|
||||
_CJK_RANGES: Final = (
|
||||
(0x3400, 0x4DBF),
|
||||
(0x4E00, 0x9FFF),
|
||||
(0xF900, 0xFAFF),
|
||||
(0x3040, 0x309F),
|
||||
(0x30A0, 0x30FF),
|
||||
(0x31F0, 0x31FF),
|
||||
)
|
||||
|
||||
_CJK_NO_BREAK_BEFORE: Final = "、。，．！？：；・ー…」』）〉》】〕"
|
||||
|
||||
_CJK_NO_BREAK_AFTER: Final = "「『（〈《【〔"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SubtitleToken:
|
||||
|
|
@ -31,69 +48,138 @@ class SubtitleCue:
|
|||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CueAccumulator:
|
||||
texts: tuple[str, ...] = ()
|
||||
start_ms: int | None = None
|
||||
end_ms: int | None = None
|
||||
speaker: str | int | None = None
|
||||
class _Word:
|
||||
text: str
|
||||
start_ms: int | None
|
||||
end_ms: int | None
|
||||
speaker: str | int | None
|
||||
|
||||
|
||||
def _completed_cue(accumulator: _CueAccumulator) -> tuple[SubtitleCue, ...]:
|
||||
if not accumulator.texts or accumulator.start_ms is None:
|
||||
return ()
|
||||
text: Final = "".join(accumulator.texts).strip()
|
||||
if not text:
|
||||
return ()
|
||||
end_ms: Final = accumulator.end_ms if accumulator.end_ms is not None else accumulator.start_ms
|
||||
return (SubtitleCue(start_ms=accumulator.start_ms, end_ms=end_ms, text=text),)
|
||||
def _is_cjk(ch: str) -> bool:
|
||||
cp: Final = ord(ch)
|
||||
return any(lo <= cp <= hi for lo, hi in _CJK_RANGES)
|
||||
|
||||
|
||||
def _cue_break_reached(accumulator: _CueAccumulator, token: SubtitleToken) -> bool:
|
||||
if len(accumulator.texts) >= CUE_MAX_TOKENS:
|
||||
return True
|
||||
def _is_cjk_word_boundary(prev_ch: str, next_ch: str) -> bool:
|
||||
if not (_is_cjk(prev_ch) or _is_cjk(next_ch)):
|
||||
return False
|
||||
return next_ch not in _CJK_NO_BREAK_BEFORE and prev_ch not in _CJK_NO_BREAK_AFTER
|
||||
|
||||
|
||||
def _text_width(text: str) -> int:
|
||||
return sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in text)
|
||||
|
||||
|
||||
def _starts_new_word(prev: SubtitleToken, token: SubtitleToken) -> bool:
|
||||
prev_last: Final = prev.text[-1:]
|
||||
first: Final = token.text[0]
|
||||
return (
|
||||
accumulator.start_ms is not None
|
||||
and token.start_ms is not None
|
||||
and token.start_ms - accumulator.start_ms >= CUE_MAX_DURATION_MS
|
||||
first.isspace()
|
||||
or prev_last.isspace()
|
||||
or token.speaker != prev.speaker
|
||||
or _is_cjk_word_boundary(prev_last, first)
|
||||
)
|
||||
|
||||
|
||||
_AbsorbStep = tuple[tuple[SubtitleCue, ...], _CueAccumulator]
|
||||
|
||||
|
||||
def _absorb_token(accumulator: _CueAccumulator, token: SubtitleToken) -> _AbsorbStep:
|
||||
if token.start_ms is None and accumulator.start_ms is None:
|
||||
return (), accumulator
|
||||
if token.speaker is not None and token.speaker != accumulator.speaker:
|
||||
return _completed_cue(accumulator), _CueAccumulator(
|
||||
texts=(token.text,),
|
||||
start_ms=token.start_ms,
|
||||
end_ms=token.end_ms,
|
||||
speaker=token.speaker,
|
||||
)
|
||||
if _cue_break_reached(accumulator, token):
|
||||
return _completed_cue(accumulator), _CueAccumulator(
|
||||
texts=(token.text,),
|
||||
start_ms=token.start_ms,
|
||||
end_ms=token.end_ms,
|
||||
speaker=accumulator.speaker,
|
||||
)
|
||||
return (), _CueAccumulator(
|
||||
texts=(*accumulator.texts, token.text),
|
||||
start_ms=accumulator.start_ms if accumulator.start_ms is not None else token.start_ms,
|
||||
end_ms=token.end_ms if token.end_ms is not None else accumulator.end_ms,
|
||||
speaker=accumulator.speaker,
|
||||
def _build_word(group: Sequence[SubtitleToken]) -> _Word:
|
||||
return _Word(
|
||||
text="".join(t.text for t in group),
|
||||
start_ms=next((t.start_ms for t in group if t.start_ms is not None), None),
|
||||
end_ms=next((t.end_ms for t in reversed(group) if t.end_ms is not None), None),
|
||||
speaker=group[0].speaker,
|
||||
)
|
||||
|
||||
|
||||
def _absorb_step(carry: _AbsorbStep, token: SubtitleToken) -> _AbsorbStep:
|
||||
return _absorb_token(carry[1], token)
|
||||
def _merge_tokens_into_words(tokens: Sequence[SubtitleToken]) -> tuple[_Word, ...]:
|
||||
"""
|
||||
Merge subword tokens (e.g. ``"Hel"``, ``"lo"``) into whole words.
|
||||
|
||||
A token starts a new word when its text begins with whitespace, when the
|
||||
previous token's text ends with whitespace, when the speaker changes, or
|
||||
at a CJK character boundary (CJK scripts carry no spaces, so without this
|
||||
an entire utterance would fuse into a single unbreakable "word"; CJK
|
||||
punctuation stays attached to the preceding character per kinsoku rules).
|
||||
Each word carries the first/last available timestamps of its tokens.
|
||||
"""
|
||||
kept: Final = tuple(t for t in tokens if t.text != "")
|
||||
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 _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 next_start(start: int, index: int) -> int:
|
||||
if words[index - 1].text.rstrip().endswith(_SENTENCE_END_CHARS):
|
||||
return index
|
||||
if _should_break(words[start:index], words[index]):
|
||||
return index
|
||||
return start
|
||||
|
||||
if not words:
|
||||
return ()
|
||||
return tuple(start for start, _ in groupby(accumulate(range(1, len(words)), next_start, initial=0)))
|
||||
|
||||
|
||||
def _build_cue(ws: Sequence[_Word]) -> SubtitleCue | 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 SubtitleCue(start_ms=start, end_ms=end if end is not None else start, text=text)
|
||||
|
||||
|
||||
def group_subtitle_tokens_into_cues(tokens: Sequence[SubtitleToken]) -> tuple[SubtitleCue, ...]:
|
||||
steps: Final = tuple(accumulate(tokens, _absorb_step, initial=((), _CueAccumulator())))
|
||||
completed: Final = chain.from_iterable(emitted for emitted, _ in steps)
|
||||
return (*completed, *_completed_cue(steps[-1][1]))
|
||||
"""
|
||||
Group transcription tokens into subtitle cues aligned to the actual speech.
|
||||
|
||||
Cues only ever break at word boundaries (tokens may be subwords, so they
|
||||
are first merged into words). A new cue starts when:
|
||||
- the speaker changes (if diarization is on),
|
||||
- a silence gap of at least CUE_GAP_MS separates two words, so
|
||||
subtitles never bridge pauses in speech,
|
||||
- adding the next word would exceed CUE_MAX_CHARS of display width
|
||||
(~two subtitle lines; East-Asian wide characters count double), or
|
||||
- adding the next word would make the cue span more than
|
||||
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, and a cue
|
||||
whose words carry no timestamps at all is dropped.
|
||||
"""
|
||||
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 _format_timestamp(total_ms: int, millis_separator: str) -> str:
|
||||
|
|
|
|||
|
|
@ -138,13 +138,25 @@ def _soniox_token_to_subtitle_token(token: SonioxToken) -> SubtitleToken:
|
|||
)
|
||||
|
||||
|
||||
def _subtitle_tokens(tokens: Sequence[SonioxToken]) -> tuple[SubtitleToken, ...]:
|
||||
"""
|
||||
Convert Soniox tokens for subtitle rendering, excluding translation tokens
|
||||
(``translation_status == "translation"``): Soniox does not timestamp them,
|
||||
so they cannot be aligned to the audio and would otherwise mix translated
|
||||
text into original-language cues.
|
||||
"""
|
||||
return tuple(
|
||||
_soniox_token_to_subtitle_token(token) for token in tokens if token.get("translation_status") != "translation"
|
||||
)
|
||||
|
||||
|
||||
def render_soniox_tokens_as_srt(tokens: Sequence[SonioxToken]) -> str:
|
||||
"""
|
||||
Render Soniox tokens as SRT (SubRip) subtitle format.
|
||||
|
||||
Returns an empty string if no tokens have timestamp data.
|
||||
"""
|
||||
return render_subtitle_tokens_as_srt(tuple(_soniox_token_to_subtitle_token(token) for token in tokens))
|
||||
return render_subtitle_tokens_as_srt(_subtitle_tokens(tokens))
|
||||
|
||||
|
||||
def render_soniox_tokens_as_vtt(tokens: Sequence[SonioxToken]) -> str:
|
||||
|
|
@ -153,4 +165,4 @@ def render_soniox_tokens_as_vtt(tokens: Sequence[SonioxToken]) -> str:
|
|||
|
||||
Returns the VTT header even if no cues are present.
|
||||
"""
|
||||
return render_subtitle_tokens_as_vtt(tuple(_soniox_token_to_subtitle_token(token) for token in tokens))
|
||||
return render_subtitle_tokens_as_vtt(_subtitle_tokens(tokens))
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
|
||||
SubtitleToken,
|
||||
_merge_tokens_into_words,
|
||||
render_subtitle_tokens_as_srt,
|
||||
render_subtitle_tokens_as_vtt,
|
||||
synthesize_subtitle_document,
|
||||
|
|
@ -23,25 +24,59 @@ class TestRenderSubtitleTokensAsSrt:
|
|||
"1\n00:00:00,000 --> 00:00:01,000\nHi.\n\n2\n00:00:01,500 --> 00:00:02,500\nHey.\n"
|
||||
)
|
||||
|
||||
def test_token_cap_starts_a_new_cue_after_15_tokens(self):
|
||||
tokens = tuple(
|
||||
SubtitleToken(text=f"{index} ", start_ms=index * 100, end_ms=index * 100 + 100) for index in range(16)
|
||||
def test_width_budget_starts_a_new_cue_at_word_boundaries(self):
|
||||
tokens = tuple(SubtitleToken(text="abcdefghi ", start_ms=i * 100, end_ms=i * 100 + 90) for i in range(20))
|
||||
result = render_subtitle_tokens_as_srt(tokens)
|
||||
texts = [cue.split("\n", 2)[2] for cue in result.strip().split("\n\n")]
|
||||
assert len(texts) == 3
|
||||
assert all(len(text) <= 84 for text in texts)
|
||||
assert all(set(text.split()) == {"abcdefghi"} for text in texts)
|
||||
|
||||
def test_duration_cap_starts_a_new_cue_before_word_crossing_7000ms(self):
|
||||
tokens = (
|
||||
SubtitleToken(text="Alpha ", start_ms=0, end_ms=3400),
|
||||
SubtitleToken(text="beta ", start_ms=3400, end_ms=6800),
|
||||
SubtitleToken(text="gamma", start_ms=6800, end_ms=7400),
|
||||
)
|
||||
assert render_subtitle_tokens_as_srt(tokens) == (
|
||||
"1\n00:00:00,000 --> 00:00:01,500\n0 1 2 3 4 5 6 7 8 9 10 11 12 13 14\n"
|
||||
"\n2\n00:00:01,500 --> 00:00:01,600\n15\n"
|
||||
"1\n00:00:00,000 --> 00:00:06,800\nAlpha beta\n\n2\n00:00:06,800 --> 00:00:07,400\ngamma\n"
|
||||
)
|
||||
|
||||
def test_duration_cap_starts_a_new_cue_at_5000ms(self):
|
||||
def test_silence_gap_starts_a_new_cue(self):
|
||||
tokens = (
|
||||
SubtitleToken(text="Alpha ", start_ms=0, end_ms=400),
|
||||
SubtitleToken(text="beta ", start_ms=2000, end_ms=2400),
|
||||
SubtitleToken(text="gamma.", start_ms=5000, end_ms=5400),
|
||||
SubtitleToken(text="beta", start_ms=2000, end_ms=2400),
|
||||
)
|
||||
assert render_subtitle_tokens_as_srt(tokens) == (
|
||||
"1\n00:00:00,000 --> 00:00:02,400\nAlpha beta\n\n2\n00:00:05,000 --> 00:00:05,400\ngamma.\n"
|
||||
"1\n00:00:00,000 --> 00:00:00,400\nAlpha\n\n2\n00:00:02,000 --> 00:00:02,400\nbeta\n"
|
||||
)
|
||||
|
||||
def test_sentence_final_punctuation_starts_a_new_cue(self):
|
||||
tokens = (
|
||||
SubtitleToken(text="Done. ", start_ms=0, end_ms=400),
|
||||
SubtitleToken(text="Next", start_ms=500, end_ms=800),
|
||||
)
|
||||
assert render_subtitle_tokens_as_srt(tokens) == (
|
||||
"1\n00:00:00,000 --> 00:00:00,400\nDone.\n\n2\n00:00:00,500 --> 00:00:00,800\nNext\n"
|
||||
)
|
||||
|
||||
def test_subword_tokens_merge_into_words_before_grouping(self):
|
||||
tokens = (
|
||||
SubtitleToken(text=" hel", start_ms=0, end_ms=150),
|
||||
SubtitleToken(text="lo", start_ms=150, end_ms=300),
|
||||
SubtitleToken(text=" world.", start_ms=350, end_ms=600),
|
||||
)
|
||||
assert render_subtitle_tokens_as_srt(tokens) == "1\n00:00:00,000 --> 00:00:00,600\nhello world.\n"
|
||||
|
||||
def test_cjk_tokens_merge_and_keep_punctuation_attached(self):
|
||||
tokens = (
|
||||
SubtitleToken(text="編", start_ms=0, end_ms=100),
|
||||
SubtitleToken(text="集", start_ms=100, end_ms=200),
|
||||
SubtitleToken(text="、", start_ms=200, end_ms=250),
|
||||
SubtitleToken(text="保存", start_ms=250, end_ms=400),
|
||||
)
|
||||
assert [word.text for word in _merge_tokens_into_words(tokens)] == ["編", "集、", "保存"]
|
||||
|
||||
def test_timestampless_token_joins_the_current_cue(self):
|
||||
tokens = (
|
||||
SubtitleToken(text="Hello ", start_ms=0, end_ms=500),
|
||||
|
|
|
|||
|
|
@ -369,6 +369,191 @@ class TestRenderSonioxTokensAsSrt:
|
|||
assert "01:01:01,000" in result
|
||||
|
||||
|
||||
def _subword_tokens(words, start_ms=0, subword_ms=150, inter_word_gap_ms=50):
|
||||
tokens = []
|
||||
t = start_ms
|
||||
for word in words:
|
||||
halves = [word[: len(word) // 2], word[len(word) // 2 :]] if len(word) > 3 else [word]
|
||||
for i, piece in enumerate(halves):
|
||||
text = (" " + piece) if i == 0 else piece
|
||||
tokens.append({"text": text, "start_ms": t, "end_ms": t + subword_ms})
|
||||
t += subword_ms
|
||||
t += inter_word_gap_ms
|
||||
return tokens, t
|
||||
|
||||
|
||||
class TestCueGroupingAlignment:
|
||||
def test_should_split_cue_on_silence_gap_with_exact_timestamps(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
before, t = _subword_tokens(["hello", "there"])
|
||||
after, _ = _subword_tokens(["welcome", "back"], start_ms=t + 5000)
|
||||
result = render_soniox_tokens_as_srt(before + after)
|
||||
cues = result.strip().split("\n\n")
|
||||
assert len(cues) == 2
|
||||
assert "00:00:00,000 --> 00:00:00,650" in cues[0]
|
||||
assert "hello there" in cues[0]
|
||||
assert "00:00:05,700 --> 00:00:06,350" in cues[1]
|
||||
assert "welcome back" in cues[1]
|
||||
|
||||
def test_should_not_bridge_pause_shorter_than_old_duration_cap(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
before, t = _subword_tokens(["first", "part"])
|
||||
after, _ = _subword_tokens(["second", "part"], start_ms=t + 3000)
|
||||
result = render_soniox_tokens_as_srt(before + after)
|
||||
cues = result.strip().split("\n\n")
|
||||
assert len(cues) == 2
|
||||
assert "first part" in cues[0]
|
||||
assert "second part" in cues[1]
|
||||
|
||||
def test_should_never_split_mid_word(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens, _ = _subword_tokens(["hello"] * 20)
|
||||
result = render_soniox_tokens_as_srt(tokens)
|
||||
text_lines = [
|
||||
line for line in result.split("\n") if line and "-->" not in line and not line.isdigit()
|
||||
]
|
||||
assert len(text_lines) >= 2
|
||||
for line in text_lines:
|
||||
assert set(line.split()) == {"hello"}
|
||||
|
||||
def test_should_split_after_sentence_final_punctuation(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens, _ = _subword_tokens(["That", "is", "done.", "Next", "topic"])
|
||||
result = render_soniox_tokens_as_srt(tokens)
|
||||
cues = result.strip().split("\n\n")
|
||||
assert len(cues) == 2
|
||||
assert cues[0].endswith("That is done.")
|
||||
assert cues[1].endswith("Next topic")
|
||||
|
||||
def test_should_split_on_char_budget_at_word_boundary(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens, _ = _subword_tokens(["wonderful"] * 12)
|
||||
result = render_soniox_tokens_as_srt(tokens)
|
||||
text_lines = [
|
||||
line for line in result.split("\n") if line and "-->" not in line and not line.isdigit()
|
||||
]
|
||||
assert len(text_lines) >= 2
|
||||
for line in text_lines:
|
||||
assert len(line) <= 84
|
||||
assert set(line.split()) == {"wonderful"}
|
||||
|
||||
def test_should_exclude_untimestamped_translation_tokens_from_cues(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [
|
||||
{"text": " Good", "start_ms": 0, "end_ms": 200, "translation_status": "original", "language": "en"},
|
||||
{"text": " Guten", "translation_status": "translation", "language": "de", "source_language": "en"},
|
||||
{"text": " morning.", "start_ms": 250, "end_ms": 600, "translation_status": "original", "language": "en"},
|
||||
]
|
||||
result = render_soniox_tokens_as_srt(tokens)
|
||||
assert "Good morning." in result
|
||||
assert "Guten" not in result
|
||||
assert "00:00:00,000 --> 00:00:00,600" in result
|
||||
|
||||
def test_should_split_before_word_whose_end_crosses_duration_cap(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [{"text": " hm", "start_ms": i * 650, "end_ms": i * 650 + 600} for i in range(10)] + [
|
||||
{"text": " boom", "start_ms": 6900, "end_ms": 7600}
|
||||
]
|
||||
result = render_soniox_tokens_as_srt(tokens)
|
||||
cues = result.strip().split("\n\n")
|
||||
assert len(cues) == 2
|
||||
assert "00:00:00,000 --> 00:00:06,450" in cues[0]
|
||||
assert "00:00:06,900 --> 00:00:07,600" in cues[1]
|
||||
assert cues[1].endswith("boom")
|
||||
|
||||
def test_should_keep_untimestamped_word_in_cue(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [
|
||||
{"text": " uh", "start_ms": None, "end_ms": None},
|
||||
{"text": " hello", "start_ms": 100, "end_ms": 500},
|
||||
]
|
||||
result = render_soniox_tokens_as_srt(tokens)
|
||||
assert "uh hello" in result
|
||||
assert "00:00:00,100 --> 00:00:00,500" in result
|
||||
|
||||
|
||||
def _cue_texts(srt: str) -> list:
|
||||
return [cue.split("\n", 2)[2] for cue in srt.strip().split("\n\n")]
|
||||
|
||||
|
||||
class TestMultilingualCueGrouping:
|
||||
def test_should_split_spaceless_chinese_on_width_budget(self):
|
||||
from litellm.litellm_core_utils.audio_utils.subtitle_utils import _text_width
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [{"text": "你好", "start_ms": i * 100, "end_ms": i * 100 + 90} for i in range(60)]
|
||||
result = render_soniox_tokens_as_srt(tokens)
|
||||
texts = _cue_texts(result)
|
||||
assert len(texts) >= 3
|
||||
for text in texts:
|
||||
assert _text_width(text) <= 84
|
||||
assert set(text) <= {"你", "好"}
|
||||
|
||||
def test_should_split_japanese_after_sentence_end_and_keep_punctuation_attached(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [
|
||||
{"text": "今日は", "start_ms": 0, "end_ms": 300},
|
||||
{"text": "いい", "start_ms": 300, "end_ms": 500},
|
||||
{"text": "天気です", "start_ms": 500, "end_ms": 900},
|
||||
{"text": "。", "start_ms": 900, "end_ms": 950},
|
||||
{"text": "明日も", "start_ms": 1000, "end_ms": 1300},
|
||||
{"text": "晴れ", "start_ms": 1300, "end_ms": 1500},
|
||||
]
|
||||
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
|
||||
assert texts == ["今日はいい天気です。", "明日も晴れ"]
|
||||
|
||||
def test_should_split_arabic_after_arabic_question_mark(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [
|
||||
{"text": " كيف", "start_ms": 0, "end_ms": 300},
|
||||
{"text": " حالك؟", "start_ms": 300, "end_ms": 700},
|
||||
{"text": " أنا", "start_ms": 800, "end_ms": 1000},
|
||||
{"text": " بخير", "start_ms": 1000, "end_ms": 1300},
|
||||
]
|
||||
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
|
||||
assert texts == ["كيف حالك؟", "أنا بخير"]
|
||||
|
||||
def test_should_split_after_devanagari_and_urdu_terminators(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [
|
||||
{"text": " नमस्ते।", "start_ms": 0, "end_ms": 400},
|
||||
{"text": " آپ", "start_ms": 500, "end_ms": 700},
|
||||
{"text": " ٹھیک۔", "start_ms": 700, "end_ms": 1100},
|
||||
{"text": " शुभ", "start_ms": 1200, "end_ms": 1400},
|
||||
]
|
||||
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
|
||||
assert texts == ["नमस्ते।", "آپ ٹھیک۔", "शुभ"]
|
||||
|
||||
def test_should_split_russian_after_sentence_end(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [
|
||||
{"text": " Как", "start_ms": 0, "end_ms": 200},
|
||||
{"text": " дела?", "start_ms": 200, "end_ms": 600},
|
||||
{"text": " Хорошо.", "start_ms": 700, "end_ms": 1200},
|
||||
]
|
||||
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
|
||||
assert texts == ["Как дела?", "Хорошо."]
|
||||
|
||||
def test_should_not_split_latin_text_within_width_budget(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
|
||||
|
||||
tokens = [{"text": f" word{i}", "start_ms": i * 100, "end_ms": i * 100 + 90} for i in range(12)]
|
||||
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
|
||||
assert len(texts) == 1
|
||||
|
||||
|
||||
class TestRenderSonioxTokensAsVtt:
|
||||
def test_should_render_basic_vtt_with_header(self):
|
||||
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_vtt
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue