diff --git a/examples/subtitle/generate_subtitle.py b/examples/subtitle/generate_subtitle.py index 1bc865e24..bfa21c6b6 100644 --- a/examples/subtitle/generate_subtitle.py +++ b/examples/subtitle/generate_subtitle.py @@ -14,7 +14,7 @@ import os import re -from funasr.cli import merge_subtitle_segments +from funasr.cli import _sentence_timestamp_words, merge_subtitle_segments def clean_text(text): @@ -112,12 +112,22 @@ def main(): segments = [] result_item = result[0] - for seg in result_item.get("sentence_info", []) or []: + sentence_words = _sentence_timestamp_words(result_item) + for index, seg in enumerate(result_item.get("sentence_info", []) or []): text = clean_text(seg.get("sentence") or seg.get("text", "")) start = int(seg.get("start", 0) or 0) end = int(seg.get("end", 0) or 0) if text and end > start: - segments.append({"start": start, "end": end, "text": text, "spk": seg.get("spk")}) + item = { + "start": start, + "end": end, + "text": text, + "spk": seg.get("spk"), + "timestamp": seg.get("timestamp") or seg.get("timestamps"), + } + if sentence_words[index]: + item["words"] = sentence_words[index] + segments.append(item) if not segments: text = clean_text(result_item.get("text", "")) diff --git a/funasr/cli.py b/funasr/cli.py index d845d88b7..95229315c 100644 --- a/funasr/cli.py +++ b/funasr/cli.py @@ -6,6 +6,7 @@ import re import sys import time +import unicodedata MODEL_CONFIGS = { "sensevoice": {"model": "iic/SenseVoiceSmall", "vad_model": "fsmn-vad", "vad_kwargs": {"max_single_segment_time": 30000}}, @@ -74,6 +75,223 @@ def _join_subtitle_text(left, right): return left + right +def _subtitle_token_spans(text): + spans = [] + pending_start = None + index = 0 + while index < len(text): + char = text[index] + if char.isspace(): + index += 1 + continue + if unicodedata.category(char).startswith("P"): + if spans: + spans[-1][1] = index + 1 + elif pending_start is None: + pending_start = index + index += 1 + continue + + start = index + if char.isascii() and (char.isalnum() or char in "_'"): + index += 1 + while index < len(text): + char = text[index] + if not (char.isascii() and (char.isalnum() or char in "_'")): + break + index += 1 + elif _is_supported_subtitle_character(char): + index += 1 + else: + return [] + if pending_start is not None: + start = pending_start + pending_start = None + spans.append([start, index]) + + if pending_start is not None and spans: + spans[-1][1] = len(text) + return spans + + +def _is_supported_subtitle_character(char): + codepoint = ord(char) + return ( + 0x3400 <= codepoint <= 0x4DBF + or 0x4E00 <= codepoint <= 0x9FFF + or 0xF900 <= codepoint <= 0xFAFF + or 0x20000 <= codepoint <= 0x323AF + or 0x3040 <= codepoint <= 0x30FF + or 0x31F0 <= codepoint <= 0x31FF + or 0xFF66 <= codepoint <= 0xFF9D + or 0x1100 <= codepoint <= 0x11FF + or 0x3130 <= codepoint <= 0x318F + or 0xAC00 <= codepoint <= 0xD7AF + ) + + +def _subtitle_word_spans(text, words): + spans = [] + cursor = 0 + for raw_word in words: + word = str(raw_word).lstrip("▁").strip() + if not word: + return [] + start = text.find(word, cursor) + if start < 0 or any( + not (char.isspace() or unicodedata.category(char).startswith("P")) + for char in text[cursor:start] + ): + return [] + if spans: + spans[-1][1] = start + elif any( + not (char.isspace() or unicodedata.category(char).startswith("P")) + for char in text[:start] + ): + return [] + end = start + len(word) + spans.append([0 if not spans and start else start, end]) + cursor = end + + if any( + not (char.isspace() or unicodedata.category(char).startswith("P")) + for char in text[cursor:] + ): + return [] + if spans: + spans[-1][1] = len(text) + return spans + + +def _timestamp_pair(item): + if not isinstance(item, (list, tuple)) or len(item) < 2: + return None + try: + start = int(item[0]) + end = int(item[1]) + except (TypeError, ValueError, OverflowError): + return None + return [start, end] if end > start else None + + +def _timestamps_are_ordered(timestamps): + return bool(timestamps) and all( + timestamp is not None + and timestamp[0] >= 0 + and (index == 0 or timestamp[0] >= timestamps[index - 1][1]) + for index, timestamp in enumerate(timestamps) + ) + + +def _sentence_timestamp_words(result): + sentence_info = result.get("sentence_info", []) or [] + words = result.get("words", []) or [] + raw_timestamps = result.get("timestamp") or result.get("timestamps") or [] + timestamps = [_timestamp_pair(item) for item in raw_timestamps] + if not words or len(words) != len(timestamps) or not _timestamps_are_ordered( + timestamps + ): + return [None] * len(sentence_info) + + mapped_words = [] + cursor = 0 + for sentence in sentence_info: + local_timestamps = [ + _timestamp_pair(item) + for item in ( + sentence.get("timestamp") or sentence.get("timestamps") or [] + ) + ] + if not _timestamps_are_ordered(local_timestamps): + mapped_words.append(None) + continue + + local_cursor = cursor + selected = [] + for timestamp in local_timestamps: + while ( + local_cursor < len(timestamps) + and timestamps[local_cursor] != timestamp + ): + local_cursor += 1 + if local_cursor == len(timestamps): + selected = [] + break + selected.append(words[local_cursor]) + local_cursor += 1 + if len(selected) == len(local_timestamps): + mapped_words.append(selected) + cursor = local_cursor + else: + mapped_words.append(None) + return mapped_words + + +def _split_subtitle_segment(segment, max_duration_ms, max_chars): + text = str(segment.get("text", "")) + start = int(segment.get("start", 0) or 0) + end = int(segment.get("end", start) or start) + if not text or (end - start <= max_duration_ms and len(text) <= max_chars): + return [dict(segment)] + + raw_timestamps = segment.get("timestamp") or segment.get("timestamps") or [] + timestamps = [_timestamp_pair(item) for item in raw_timestamps] + if not _timestamps_are_ordered(timestamps): + return [dict(segment)] + + words = segment.get("words") or [] + token_spans = ( + _subtitle_word_spans(text, words) if words else _subtitle_token_spans(text) + ) + if not timestamps or len(timestamps) != len(token_spans): + return [dict(segment)] + for index, span in enumerate(token_spans): + token_text = text[span[0] : span[1]].strip() + if ( + timestamps[index][1] - timestamps[index][0] > max_duration_ms + or len(token_text) > max_chars + ): + return [dict(segment)] + + cues = [] + token_start = 0 + while token_start < len(token_spans): + token_end = token_start + while token_end < len(token_spans): + candidate_token_end = token_end + 1 + candidate_text = text[ + token_spans[token_start][0] : token_spans[candidate_token_end - 1][1] + ].strip() + candidate_duration = ( + timestamps[candidate_token_end - 1][1] + - timestamps[token_start][0] + ) + exceeds_limit = ( + candidate_duration > max_duration_ms + or len(candidate_text) > max_chars + ) + if exceeds_limit and token_end > token_start: + break + token_end = candidate_token_end + if exceeds_limit: + break + + cue = dict(segment) + cue["text"] = text[ + token_spans[token_start][0] : token_spans[token_end - 1][1] + ].strip() + cue["start"] = timestamps[token_start][0] + cue["end"] = timestamps[token_end - 1][1] + cue["timestamp"] = timestamps[token_start:token_end] + cue.pop("timestamps", None) + cue.pop("words", None) + cues.append(cue) + token_start = token_end + + return cues + + def merge_subtitle_segments( segments, max_gap_ms=500, max_duration_ms=8000, max_chars=42 ): @@ -107,6 +325,22 @@ def combine(group): for item in group[1:]: text = _join_subtitle_text(text, item.get("text", "")) cue["text"] = text + if any(item.get("timestamp") for item in group): + cue["timestamp"] = [ + timestamp + for item in group + for timestamp in item.get("timestamp", []) + ] + if len(group) > 1 and any(item.get("words") for item in group): + if all( + isinstance(item.get("words"), list) + and item["words"] + and len(item["words"]) == len(item.get("timestamp", [])) + for item in group + ): + cue["words"] = [word for item in group for word in item["words"]] + else: + cue.pop("words", None) return cue def pack(chain): @@ -130,11 +364,13 @@ def pack(chain): merged = [] chain = [] for source in segments: - current = dict(source) - if chain and not can_follow(chain[-1], current): - merged.extend(pack(chain)) - chain = [] - chain.append(current) + for current in _split_subtitle_segment( + source, max_duration_ms=max_duration_ms, max_chars=max_chars + ): + if chain and not can_follow(chain[-1], current): + merged.extend(pack(chain)) + chain = [] + chain.append(current) if chain: merged.extend(pack(chain)) return merged @@ -313,8 +549,20 @@ def main(): text = clean_text(result[0].get("text", "")) segments = [] if "sentence_info" in result[0]: - for seg in result[0]["sentence_info"]: - s = {"start": seg.get("start", 0), "end": seg.get("end", 0), "text": clean_text(seg.get("sentence") or seg.get("text", ""))} + sentence_words = _sentence_timestamp_words(result[0]) + for index, seg in enumerate(result[0]["sentence_info"]): + s = { + "start": seg.get("start", 0), + "end": seg.get("end", 0), + "text": clean_text(seg.get("sentence") or seg.get("text", "")), + "timestamp": seg.get("timestamp") or seg.get("timestamps"), + } + if ( + args.output_format == "srt" + and args.subtitle_segment_mode == "readable" + and sentence_words[index] + ): + s["words"] = sentence_words[index] if args.spk and "spk" in seg: s["speaker"] = seg["spk"] segments.append(s) diff --git a/tests/test_cli.py b/tests/test_cli.py index dafbd25cf..ef7aafa38 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -299,3 +299,250 @@ def test_merge_subtitle_segments_preserves_hard_boundaries(): ] assert cli.merge_subtitle_segments(segments) == segments + + +def test_merge_subtitle_segments_splits_overlong_source_with_token_timestamps(): + text = "甲" * 60 + timestamps = [[index * 300, index * 300 + 120] for index in range(len(text))] + segments = [ + { + "start": timestamps[0][0], + "end": timestamps[-1][1], + "text": text, + "timestamp": timestamps, + } + ] + + cues = cli.merge_subtitle_segments(segments) + + assert len(cues) > 1 + assert "".join(cue["text"] for cue in cues) == text + assert [timestamp for cue in cues for timestamp in cue["timestamp"]] == timestamps + assert all(cue["end"] - cue["start"] <= 8000 for cue in cues) + assert all(len(cue["text"]) <= 42 for cue in cues) + + +def test_merge_subtitle_segments_keeps_word_timestamp_boundaries(): + timestamps = [[0, 900], [1000, 1900], [2000, 2900]] + segments = [ + { + "start": 0, + "end": 2900, + "text": "hello, world! again", + "timestamp": timestamps, + } + ] + + assert cli.merge_subtitle_segments( + segments, max_duration_ms=2000, max_chars=13 + ) == [ + { + "start": 0, + "end": 1900, + "text": "hello, world!", + "timestamp": timestamps[:2], + }, + { + "start": 2000, + "end": 2900, + "text": "again", + "timestamp": timestamps[2:], + }, + ] + + +def test_merge_subtitle_segments_preserves_overlong_source_without_timestamps(): + segment = {"start": 0, "end": 12000, "text": "没有时间戳的长句" * 8} + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_preserves_unalignable_token_timestamps(): + segment = { + "start": 0, + "end": 12000, + "text": "어디 통화하고 있는 거네 지금", + "timestamp": [[index * 1000, index * 1000 + 900] for index in range(8)], + } + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_keeps_indivisible_overlong_token(): + segment = { + "start": 0, + "end": 12000, + "text": "supercalifragilisticexpialidocious", + "timestamps": [[0, 12000]], + "words": ["supercalifragilisticexpialidocious"], + } + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_preserves_source_around_indivisible_token(): + segment = { + "start": 0, + "end": 14000, + "text": "before supercalifragilisticexpialidocious after", + "timestamps": [[0, 900], [1000, 13000], [13100, 14000]], + "words": ["before", "supercalifragilisticexpialidocious", "after"], + } + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_rejects_partially_invalid_timestamps(): + segment = { + "start": 0, + "end": 12000, + "text": "hello world", + "timestamp": [[0, 900], [1000], [11100, 12000]], + } + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_rejects_out_of_order_timestamps(): + segment = { + "start": 0, + "end": 12000, + "text": "hello world again", + "timestamp": [[0, 1000], [5000, 6000], [2000, 3000]], + } + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_rejects_negative_timestamps(): + segment = { + "start": 0, + "end": 12000, + "text": "hello world", + "timestamp": [[-100, 3900], [4000, 7900]], + } + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_rejects_non_finite_timestamps(): + segment = { + "start": 0, + "end": 12000, + "text": "hello", + "timestamp": [[0, float("inf")]], + } + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_does_not_infer_unsupported_script_surfaces(): + segment = { + "start": 0, + "end": 12000, + "text": "مرح", + "timestamp": [[0, 3900], [4000, 7900], [8000, 12000]], + } + + assert cli.merge_subtitle_segments([segment]) == [segment] + + +def test_merge_subtitle_segments_uses_explicit_word_surfaces(): + timestamps = [[index * 1000, index * 1000 + 900] for index in range(8)] + segment = { + "start": 0, + "end": 7900, + "text": "어디 통화하고 있는 거네 지금", + "timestamp": timestamps, + "words": ["어", "디", "통화", "하고", "있는", "거", "네", "지금"], + } + + assert cli.merge_subtitle_segments( + [segment], max_duration_ms=3900, max_chars=42 + ) == [ + { + "start": 0, + "end": 3900, + "text": "어디 통화하고", + "timestamp": timestamps[:4], + }, + { + "start": 4000, + "end": 7900, + "text": "있는 거네 지금", + "timestamp": timestamps[4:], + }, + ] + + +def test_sentence_timestamp_words_tracks_global_word_surfaces(): + result = { + "words": ["hello", "世", "界", "▁진짜"], + "timestamp": [[0, 100], [110, 210], [220, 320], [330, 430]], + "sentence_info": [ + {"timestamp": [[0, 100], [110, 210], [220, 320]]}, + {"timestamp": [[330, 430]]}, + ], + } + + assert cli._sentence_timestamp_words(result) == [ + ["hello", "世", "界"], + ["▁진짜"], + ] + + +def test_merge_subtitle_segments_combines_explicit_word_surfaces(): + segments = [ + { + "start": 0, + "end": 900, + "text": "hello,", + "timestamp": [[0, 900]], + "words": ["hello"], + }, + { + "start": 1000, + "end": 1900, + "text": "world", + "timestamp": [[1000, 1900]], + "words": ["world"], + }, + ] + + assert cli.merge_subtitle_segments(segments) == [ + { + "start": 0, + "end": 1900, + "text": "hello,world", + "timestamp": [[0, 900], [1000, 1900]], + "words": ["hello", "world"], + } + ] + + +def test_merge_subtitle_segments_drops_incomplete_word_surfaces(): + segments = [ + { + "start": 0, + "end": 900, + "text": "hello,", + "timestamp": [[0, 900]], + "words": ["hello"], + }, + { + "start": 1000, + "end": 1900, + "text": "world", + "timestamp": [[1000, 1900]], + "words": [], + }, + ] + + assert cli.merge_subtitle_segments(segments) == [ + { + "start": 0, + "end": 1900, + "text": "hello,world", + "timestamp": [[0, 900], [1000, 1900]], + } + ]