[Feat] Streaming TTS word-level timestamps via forced aligner (#4034)

Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

tests/utils/test_forced_aligner.py

@@ -0,0 +1,165 @@
+import numpy as np
+import pytest
+
+from vllm_omni.utils import forced_aligner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_decode_timestamps_maps_boundary_bins_to_words():
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0
+    logits[1, 2] = 1.0
+    logits[2, 2] = 1.0
+    logits[3, 4] = 1.0
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == [
+        forced_aligner.WordTimestamp("hello", 0, 400),
+        forced_aligner.WordTimestamp("world", 400, 800),
+    ]
+
+
+def test_decode_timestamps_repairs_non_monotonic_bins():
+    # Word 1's end bin (1) dips below its start bin (2); fix_timestamp should
+    # snap it back so the frame never carries end_ms < start_ms.
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0  # word0 start -> bin 0
+    logits[1, 2] = 1.0  # word0 end   -> bin 2
+    logits[2, 1] = 1.0  # word1 start -> bin 1 (out of order)
+    logits[3, 3] = 1.0  # word1 end   -> bin 3
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        timestamp_segment_time_ms=200,
+        audio_duration_ms=1000,
+    )
+
+    for ts in timestamps:
+        assert ts.end_ms >= ts.start_ms
+    # bins [0, 2, 1, 3] -> repaired [0, 2, 2, 3] -> ms x200
+    assert timestamps == [
+        forced_aligner.WordTimestamp("hello", 0, 400),
+        forced_aligner.WordTimestamp("world", 400, 600),
+    ]
+
+
+def test_decode_timestamps_clamps_to_audio_duration_and_stays_ordered():
+    # Bins run past the 1000 ms audio; the decoder must cap each bound at the
+    # audio length and keep words non-overlapping (start <= end, monotonic).
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0  # word0 start -> bin 0
+    logits[1, 3] = 1.0  # word0 end   -> bin 3 (1200 ms, past audio)
+    logits[2, 3] = 1.0  # word1 start -> bin 3 (1200 ms, past audio)
+    logits[3, 4] = 1.0  # word1 end   -> bin 4 (1600 ms, past audio)
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["a", "b"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        timestamp_segment_time_ms=400,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == [
+        forced_aligner.WordTimestamp("a", 0, 1000),
+        forced_aligner.WordTimestamp("b", 1000, 1000),
+    ]
+
+
+def test_decode_timestamps_rejects_marker_count_mismatch():
+    logits = np.zeros((2, 5), dtype=np.float32)
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1],
+        classify_num=5,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == []
+
+
+def test_build_config_from_yaml(tmp_path):
+    cfg = tmp_path / "forced_aligner.yaml"
+    cfg.write_text(
+        """
+forced_aligner:
+  model: Qwen/Qwen3-ForcedAligner-0.6B
+  gpu_memory_utilization: 0.42
+  dtype: float16
+  max_model_len: 2048
+  trust_remote_code: false
+""",
+        encoding="utf-8",
+    )
+    args = type("Args", (), {"forced_aligner": None, "forced_aligner_config": str(cfg)})()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out == forced_aligner.ForcedAlignerConfig(
+        model="Qwen/Qwen3-ForcedAligner-0.6B",
+        runner="pooling",
+        architecture="Qwen3ASRForcedAlignerForTokenClassification",
+        pooling_task="token_classify",
+        gpu_memory_utilization=0.42,
+        dtype="float16",
+        max_model_len=2048,
+        trust_remote_code=False,
+    )
+
+
+def test_build_config_cli_model_overrides_yaml(tmp_path):
+    cfg = tmp_path / "forced_aligner.yaml"
+    cfg.write_text(
+        "forced_aligner:\n  model: old\n  gpu_memory_utilization: 0.2\n  dtype: float16\n",
+        encoding="utf-8",
+    )
+    args = type(
+        "Args",
+        (),
+        {
+            "forced_aligner": "new",
+            "forced_aligner_config": str(cfg),
+        },
+    )()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out is not None
+    # --forced-aligner overrides the YAML model; gpu_memory_utilization/dtype
+    # come from the user YAML (there is no longer a CLI flag for gpu mem).
+    assert out.model == "new"
+    assert out.gpu_memory_utilization == 0.2
+    assert out.dtype == "float16"
+    assert out.runner == "pooling"
+
+
+def test_build_config_from_cli_model_uses_default_yaml():
+    args = type(
+        "Args",
+        (),
+        {
+            "forced_aligner": "local-aligner",
+            "forced_aligner_config": None,
+        },
+    )()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out is not None
+    assert out.model == "local-aligner"
+    assert out.runner == "pooling"
+    assert out.architecture == "Qwen3ASRForcedAlignerForTokenClassification"

tests/utils/test_qwen3_force_align_processor.py

@@ -0,0 +1,65 @@
+import pytest
+
+from vllm_omni.utils import qwen3_force_align_processor as processor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_build_prompt_has_boundary_timestamp_markers():
+    prompt = processor.build_prompt(["hello", "world"])
+
+    assert prompt.count("<timestamp>") == 4
+    assert "hello<timestamp><timestamp>world" in prompt
+    assert prompt.endswith("<|im_start|>assistant\n")
+
+
+@pytest.mark.parametrize(
+    "text,expected",
+    [
+        # Punctuation inside a whitespace segment is stripped, not split on:
+        # the hand-rolled splitter used to break these into many fake words.
+        ("U.S.A", ["USA"]),
+        ("hello, world!", ["hello", "world"]),
+        ("don't stop", ["don't", "stop"]),
+        # CJK characters peel off as individual tokens, Latin runs stay whole.
+        ("你好world", ["你", "好", "world"]),
+        ("我爱 China", ["我", "爱", "China"]),
+        ("   spaced   out   ", ["spaced", "out"]),
+    ],
+)
+def test_tokenize_space_lang_matches_official_segmentation(text, expected):
+    assert processor._tokenize_space_lang(text) == expected
+
+
+def test_segment_words_falls_back_to_port_without_qwen_asr(monkeypatch):
+    # With qwen_asr unavailable, segmentation must use the built-in port and
+    # still produce the faithful result for the common (non-JP/KO) case.
+    monkeypatch.setattr(processor, "_get_official_processor", lambda: None)
+
+    assert processor.segment_words("U.S.A is here", "auto") == ["USA", "is", "here"]
+
+
+def test_fix_timestamp_repairs_dip_onto_monotonic_sequence():
+    # 200 dips below the preceding 400; it snaps back up to 400.
+    assert processor.fix_timestamp([0, 400, 200, 800]) == [0, 400, 400, 800]
+
+
+def test_fix_timestamp_passes_through_monotonic():
+    assert processor.fix_timestamp([0, 100, 100, 250]) == [0, 100, 100, 250]
+
+
+def test_resolve_timestamp_token_id_defaults_to_marker_token():
+    # Regression: the default must resolve the same <timestamp> marker that
+    # build_prompt inserts, not None (which would degrade every request).
+    seen = {}
+
+    class FakeTokenizer:
+        def convert_tokens_to_ids(self, token):
+            seen["token"] = token
+            return 151705
+
+    tid = processor.resolve_timestamp_token_id(FakeTokenizer())
+
+    assert tid == 151705
+    assert seen["token"] == processor.TIMESTAMP_TOKEN
+    assert processor.TIMESTAMP_TOKEN in processor.build_prompt(["hello", "world"])