diff --git a/kb_workflow/srt/corrector.py b/kb_workflow/srt/corrector.py new file mode 100644 index 0000000..5609578 --- /dev/null +++ b/kb_workflow/srt/corrector.py @@ -0,0 +1,157 @@ +"""ASR Corrector Module + +使用 LLM 纠正语音识别错误。复用 SimpleAIClient。 + +功能: +- 纠正明显的语音识别错误(同音字、漏字、错字) +- 分批处理长文本 +- AI 不可用时返回原文 +""" + +from __future__ import annotations + +import logging +from typing import Optional, Protocol + +logger = logging.getLogger(__name__) + + +class AIClientProtocol(Protocol): + """AI 客户端协议""" + + def generate(self, prompt: str, max_tokens: int = 16000) -> str: + """生成文本""" + ... + + +# 纠错 Prompt 模板 +CORRECTION_PROMPT = """你是一个文本校对专家。以下是语音识别生成的字幕,可能包含识别错误。 +请只纠正明显的错误(如同音字、漏字、错字),不要改变原意和风格。 + +原文: +{text} + +纠正后的文本(只输出纠正后的文本,不要解释):""" + + +class ASRCorrector: + """ASR 纠错器 + + 使用 LLM 纠正语音识别中的明显错误。 + + Attributes: + ai_client: AI 客户端(必须实现 generate 方法) + batch_size: 分批处理的字符数,默认 5000 + """ + + def __init__( + self, + ai_client: Optional[AIClientProtocol] = None, + batch_size: int = 5000 + ): + """初始化纠错器 + + Args: + ai_client: AI 客户端实例,如果为 None 则不进行纠错 + batch_size: 分批处理的字符数 + """ + self.ai_client = ai_client + self.batch_size = batch_size + + def correct(self, text: str) -> str: + """纠正文本中的语音识别错误 + + Args: + text: 原始文本 + + Returns: + 纠正后的文本。如果 AI 不可用或文本为空,返回原文。 + + Example: + >>> corrector = ASRCorrector(ai_client=my_client) + >>> result = corrector.correct("什么是机器学习") + >>> print(result) + 什么是机器学习 + """ + # 空文本直接返回 + if not text or not text.strip(): + return text or "" + + # AI 不可用,返回原文 + if self.ai_client is None: + logger.info("AI client not available, returning original text") + return text + + # 短文本直接处理 + if len(text) <= self.batch_size: + return self._correct_batch(text) + + # 长文本分批处理 + return self._correct_long_text(text) + + def _correct_batch(self, text: str) -> str: + """处理单个批次的文本 + + Args: + text: 文本内容 + + Returns: + 纠正后的文本 + """ + try: + prompt = CORRECTION_PROMPT.format(text=text) + result = self.ai_client.generate(prompt, max_tokens=len(text) + 500) + return result.strip() + except Exception as e: + logger.warning(f"ASR correction failed: {e}, returning original text") + return text + + def _correct_long_text(self, text: str) -> str: + """分批处理长文本 + + 按句子分割,每批不超过 batch_size 字符。 + + Args: + text: 长文本内容 + + Returns: + 纠正后的文本 + """ + # 按句子分割(简单分割,基于句号、问号、感叹号) + sentences = [] + current = "" + for char in text: + current += char + if char in '。!?\n': + sentences.append(current) + current = "" + + # 添加最后一个不完整的句子 + if current: + sentences.append(current) + + # 分批处理 + batches = [] + current_batch = [] + current_length = 0 + + for sentence in sentences: + if current_length + len(sentence) > self.batch_size: + if current_batch: + batches.append(''.join(current_batch)) + current_batch = [sentence] + current_length = len(sentence) + else: + current_batch.append(sentence) + current_length += len(sentence) + + if current_batch: + batches.append(''.join(current_batch)) + + # 处理每个批次 + corrected_batches = [] + for batch in batches: + corrected = self._correct_batch(batch) + corrected_batches.append(corrected) + + return ''.join(corrected_batches)