diff --git a/.github/workflows/sensevoice-container.yml b/.github/workflows/sensevoice-container.yml index 9131f95..a0b9c68 100644 --- a/.github/workflows/sensevoice-container.yml +++ b/.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 diff --git a/README.md b/README.md index 3db566e..018a9e8 100644 --- a/README.md +++ b/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 diff --git a/README_zh.md b/README_zh.md index d501ec8..23e0d88 100644 --- a/README_zh.md +++ b/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 diff --git a/demo2.py b/demo2.py index 341d67d..66d10f1 100644 --- a/demo2.py +++ b/demo2.py @@ -32,6 +32,10 @@ ) 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}") diff --git a/model.py b/model.py index b32b931..19fe279 100644 --- a/model.py +++ b/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 @@ def inference( from itertools import groupby timestamp = [] tokens = tokenizer.text2tokens(text)[4:] - - logits_speech = self.ctc.softmax(encoder_out)[i, 4:encoder_out_lens[i].item(), :] + 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.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 diff --git a/tests/test_model_timestamps.py b/tests/test_model_timestamps.py new file mode 100644 index 0000000..722a98b --- /dev/null +++ b/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()