129 lines
3.7 KiB
Python
129 lines
3.7 KiB
Python
"""LLM post-correction for ASR transcripts (Ark / DeepSeek)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from typing import Callable
|
|
|
|
# Soft domain hints for the model — not string-replace rules.
|
|
DEFAULT_CONTEXT = (
|
|
"这是个人助理「小荷」（Xiaohe）客户端的语音输入。"
|
|
"识别结果可能把专有名词听错成同音字（例如产品名）。"
|
|
)
|
|
|
|
DEFAULT_ARK_MODEL = os.environ.get(
|
|
"MYTOOLKIT_ASR_CORRECT_MODEL", "doubao-seed-1-6-flash-250828"
|
|
)
|
|
DEFAULT_DEEPSEEK_MODEL = os.environ.get(
|
|
"MYTOOLKIT_ASR_CORRECT_DEEPSEEK_MODEL", "deepseek-chat"
|
|
)
|
|
|
|
|
|
def _system_prompt(context: str) -> str:
|
|
system = (
|
|
"你是语音识别结果校对助手。根据上下文，轻量修正明显的识别错误"
|
|
"（同音字、专有名词、标点），保留原意与口语风格。"
|
|
"不要新增原句没有的信息，不要翻译，不要解释。"
|
|
"只输出修正后的完整句子。"
|
|
)
|
|
if context.strip():
|
|
system += f"\n背景：{context.strip()}"
|
|
return system
|
|
|
|
|
|
def _unwrap(text: str) -> str:
|
|
out = (text or "").strip()
|
|
if len(out) >= 2 and out[0] == out[-1] and out[0] in "\"'“”":
|
|
out = out[1:-1].strip()
|
|
return out
|
|
|
|
|
|
def _chat_correct(
|
|
*,
|
|
base_url: str,
|
|
api_key: str,
|
|
model: str,
|
|
text: str,
|
|
context: str,
|
|
timeout_sec: float,
|
|
) -> str:
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI(base_url=base_url, api_key=api_key, timeout=timeout_sec)
|
|
resp = client.chat.completions.create(
|
|
model=model,
|
|
temperature=0.1,
|
|
max_tokens=min(800, max(64, len(text) * 3)),
|
|
messages=[
|
|
{"role": "system", "content": _system_prompt(context)},
|
|
{"role": "user", "content": text},
|
|
],
|
|
)
|
|
return _unwrap(resp.choices[0].message.content or "") or text
|
|
|
|
|
|
def correct_asr_text(
|
|
text: str,
|
|
*,
|
|
context: str = DEFAULT_CONTEXT,
|
|
ark_api_key: str | None = None,
|
|
deepseek_api_key: str | None = None,
|
|
ark_model: str | None = None,
|
|
deepseek_model: str | None = None,
|
|
timeout_sec: float = 20.0,
|
|
on_skip: Callable[[str], None] | None = None,
|
|
) -> str:
|
|
"""Correct an ASR transcript with a chat model. Returns corrected text.
|
|
|
|
Tries Ark first (if key present), then DeepSeek. On failure, returns the
|
|
original text. Domain hints are prompt context only — no hardcoded replaces.
|
|
"""
|
|
cleaned = (text or "").strip()
|
|
if not cleaned:
|
|
return ""
|
|
|
|
attempts: list[tuple[str, str, str, str]] = []
|
|
if ark_api_key:
|
|
attempts.append(
|
|
(
|
|
"ark",
|
|
"https://ark.cn-beijing.volces.com/api/v3",
|
|
ark_api_key,
|
|
(ark_model or DEFAULT_ARK_MODEL).strip() or DEFAULT_ARK_MODEL,
|
|
)
|
|
)
|
|
if deepseek_api_key:
|
|
attempts.append(
|
|
(
|
|
"deepseek",
|
|
"https://api.deepseek.com",
|
|
deepseek_api_key,
|
|
(deepseek_model or DEFAULT_DEEPSEEK_MODEL).strip()
|
|
or DEFAULT_DEEPSEEK_MODEL,
|
|
)
|
|
)
|
|
|
|
if not attempts:
|
|
if on_skip:
|
|
on_skip("no LLM API key for ASR correction")
|
|
return cleaned
|
|
|
|
last_err = ""
|
|
for name, base_url, key, model in attempts:
|
|
try:
|
|
return _chat_correct(
|
|
base_url=base_url,
|
|
api_key=key,
|
|
model=model,
|
|
text=cleaned,
|
|
context=context,
|
|
timeout_sec=timeout_sec,
|
|
)
|
|
except Exception as exc: # noqa: BLE001
|
|
last_err = f"{name}: {exc}"
|
|
continue
|
|
|
|
if on_skip:
|
|
on_skip(last_err or "ASR correction failed")
|
|
return cleaned
|