fix(subtitles): split overlong source cues

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", ""))

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)