Fix CLI SRT sentence segmentation (#3254)

Author: LauraGPT

funasr/cli.py

@@ -20,9 +20,12 @@ def clean_text(text):
 
 
 def _srt_time(ms):
-    s = ms / 1000.0
-    h, m, sec = int(s // 3600), int((s % 3600) // 60), int(s % 60)
-    return f"{h:02d}:{m:02d}:{sec:02d},{int((s % 1) * 1000):03d}"
+    ms = max(0, int(round(ms)))
+    h = ms // 3600000
+    m = (ms % 3600000) // 60000
+    sec = (ms % 60000) // 1000
+    ms_rem = ms % 1000
+    return f"{h:02d}:{m:02d}:{sec:02d},{ms_rem:03d}"
 
 
 def format_srt(segments):
@@ -39,6 +42,38 @@ def format_tsv(segments):
     return "\n".join(lines)
 
 
+def _parse_ms(value, scale=1):
+    if value is None:
+        return None
+    try:
+        return int(float(value) * scale)
+    except (TypeError, ValueError):
+        return None
+
+
+def _timestamp_bounds_ms(result):
+    bounds = []
+    for key in ("timestamp", "timestamps"):
+        for ts in result.get(key, []) or []:
+            if isinstance(ts, dict):
+                start = ts.get("start_time", ts.get("start"))
+                end = ts.get("end_time", ts.get("end"))
+                start_ms = _parse_ms(start, 1000)
+                end_ms = _parse_ms(end, 1000)
+            elif isinstance(ts, (list, tuple)) and len(ts) >= 2:
+                start_ms = _parse_ms(ts[0])
+                end_ms = _parse_ms(ts[1])
+            else:
+                continue
+            if start_ms is None or end_ms is None:
+                continue
+            if end_ms > start_ms:
+                bounds.append((start_ms, end_ms))
+    if not bounds:
+        return None
+    return min(start for start, _ in bounds), max(end for _, end in bounds)
+
+
 def _format_output(text, segments, timestamps, fmt, audio_path, model_name, language, elapsed):
     if fmt == "text":
         return text
@@ -58,8 +93,12 @@ def _format_output(text, segments, timestamps, fmt, audio_path, model_name, lang
     elif fmt == "srt":
         if segments:
             return format_srt(segments)
-        # No per-sentence timestamps: emit one valid cue spanning the whole audio
-        # (instead of a bogus 99:59:59 end time).
+        # No per-sentence timestamps: emit one valid cue spanning the known
+        # timestamp/audio bounds instead of a bogus 99:59:59 end time.
+        timestamp_bounds = _timestamp_bounds_ms({"timestamp": timestamps})
+        if timestamp_bounds:
+            start_ms, end_ms = timestamp_bounds
+            return f"1\n{_srt_time(start_ms)} --> {_srt_time(end_ms)}\n{text}\n"
         try:
             import soundfile as sf
             dur_ms = int(sf.info(audio_path).duration * 1000)
@@ -115,8 +154,9 @@ def main():
     config["hub"] = args.hub
     if args.spk and "spk_model" not in config:
         config["spk_model"] = "cam++"
-    if "punc_model" not in config and args.model not in ("fun-asr-nano", "sensevoice"):
-        config["punc_model"] = "ct-punc"
+    if "punc_model" not in config and args.model != "fun-asr-nano":
+        if args.model != "sensevoice" or args.output_format in ("srt", "tsv"):
+            config["punc_model"] = "ct-punc"
 
     t_load = time.time()
     model = AutoModel(device=device, disable_update=True, **config)
@@ -147,6 +187,15 @@ def main():
             else:
                 gen_kw["hotwords"] = hotwords
 
+        if args.output_format in ("srt", "tsv"):
+            gen_kw.update(
+                {
+                    "sentence_timestamp": True,
+                    "output_timestamp": True,
+                    "return_time_stamps": True,
+                }
+            )
+
         result = model.generate(**gen_kw)
         elapsed = time.time() - t0
 
@@ -159,7 +208,9 @@ def main():
                     s["speaker"] = seg["spk"]
                 segments.append(s)
 
-        timestamps = result[0].get("timestamps") if args.timestamps else None
+        timestamps = result[0].get("timestamps") or result[0].get("timestamp")
+        if not args.timestamps and args.output_format not in ("srt", "tsv"):
+            timestamps = None
         output = _format_output(text, segments, timestamps, args.output_format, audio_path, args.model, args.language, elapsed)
 
         if args.output_dir:

tests/test_cli.py

@@ -19,6 +19,22 @@ def generate(self, **kwargs):
         return [{"text": "hello"}]
 
 
+class SubtitleAutoModel(DummyAutoModel):
+    def generate(self, **kwargs):
+        self.generate_kwargs = kwargs
+        if kwargs.get("sentence_timestamp") and kwargs.get("output_timestamp"):
+            return [
+                {
+                    "text": "<|zh|>第一句。第二句。",
+                    "sentence_info": [
+                        {"start": 0, "end": 1200, "text": "<|zh|>第一句。"},
+                        {"start": 1200, "end": 2600, "sentence": "第二句。"},
+                    ],
+                }
+            ]
+        return [{"text": "<|zh|>第一句。第二句。"}]
+
+
 def test_cli_passes_hub_to_auto_model(tmp_path):
     audio_path = tmp_path / "sample.wav"
     audio_path.write_bytes(b"not a real wav")
@@ -69,3 +85,61 @@ def test_cli_routes_multiple_hotwords_to_paraformer_hotword(tmp_path):
     generate_kwargs = DummyAutoModel.instances[0].generate_kwargs
     assert generate_kwargs["hotword"] == "FunASR ModelScope"
     assert "hotwords" not in generate_kwargs
+
+
+def test_timestamp_bounds_skip_malformed_entries():
+    assert cli._timestamp_bounds_ms(
+        {
+            "timestamp": [
+                [None, 1000],
+                ["bad", "2000"],
+                ["1200.0", "2600.0"],
+                {"start": "0.5", "end": "1.0"},
+                {"start_time": "oops", "end_time": "2.0"},
+            ]
+        }
+    ) == (500, 2600)
+
+
+def test_cli_srt_requests_sentence_timestamps_and_writes_segmented_output(tmp_path):
+    audio_path = tmp_path / "sample.wav"
+    audio_path.write_bytes(b"not a real wav")
+    out_dir = tmp_path / "subs"
+    fake_torch = types.SimpleNamespace(
+        cuda=types.SimpleNamespace(is_available=lambda: False),
+    )
+
+    SubtitleAutoModel.instances = []
+    argv = [
+        "funasr",
+        str(audio_path),
+        "--output-format",
+        "srt",
+        "--output-dir",
+        str(out_dir),
+        "--lang",
+        "zh",
+    ]
+
+    with (
+        patch.object(sys, "argv", argv),
+        patch.dict(sys.modules, {"torch": fake_torch}),
+        patch("funasr.AutoModel", SubtitleAutoModel),
+        redirect_stdout(io.StringIO()),
+    ):
+        cli.main()
+
+    instance = SubtitleAutoModel.instances[0]
+    assert instance.kwargs["punc_model"] == "ct-punc"
+    assert instance.generate_kwargs["language"] == "zh"
+    assert instance.generate_kwargs["sentence_timestamp"] is True
+    assert instance.generate_kwargs["output_timestamp"] is True
+    assert instance.generate_kwargs["return_time_stamps"] is True
+    assert (out_dir / "sample.srt").read_text(encoding="utf-8") == (
+        "1\n"
+        "00:00:00,000 --> 00:00:01,200\n"
+        "第一句。\n\n"
+        "2\n"
+        "00:00:01,200 --> 00:00:02,600\n"
+        "第二句。\n"
+    )