fix: align repository SenseVoice timestamps with FunASR

Signed-off-by: LauraGPT <18321252+LauraGPT@users.noreply.github.com>

.github/workflows/sensevoice-container.yml

@@ -18,6 +18,8 @@ on:
       - tests/test_container_contract.py
       - tests/test_docker_context.sh
       - tests/test_device_env.py
+      - tests/test_model_timestamps.py
+      - utils/ctc_alignment.py
   push:
     branches:
       - main
@@ -39,6 +41,8 @@ on:
       - tests/test_container_contract.py
       - tests/test_docker_context.sh
       - tests/test_device_env.py
+      - tests/test_model_timestamps.py
+      - utils/ctc_alignment.py
   workflow_dispatch:
 
 permissions:
@@ -58,6 +62,13 @@ jobs:
           persist-credentials: false
       - name: Validate container contract
         run: python -m unittest tests.test_container_contract tests.test_device_env
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install CPU PyTorch for model contracts
+        run: python -m pip install torch==2.12.1 --index-url https://download.pytorch.org/whl/cpu
+      - name: Validate module and timestamp contracts
+        run: python -m unittest tests.test_model_timestamps -v
 
   build:
     needs: contract

README.md

@@ -211,7 +211,17 @@ for sent in res[0]["sentence_info"]:
     print(f"Speaker {sent['spk']}: [{sent['start']}ms - {sent['end']}ms] {text}")
 ```
 
-> Note: Requires installing FunASR from source: `pip install git+https://github.com/modelscope/FunASR.git`
+Use the current repository `model.py` with FunASR 1.4.15 or newer. A local source
+update is needed when using `remote_code="./model.py"`; upgrading the Python
+package alone does not update that file. This composition was tested on a
+fixed public sample, not validated for diarization accuracy.
+
+With `output_timestamp=True`, `timestamp` contains `[start_ms, end_ms]` pairs
+aligned one-to-one with `words`. Older copies of this repository returned
+`[token, start_seconds, end_seconds]` triples, which are incompatible with
+FunASR's VAD timestamp aggregation. Direct callers must now read the token or
+word from `words`, not from `timestamp[i][0]`. See [demo2.py](./demo2.py).
+Speaker labels are anonymous clusters, not recognized personal identities.
 
 If all inputs are short audios (<30s), and batch inference is needed to speed up inference efficiency, the VAD model can be removed, and `batch_size` can be set accordingly.
 ```python

README_zh.md

@@ -116,6 +116,13 @@ SenseVoiceSmall 示例与 FunASR 组合说话人分离路径需要 `funasr>=1.3.
 
 支持常见格式音频输入。长录音必须先分段再送入编码器；下例使用 FSMN-VAD 完成分段。
 
+使用 `remote_code="./model.py"` 时，升级 FunASR 包不会同步更新本地的
+`model.py`，请同时更新仓库源码。当前时间戳格式为 `timestamp=[[开始毫秒, 结束毫秒], ...]`，
+与 `words` 一一对应；旧版的 `[词元, 开始秒, 结束秒]` 三元组不兼容 VAD 聚合。
+直接调用模型的代码也需改为从 `words` 读取文字，参见 [demo2.py](./demo2.py)。
+说话人组合示例见 [Speaker Diarization](./README.md#speaker-diarization)，已用
+FunASR 1.4.15 和固定公开样本验证流程；这不是说话人准确率或真实身份识别验证。
+
 ```python
 from funasr import AutoModel
 from funasr.utils.postprocess_utils import rich_transcription_postprocess

demo2.py

@@ -32,6 +32,10 @@ res = m.inference(
 )
 
 timestamp = res[0][0]["timestamp"]
+words = res[0][0]["words"]
 text = rich_transcription_postprocess(res[0][0]["text"])
 print(text)
-print(timestamp)
+if len(words) != len(timestamp):
+    raise ValueError("Words and timestamps must have equal lengths")
+for word, (start_ms, end_ms) in zip(words, timestamp):
+    print(f"{start_ms}-{end_ms} ms: {word}")

model.py

@@ -19,7 +19,7 @@ class SinusoidalPositionEncoder(torch.nn.Module):
     """ """
 
     def __init__(self, d_model=80, dropout_rate=0.1):
-        pass
+        super().__init__()
 
     def encode(
         self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32
@@ -887,40 +887,84 @@ class SenseVoiceSmall(nn.Module):
                 from itertools import groupby
                 timestamp = []
                 tokens = tokenizer.text2tokens(text)[4:]
+                token_ids = []
+                token_owners = []
+                for owner, ids in enumerate(tokenizer.tokens2ids(tokens)):
+                    ids = ids or [124]
+                    token_ids.extend(ids)
+                    token_owners.extend([owner] * len(ids))
+                if not token_ids:
+                    results.append({"key": key[i], "text": text, "timestamp": [], "words": []})
+                    continue
 
-                logits_speech = self.ctc.softmax(encoder_out)[i, 4:encoder_out_lens[i].item(), :]
+                logits_speech = self.ctc.log_softmax(encoder_out)[i, 4:encoder_out_lens[i].item(), :]
 
                 pred = logits_speech.argmax(-1).cpu()
                 logits_speech[pred==self.blank_id, self.blank_id] = 0
 
                 align = ctc_forced_align(
                     logits_speech.unsqueeze(0).float(),
-                    torch.Tensor(token_int[4:]).unsqueeze(0).long().to(logits_speech.device),
+                    torch.tensor(token_ids).unsqueeze(0).long().to(logits_speech.device),
                     (encoder_out_lens-4).long()[i],
-                    torch.tensor(len(token_int)-4).unsqueeze(0).long().to(logits_speech.device),
+                    torch.tensor(len(token_ids)).unsqueeze(0).long().to(logits_speech.device),
                     ignore_id=self.ignore_id,
                 )
 
-                pred = groupby(align[0, :encoder_out_lens[0]])
+                pred = groupby(align[0, :encoder_out_lens[i]])
                 _start = 0
                 token_id = 0
                 ts_max = encoder_out_lens[i] - 4
                 for pred_token, pred_frame in pred:
                     _end = _start + len(list(pred_frame))
-                    if pred_token != 0 and token_id < len(tokens):
+                    if pred_token != 0:
+                        if token_id >= len(token_ids) or int(pred_token) != token_ids[token_id]:
+                            raise ValueError("CTC alignment does not match the decoded token sequence")
+                        owner = token_owners[token_id]
                         ts_left = max((_start*60-30)/1000, 0)
                         ts_right = min((_end*60-30)/1000, (ts_max*60-30)/1000)
-                        timestamp.append([tokens[token_id], ts_left, ts_right])
+                        if token_id and owner == token_owners[token_id - 1]:
+                            timestamp[-1][2] = ts_right
+                        else:
+                            timestamp.append([tokens[owner], ts_left, ts_right])
                         token_id += 1
                     _start = _end
+                if token_id != len(token_ids):
+                    raise ValueError("CTC alignment did not cover all decoded tokens")
 
-                result_i = {"key": key[i], "text": text, "timestamp": timestamp}
+                timestamp, words = self.post(timestamp)
+                result_i = {"key": key[i], "text": text, "timestamp": timestamp, "words": words}
                 results.append(result_i)
             else:
                 result_i = {"key": key[i], "text": text}
                 results.append(result_i)
         return results, meta_data
 
+    def post(self, timestamp):
+        """Convert aligned subword seconds to word-aligned millisecond pairs."""
+        pairs, words = [], []
+        previous = None
+        for token, start, end in timestamp:
+            if token == "▁":
+                previous = None
+                continue
+            starts_word = token.startswith("▁")
+            word = token[1:] if starts_word else token
+            start, end = int(start * 1000), int(end * 1000)
+            if (
+                not starts_word
+                and previous is not None
+                and previous.isascii() and previous.isalpha()
+                and word.isascii() and word.isalpha()
+            ):
+                word = previous + word
+                pairs[-1][1] = end
+                words[-1] = word
+            else:
+                pairs.append([start, end])
+                words.append(word)
+            previous = word
+        return pairs, words
+
     def export(self, **kwargs):
         from export_meta import export_rebuild_model
 

tests/test_model_timestamps.py

@@ -0,0 +1,168 @@
+"""Exercise production helpers without downloading checkpoint dependencies."""
+import ast
+from pathlib import Path
+import unittest
+from types import SimpleNamespace
+
+import torch
+from utils.ctc_alignment import ctc_forced_align
+
+
+SOURCE = Path(__file__).resolve().parents[1] / "model.py"
+
+
+def load_class(name):
+    node = next(
+        node for node in ast.parse(SOURCE.read_text()).body
+        if isinstance(node, ast.ClassDef) and node.name == name
+    )
+    namespace = {"torch": torch}
+    exec(compile(ast.Module(body=[node], type_ignores=[]), str(SOURCE), "exec"), namespace)
+    return namespace[name]
+
+
+def load_model_methods():
+    node = next(
+        node for node in ast.parse(SOURCE.read_text()).body
+        if isinstance(node, ast.ClassDef) and node.name == "SenseVoiceSmall"
+    )
+    methods = [item for item in node.body if isinstance(item, ast.FunctionDef)
+               and item.name in ("inference", "post")]
+    namespace = {
+        "torch": torch,
+        "ctc_forced_align": ctc_forced_align,
+    }
+    exec(compile(ast.Module(body=methods, type_ignores=[]), str(SOURCE), "exec"), namespace)
+    return namespace
+
+
+class PositionEncoderTest(unittest.TestCase):
+    def test_module_supports_checkpoint_state_and_device_transfer(self):
+        encoder = load_class("SinusoidalPositionEncoder")(d_model=8)
+        self.assertEqual(dict(encoder.state_dict()), {})
+        encoder.load_state_dict({})
+        encoder.to("cpu")
+
+    def test_call_matches_forward_without_resizing_the_input(self):
+        encoder = load_class("SinusoidalPositionEncoder")(d_model=8)
+        speech = torch.zeros(2, 5, 8)
+        output = encoder(speech)
+        self.assertEqual(output.shape, speech.shape)
+        self.assertTrue(torch.equal(output, encoder.forward(speech)))
+        self.assertTrue(torch.equal(speech, torch.zeros_like(speech)))
+
+
+class TimestampContractTest(unittest.TestCase):
+    def test_multiple_ids_keep_their_original_token_time_range(self):
+        methods = load_model_methods()
+        logits = torch.full((1, 9, 8), -10.0)
+        for frame, token in enumerate([1, 2, 3, 4, 5, 6, 0, 7, 0]):
+            logits[0, frame, token] = 10
+        model = SimpleNamespace(
+            lid_dict={"zh": 0}, textnorm_dict={"woitn": 0},
+            embed=lambda ids: torch.zeros(*ids.shape, 8),
+            encoder=lambda speech, lengths: (torch.zeros(1, 9, 8), torch.tensor([9])),
+            ctc=SimpleNamespace(log_softmax=lambda x: logits, softmax=lambda x: logits.softmax(-1)),
+            blank_id=0, ignore_id=-1,
+        )
+        model.post = lambda timestamps: methods["post"](model, timestamps)
+        tokenizer = SimpleNamespace(
+            decode=lambda ids: "A B",
+            text2tokens=lambda text: ["zh", "neutral", "speech", "woitn", "▁A", "▁B"],
+            tokens2ids=lambda tokens: [[5, 6], [7]],
+        )
+        results, _ = methods["inference"](
+            model, torch.zeros(1, 5, 8), torch.tensor([5]), key=["sample"],
+            tokenizer=tokenizer, data_type="fbank", device="cpu",
+            language="zh", output_timestamp=True,
+        )
+        self.assertEqual(results[0]["timestamp"], [[0, 90], [150, 210]])
+        self.assertEqual(results[0]["words"], ["A", "B"])
+
+        tokenizer.decode = lambda ids: ""
+        tokenizer.text2tokens = lambda text: ["zh", "neutral", "speech", "woitn"]
+        tokenizer.tokens2ids = lambda tokens: []
+        results, _ = methods["inference"](
+            model, torch.zeros(1, 5, 8), torch.tensor([5]), key=["empty"],
+            tokenizer=tokenizer, data_type="fbank", device="cpu",
+            language="zh", output_timestamp=True,
+        )
+        self.assertEqual(results, [{"key": "empty", "text": "", "timestamp": [], "words": []}])
+
+    def test_batch_alignment_uses_each_items_length(self):
+        methods = load_model_methods()
+        logits = torch.full((2, 12, 6), -10.0)
+        for item in range(2):
+            for frame, token in enumerate([1, 2, 3, 4, 5, 0, 5, 0, 5, 0, 5, 0]):
+                logits[item, frame, token] = 10
+        model = SimpleNamespace(
+            lid_dict={"zh": 0}, textnorm_dict={"woitn": 0},
+            embed=lambda ids: torch.zeros(*ids.shape, 8),
+            encoder=lambda speech, lengths: (torch.zeros(2, 12, 8), torch.tensor([6, 12])),
+            ctc=SimpleNamespace(log_softmax=lambda x: logits, softmax=lambda x: logits.softmax(-1)),
+            blank_id=0, ignore_id=-1,
+        )
+        model.post = lambda timestamps: methods["post"](model, timestamps)
+        tokenizer = SimpleNamespace(
+            decode=lambda ids: "你" * (len(ids) - 4),
+            text2tokens=lambda text: ["zh", "neutral", "speech", "woitn"] + list(text),
+            tokens2ids=lambda tokens: [[5] for token in tokens],
+        )
+        results, _ = methods["inference"](
+            model, torch.zeros(2, 8, 8), torch.tensor([2, 8]), key=["short", "long"],
+            tokenizer=tokenizer, data_type="fbank", device="cpu",
+            language="zh", output_timestamp=True,
+        )
+        self.assertEqual([r["key"] for r in results], ["short", "long"])
+        self.assertEqual(results[0]["timestamp"], [[0, 30]])
+        self.assertEqual(results[1]["timestamp"], [[0, 30], [90, 150], [210, 270], [330, 390]])
+        self.assertEqual(results[1]["words"], ["你"] * 4)
+
+    def test_inference_returns_word_aligned_millisecond_pairs(self):
+        methods = load_model_methods()
+        logits = torch.full((1, 8, 6), -10.0)
+        for frame, token in enumerate([1, 2, 3, 4, 5, 0, 5, 0]):
+            logits[0, frame, token] = 10
+        model = SimpleNamespace(
+            lid_dict={"zh": 0}, textnorm_dict={"woitn": 0},
+            embed=lambda ids: torch.zeros(*ids.shape, 8),
+            encoder=lambda speech, lengths: (torch.zeros(1, 8, 8), torch.tensor([8])),
+            ctc=SimpleNamespace(log_softmax=lambda x: logits, softmax=lambda x: logits.softmax(-1)),
+            blank_id=0, ignore_id=-1,
+        )
+        if "post" in methods:
+            model.post = lambda timestamps: methods["post"](model, timestamps)
+        tokenizer = SimpleNamespace(
+            decode=lambda ids: "你好",
+            text2tokens=lambda text: ["zh", "neutral", "speech", "woitn", "你", "好"],
+            tokens2ids=lambda tokens: [[5] for token in tokens],
+        )
+        results, _ = methods["inference"](
+            model, torch.zeros(1, 4, 8), torch.tensor([4]), key=["sample"],
+            tokenizer=tokenizer, data_type="fbank", device="cpu",
+            language="zh", output_timestamp=True,
+        )
+        self.assertEqual(results[0]["timestamp"], [[0, 30], [90, 150]])
+        self.assertEqual(results[0]["words"], ["你", "好"])
+
+    def test_word_boundaries_empty_input_and_nonmutation(self):
+        methods = load_model_methods()
+        self.assertIn("post", methods, "The timestamp normalization must be part of the model")
+        convert = lambda values: methods["post"](None, values)
+        self.assertEqual(convert([]), ([], []))
+        tokens = [
+            ["▁Hel", 0.0, 0.06], ["lo", 0.06, 0.12],
+            ["▁", 0.12, 0.15], ["world", 0.15, 0.24],
+            ["你", 0.24, 0.30], ["好", 0.30, 0.36], ["!", 0.36, 0.42],
+        ]
+        original = [token[:] for token in tokens]
+        self.assertEqual(
+            convert(tokens),
+            ([[0, 120], [150, 240], [240, 300], [300, 360], [360, 420]],
+             ["Hello", "world", "你", "好", "!"]),
+        )
+        self.assertEqual(tokens, original)
+
+
+if __name__ == "__main__":
+    unittest.main()