diff --git a/docs/tutorial/README.md b/docs/tutorial/README.md index 269ce7fdb..63cb541bf 100644 --- a/docs/tutorial/README.md +++ b/docs/tutorial/README.md @@ -86,6 +86,31 @@ wav_file = f"{model.model_path}/example/asr_example.wav" res = model.generate(input=wav_file, batch_size_s=300, batch_size_threshold_s=60, hotword='魔搭') print(res) ``` + +#### Postprocess hotword correction + +Model-level `hotword` / `hotwords` boosts a small set of terms during decoding. For large vocabularies (for example thousands of stock names), apply text-level correction after recognition: + +```python +from funasr import AutoModel + +model = AutoModel(model="paraformer-zh", vad_model="fsmn-vad", punc_model="ct-punc") + +res = model.generate( + input="asr_example.wav", + postprocess_hotwords={ + "科大迅飞": "科大讯飞", + "东方财富": "东方财富", + }, + postprocess_hotword_threshold=0.85, + return_postprocess_hotword_matches=True, +) +print(res[0]["text"]) +print(res[0].get("postprocess_hotword_matches")) +``` + +You can also pass `postprocess_hotword_file` with one target word per line, or an explicit mapping such as `wrong=>right`. Fuzzy matching requires optional `pypinyin` and `rapidfuzz`; explicit mappings work without them. + Notes: - Typically, the input duration for models is limited to under 30 seconds. However, when combined with `vad_model`, support for audio input of any length is enabled, not limited to the paraformer model—any audio input model can be used. - Parameters related to model can be directly specified in the definition of AutoModel; parameters related to `vad_model` can be set through `vad_kwargs`, which is a dict; similar parameters include `punc_kwargs` and `spk_kwargs`. diff --git a/docs/tutorial/README_zh.md b/docs/tutorial/README_zh.md index fdbf0e761..f02470e5b 100644 --- a/docs/tutorial/README_zh.md +++ b/docs/tutorial/README_zh.md @@ -87,6 +87,31 @@ wav_file = f"{model.model_path}/example/asr_example.wav" res = model.generate(input=wav_file, batch_size_s=300, batch_size_threshold_s=60, hotword='魔搭') print(res) ``` + +#### ASR 后处理热词替换 + +模型层 `hotword` / `hotwords` 用于在解码阶段提升少量高优先级词的识别率；如果词表很大（例如几千个股票名），更适合在识别完成后做文本级纠错： + +```python +from funasr import AutoModel + +model = AutoModel(model="paraformer-zh", vad_model="fsmn-vad", punc_model="ct-punc") + +res = model.generate( + input="asr_example.wav", + postprocess_hotwords={ + "科大迅飞": "科大讯飞", + "东方财富": "东方财富", + }, + postprocess_hotword_threshold=0.85, + return_postprocess_hotword_matches=True, +) +print(res[0]["text"]) +print(res[0].get("postprocess_hotword_matches")) +``` + +也支持热词文件 `postprocess_hotword_file`，每行一个目标词，或一行一个显式映射（`错误词=>目标词`）。模糊匹配需要额外安装 `pypinyin` 和 `rapidfuzz`；未安装时仍可使用显式映射。 + 注意： - 通常模型输入限制时长30s以下，组合`vad_model`后，支持任意时长音频输入，不局限于paraformer模型，所有音频输入模型均可以。 - `model`相关的参数可以直接在`AutoModel`定义中直接指定；与`vad_model`相关参数可以通过`vad_kwargs`来指定，类型为dict；类似的有`punc_kwargs`，`spk_kwargs`； diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index f7bcae26e..a92309198 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -28,6 +28,7 @@ from funasr.train_utils.set_all_random_seed import set_all_random_seed from funasr.train_utils.load_pretrained_model import load_pretrained_model from funasr.utils import export_utils +from funasr.utils.postprocess_hotwords import apply_postprocess_hotwords_to_results from funasr.utils import misc @@ -458,6 +459,12 @@ def generate(self, input, input_len=None, progress_callback=None, **cfg): **cfg: Runtime parameters: - cache (dict): State cache for streaming mode. Pass {} for first call. - hotword (str/list): Keywords to boost recognition accuracy. + - postprocess_hotwords (str/list/dict): Text-level hotword correction after + decoding. Unlike model-level ``hotword``, this runs on the final text. + - postprocess_hotword_file (str): Hotword file path. Each line is a target + word or an explicit mapping like ``错误词=>目标词``. + - postprocess_hotword_threshold (float): Fuzzy match threshold in [0, 1]. + - return_postprocess_hotword_matches (bool): Include replacement details. - language (str): Language hint ("auto", "zh", "en", "Chinese", etc.) - batch_size_s (int): Dynamic batch total duration in seconds. - is_final (bool): Last chunk flag for streaming mode. @@ -486,12 +493,13 @@ def generate(self, input, input_len=None, progress_callback=None, **cfg): if cfg.get("return_raw_text", self.kwargs.get("return_raw_text", False)): result["raw_text"] = copy.copy(result["text"]) result["text"] = punc_res[0]["text"] - return results + return apply_postprocess_hotwords_to_results(results, cfg) else: - return self.inference_with_vad( + results = self.inference_with_vad( input, input_len=input_len, progress_callback=progress_callback, **cfg ) + return apply_postprocess_hotwords_to_results(results, cfg) def inference( self, diff --git a/funasr/utils/postprocess_hotwords.py b/funasr/utils/postprocess_hotwords.py new file mode 100644 index 000000000..26e482697 --- /dev/null +++ b/funasr/utils/postprocess_hotwords.py @@ -0,0 +1,398 @@ +# Copyright FunASR (https://github.com/modelscope/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +"""Text-level hotword correction after ASR decoding. + +This module is intentionally separate from model-level ``hotword`` / ``hotwords`` +prompting. It runs after ASR (and punctuation / ITN when configured) and only +updates top-level ``text`` plus sentence-level ``text`` / ``sentence`` fields. +""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union + +HotwordInput = Union[str, Sequence[str], Mapping[str, str], None] + +_EXPLICIT_SEPARATORS = ("=>", "->", "→") +_TOKEN_PATTERN = re.compile(r"[\u4e00-\u9fff]|[a-zA-Z]+|[0-9]+") + +_LAZY_PINYIN = None +_PINYIN_STYLE = None +_RAPIDFUZZ_FUZZ = None + + +@dataclass(frozen=True) +class HotwordMatch: + """A single postprocess hotword replacement.""" + + original: str + replacement: str + score: float + start: int + end: int + + def as_dict(self) -> Dict[str, Any]: + return { + "original": self.original, + "replacement": self.replacement, + "score": self.score, + "start": self.start, + "end": self.end, + } + + +def _require_pypinyin(): + global _LAZY_PINYIN, _PINYIN_STYLE + if _LAZY_PINYIN is None: + try: + from pypinyin import Style, lazy_pinyin + except ImportError as exc: + raise ImportError( + "postprocess hotword fuzzy matching requires pypinyin. " + "Install it with: pip install pypinyin" + ) from exc + _LAZY_PINYIN = lazy_pinyin + _PINYIN_STYLE = Style + return _LAZY_PINYIN, _PINYIN_STYLE + + +def _require_rapidfuzz(): + global _RAPIDFUZZ_FUZZ + if _RAPIDFUZZ_FUZZ is None: + try: + from rapidfuzz import fuzz + except ImportError as exc: + raise ImportError( + "postprocess hotword fuzzy matching requires rapidfuzz. " + "Install it with: pip install rapidfuzz" + ) from exc + _RAPIDFUZZ_FUZZ = fuzz + return _RAPIDFUZZ_FUZZ + + +def _to_pinyin_key(text: str) -> str: + lazy_pinyin, style = _require_pypinyin() + return "".join(lazy_pinyin(text, style=style.NORMAL, errors="ignore")).lower() + + +def _parse_line(line: str) -> Tuple[Optional[str], Optional[str], bool]: + """Parse one hotword file line. + + Returns: + (wrong, right, is_explicit) + For fuzzy-only targets, wrong is None and right is the target word. + """ + stripped = line.strip() + if not stripped or stripped.startswith("#"): + return None, None, False + + for sep in _EXPLICIT_SEPARATORS: + if sep in stripped: + wrong, right = stripped.split(sep, 1) + wrong = wrong.strip() + right = right.strip() + if wrong and right: + return wrong, right, True + return None, None, False + + return None, stripped, False + + +def parse_hotword_file(path: str) -> Tuple[Dict[str, str], List[str]]: + if not os.path.isfile(path): + raise FileNotFoundError(f"postprocess_hotword_file not found: {path}") + + explicit: Dict[str, str] = {} + fuzzy_targets: List[str] = [] + with open(path, "r", encoding="utf-8") as f: + for line in f: + wrong, right, is_explicit = _parse_line(line) + if not right: + continue + if is_explicit and wrong is not None: + explicit[wrong] = right + else: + fuzzy_targets.append(right) + return explicit, fuzzy_targets + + +def parse_postprocess_hotwords( + postprocess_hotwords: HotwordInput, +) -> Tuple[Dict[str, str], List[str]]: + """Parse in-memory hotword config into explicit and fuzzy buckets.""" + explicit: Dict[str, str] = {} + fuzzy_targets: List[str] = [] + + if postprocess_hotwords is None: + return explicit, fuzzy_targets + + if isinstance(postprocess_hotwords, str): + for line in postprocess_hotwords.splitlines(): + wrong, right, is_explicit = _parse_line(line) + if not right: + continue + if is_explicit and wrong is not None: + explicit[wrong] = right + else: + fuzzy_targets.append(right) + return explicit, fuzzy_targets + + if isinstance(postprocess_hotwords, Mapping): + for wrong, right in postprocess_hotwords.items(): + wrong_s = str(wrong).strip() + right_s = str(right).strip() + if not right_s: + continue + if wrong_s and wrong_s != right_s: + explicit[wrong_s] = right_s + else: + fuzzy_targets.append(right_s) + return explicit, fuzzy_targets + + if isinstance(postprocess_hotwords, Sequence) and not isinstance(postprocess_hotwords, (str, bytes)): + for item in postprocess_hotwords: + if item is None: + continue + item_s = str(item).strip() + if not item_s: + continue + wrong, right, is_explicit = _parse_line(item_s) + if is_explicit and wrong is not None: + explicit[wrong] = right + elif right: + fuzzy_targets.append(right) + return explicit, fuzzy_targets + + raise TypeError( + "postprocess_hotwords must be None, str, list, or dict; " + f"got {type(postprocess_hotwords)!r}" + ) + + +def build_postprocess_hotword_matcher( + postprocess_hotwords: HotwordInput = None, + postprocess_hotword_file: Optional[str] = None, + postprocess_hotword_threshold: float = 0.85, + enable_fuzzy: bool = True, +) -> Optional["PostprocessHotwordMatcher"]: + """Compile a matcher once per ``generate()`` call.""" + explicit: Dict[str, str] = {} + fuzzy_targets: List[str] = [] + + if postprocess_hotwords is not None: + e, f = parse_postprocess_hotwords(postprocess_hotwords) + explicit.update(e) + fuzzy_targets.extend(f) + + if postprocess_hotword_file: + e, f = parse_hotword_file(postprocess_hotword_file) + explicit.update(e) + fuzzy_targets.extend(f) + + if not explicit and not fuzzy_targets: + return None + + return PostprocessHotwordMatcher( + explicit_map=explicit, + fuzzy_targets=fuzzy_targets, + threshold=postprocess_hotword_threshold, + enable_fuzzy=enable_fuzzy, + ) + + +class PostprocessHotwordMatcher: + """Compiled matcher reused across all results in one generate() call.""" + + def __init__( + self, + explicit_map: Optional[Dict[str, str]] = None, + fuzzy_targets: Optional[Iterable[str]] = None, + threshold: float = 0.85, + enable_fuzzy: bool = True, + ): + self.explicit_map = dict(explicit_map or {}) + self.threshold = float(threshold) + if not 0.0 <= self.threshold <= 1.0: + raise ValueError( + f"postprocess_hotword_threshold must be between 0.0 and 1.0, got {threshold}" + ) + self.enable_fuzzy = bool(enable_fuzzy) + + seen = set() + self.fuzzy_targets: List[str] = [] + for target in fuzzy_targets or []: + target_s = str(target).strip() + if target_s and target_s not in seen: + seen.add(target_s) + self.fuzzy_targets.append(target_s) + + self._length_buckets: Dict[int, List[Tuple[str, str]]] = {} + self._fuzz = None + if self.fuzzy_targets and self.enable_fuzzy: + self._fuzz = _require_rapidfuzz() + _require_pypinyin() + for target in self.fuzzy_targets: + bucket = self._length_buckets.setdefault(len(target), []) + bucket.append((target, _to_pinyin_key(target))) + + def apply_text(self, text: str) -> Tuple[str, List[HotwordMatch]]: + if not text: + return text, [] + + matches: List[HotwordMatch] = [] + updated = self._apply_explicit(text, matches) + if self.fuzzy_targets and self.enable_fuzzy: + updated, fuzzy_matches = self._apply_fuzzy(updated) + matches.extend(fuzzy_matches) + return updated, matches + + def apply_result(self, result: Dict[str, Any], return_matches: bool = False) -> Dict[str, Any]: + text = result.get("text", "") + if not isinstance(text, str) or not text: + if return_matches: + result["postprocess_hotword_matches"] = [] + return result + + original_timestamp = result.get("timestamp") + new_text, matches = self.apply_text(text) + result["text"] = new_text + + sentence_info = result.get("sentence_info") + if isinstance(sentence_info, list): + for sentence in sentence_info: + if not isinstance(sentence, dict): + continue + for field in ("text", "sentence"): + if field in sentence and isinstance(sentence[field], str): + corrected, _ = self.apply_text(sentence[field]) + sentence[field] = corrected + + if return_matches: + result["postprocess_hotword_matches"] = [m.as_dict() for m in matches] + + # Timestamps intentionally remain aligned to the original recognition. + if original_timestamp is not None: + result["timestamp"] = original_timestamp + return result + + def _apply_explicit(self, text: str, matches: List[HotwordMatch]) -> str: + if not self.explicit_map: + return text + + updated = text + for wrong in sorted(self.explicit_map, key=len, reverse=True): + right = self.explicit_map[wrong] + start = 0 + while True: + idx = updated.find(wrong, start) + if idx < 0: + break + end = idx + len(wrong) + matches.append( + HotwordMatch( + original=wrong, + replacement=right, + score=1.0, + start=idx, + end=end, + ) + ) + updated = updated[:idx] + right + updated[end:] + start = idx + len(right) + return updated + + def _apply_fuzzy(self, text: str) -> Tuple[str, List[HotwordMatch]]: + assert self._fuzz is not None + candidates: List[HotwordMatch] = [] + + if not self._length_buckets: + return text, [] + + min_len = min(self._length_buckets) + max_len = max(self._length_buckets) + text_len = len(text) + + for win_len in range(max(1, min_len - 1), max_len + 2): + bucket_keys = [ + length + for length in (win_len - 1, win_len, win_len + 1) + if length in self._length_buckets + ] + if not bucket_keys: + continue + + for start in range(0, text_len - win_len + 1): + end = start + win_len + segment = text[start:end] + if not segment or not _TOKEN_PATTERN.search(segment): + continue + segment_py = _to_pinyin_key(segment) + + for length in bucket_keys: + for target, target_py in self._length_buckets[length]: + if segment == target: + continue + score = self._fuzz.ratio(segment_py, target_py) / 100.0 + if score >= self.threshold: + candidates.append( + HotwordMatch( + original=segment, + replacement=target, + score=round(score, 4), + start=start, + end=end, + ) + ) + + if not candidates: + return text, [] + + selected = _select_non_overlapping(candidates) + updated = text + applied: List[HotwordMatch] = [] + for match in sorted(selected, key=lambda m: m.start, reverse=True): + updated = updated[: match.start] + match.replacement + updated[match.end :] + applied.append(match) + applied.sort(key=lambda m: m.start) + return updated, applied + + +def _select_non_overlapping(candidates: List[HotwordMatch]) -> List[HotwordMatch]: + ranked = sorted(candidates, key=lambda m: (m.score, m.end - m.start), reverse=True) + selected: List[HotwordMatch] = [] + occupied: List[Tuple[int, int]] = [] + + for candidate in ranked: + if any(not (candidate.end <= start or candidate.start >= end) for start, end in occupied): + continue + selected.append(candidate) + occupied.append((candidate.start, candidate.end)) + + return sorted(selected, key=lambda m: m.start) + + +def apply_postprocess_hotwords_to_results( + results: List[Dict[str, Any]], + cfg: Mapping[str, Any], +) -> List[Dict[str, Any]]: + """Apply compiled matcher to each result dict if configured in cfg.""" + matcher = build_postprocess_hotword_matcher( + postprocess_hotwords=cfg.get("postprocess_hotwords"), + postprocess_hotword_file=cfg.get("postprocess_hotword_file"), + postprocess_hotword_threshold=cfg.get("postprocess_hotword_threshold", 0.85), + enable_fuzzy=cfg.get("postprocess_hotword_fuzzy", True), + ) + if matcher is None: + return results + + return_matches = bool( + cfg.get("return_postprocess_hotword_matches", False) + ) + for result in results: + if isinstance(result, dict): + matcher.apply_result(result, return_matches=return_matches) + return results diff --git a/tests/test_postprocess_hotwords.py b/tests/test_postprocess_hotwords.py new file mode 100644 index 000000000..d00a79086 --- /dev/null +++ b/tests/test_postprocess_hotwords.py @@ -0,0 +1,163 @@ +import importlib.util +import os +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +_MODULE_PATH = Path(__file__).resolve().parents[1] / "funasr" / "utils" / "postprocess_hotwords.py" +_SPEC = importlib.util.spec_from_file_location("postprocess_hotwords", _MODULE_PATH) +postprocess_hotwords = importlib.util.module_from_spec(_SPEC) +sys.modules["postprocess_hotwords"] = postprocess_hotwords +assert _SPEC.loader is not None +_SPEC.loader.exec_module(postprocess_hotwords) + +HotwordMatch = postprocess_hotwords.HotwordMatch +PostprocessHotwordMatcher = postprocess_hotwords.PostprocessHotwordMatcher +apply_postprocess_hotwords_to_results = postprocess_hotwords.apply_postprocess_hotwords_to_results +build_postprocess_hotword_matcher = postprocess_hotwords.build_postprocess_hotword_matcher +parse_hotword_file = postprocess_hotwords.parse_hotword_file +parse_postprocess_hotwords = postprocess_hotwords.parse_postprocess_hotwords + + +class TestPostprocessHotwordParsing(unittest.TestCase): + def test_parse_list_and_dict(self): + explicit, fuzzy = parse_postprocess_hotwords(["科大讯飞", "东方财富"]) + self.assertEqual(explicit, {}) + self.assertEqual(fuzzy, ["科大讯飞", "东方财富"]) + + explicit, fuzzy = parse_postprocess_hotwords({"科大迅飞": "科大讯飞", "东方财富": "东方财富"}) + self.assertEqual(explicit, {"科大迅飞": "科大讯飞"}) + self.assertEqual(fuzzy, ["东方财富"]) + + def test_parse_inline_mapping(self): + explicit, fuzzy = parse_postprocess_hotwords(["撒贝你=>撒贝宁", "康辉"]) + self.assertEqual(explicit, {"撒贝你": "撒贝宁"}) + self.assertEqual(fuzzy, ["康辉"]) + + def test_parse_hotword_file(self): + with tempfile.NamedTemporaryFile("w", encoding="utf-8", delete=False) as f: + f.write("# comment\n") + f.write("科大讯飞\n") + f.write("科大迅飞=>科大讯飞\n") + path = f.name + + try: + explicit, fuzzy = parse_hotword_file(path) + self.assertEqual(explicit, {"科大迅飞": "科大讯飞"}) + self.assertEqual(fuzzy, ["科大讯飞"]) + finally: + os.unlink(path) + + +class TestPostprocessHotwordMatcher(unittest.TestCase): + def test_explicit_replace(self): + matcher = PostprocessHotwordMatcher( + explicit_map={"撒贝你": "撒贝宁"}, + enable_fuzzy=False, + ) + text, matches = matcher.apply_text("我非常喜欢撒贝你说的新闻") + self.assertEqual(text, "我非常喜欢撒贝宁说的新闻") + self.assertEqual(len(matches), 1) + self.assertEqual(matches[0].replacement, "撒贝宁") + self.assertEqual(matches[0].score, 1.0) + + def test_fuzzy_replace_with_optional_deps(self): + try: + import pypinyin # noqa: F401 + import rapidfuzz # noqa: F401 + except ImportError: + self.skipTest("pypinyin and rapidfuzz are required for fuzzy tests") + + matcher = PostprocessHotwordMatcher( + fuzzy_targets=["科大讯飞"], + threshold=0.75, + ) + text, matches = matcher.apply_text("科大迅飞的语音识别很强") + self.assertIn("科大讯飞", text) + self.assertTrue(matches) + self.assertEqual(matches[0].replacement, "科大讯飞") + + def test_missing_fuzzy_dependency_raises(self): + with mock.patch.object( + postprocess_hotwords, + "_require_pypinyin", + side_effect=ImportError("missing pypinyin"), + ): + with self.assertRaises(ImportError): + PostprocessHotwordMatcher( + fuzzy_targets=["科大讯飞"], + threshold=0.85, + ) + + def test_invalid_threshold_raises(self): + with self.assertRaises(ValueError): + PostprocessHotwordMatcher(explicit_map={"a": "b"}, threshold=1.5) + + def test_sentence_info_and_timestamp_preserved(self): + matcher = PostprocessHotwordMatcher( + explicit_map={"撒贝你": "撒贝宁"}, + enable_fuzzy=False, + ) + result = { + "text": "撒贝你主持节目", + "timestamp": [[0, 100], [100, 200]], + "sentence_info": [ + {"text": "撒贝你主持", "sentence": "撒贝你主持", "start": 0, "end": 1000}, + {"text": "节目", "sentence": "节目", "start": 1000, "end": 1500}, + ], + } + matcher.apply_result(result, return_matches=True) + self.assertEqual(result["text"], "撒贝宁主持节目") + self.assertEqual(result["sentence_info"][0]["text"], "撒贝宁主持") + self.assertEqual(result["sentence_info"][0]["sentence"], "撒贝宁主持") + self.assertEqual(result["timestamp"], [[0, 100], [100, 200]]) + self.assertEqual(result["postprocess_hotword_matches"][0]["replacement"], "撒贝宁") + + def test_matcher_reused_for_multiple_results(self): + build_calls = {"count": 0} + original_build = build_postprocess_hotword_matcher + + def counting_build(*args, **kwargs): + build_calls["count"] += 1 + return original_build(*args, **kwargs) + + results = [ + {"text": "撒贝你主持", "timestamp": [1]}, + {"text": "康灰播报", "timestamp": [2]}, + ] + cfg = { + "postprocess_hotwords": {"撒贝你": "撒贝宁", "康灰": "康辉"}, + "return_postprocess_hotword_matches": True, + } + with mock.patch.object( + postprocess_hotwords, + "build_postprocess_hotword_matcher", + side_effect=counting_build, + ): + updated = apply_postprocess_hotwords_to_results(results, cfg) + + self.assertEqual(build_calls["count"], 1) + self.assertEqual(updated[0]["text"], "撒贝宁主持") + self.assertEqual(updated[1]["text"], "康辉播报") + self.assertEqual(len(updated[0]["postprocess_hotword_matches"]), 1) + + def test_noop_when_not_configured(self): + results = [{"text": "不变", "timestamp": [1]}] + updated = apply_postprocess_hotwords_to_results(results, {}) + self.assertIs(updated, results) + self.assertEqual(updated[0]["text"], "不变") + + +class TestHotwordMatch(unittest.TestCase): + def test_as_dict(self): + match = HotwordMatch("a", "b", 0.9, 1, 2) + self.assertEqual( + match.as_dict(), + {"original": "a", "replacement": "b", "score": 0.9, "start": 1, "end": 2}, + ) + + +if __name__ == "__main__": + unittest.main()