wehub-resource-sync
399 lines
14 KiB
Python
399 lines
14 KiB
Python
# 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
|