cat-shark
5ffa2ac39e
- 规则层（不调 LLM）：横线装饰/HTML 水印 token/URL/邮箱/单双 ASCII 字符直接删 - LLM 五类分类：garbage/overlay/noise 删，repeat/dialogue 留，未识别回退保留 - 按文本去重：相同文本只调一次 LLM（忽略空白/大小写），判定一致并省调用 - 长文本保护：≥min_keep_len 时 noise 不构成删除依据 - 真实任务 run_ac7f480a3ccb 验证：非规则误删 350→193（-45%），呻吟/对话保留 - ocr-subtitle 工作流 v4：pool 钉死单线程（1/1） - 回归夹具 testdata/ocr_srt_run_ac7f480a3ccb.srt（真实 1666 条 OCR 输出）
325 lines
14 KiB
Python
325 lines
14 KiB
Python
"""LLM 字幕过滤节点。
|
||
|
||
对 OCR 识别出的 SRT 字幕做二次过滤，两级判断：
|
||
|
||
1. **确定性规则层**（不调用 LLM）：横线装饰、HTML/水印 token、URL/邮箱、
|
||
单双 ASCII 字符等 OCR 噪声直接删除——这些模式是稳定可判的，走规则
|
||
既省 token 又保证结果确定（真实数据中占删除量的 65%+）。
|
||
|
||
2. **LLM 五类分类层**：把目标字幕连同前后各 context_size 条纯文本分批
|
||
提供给 LLM，模型输出五个类别之一：
|
||
- garbage：垃圾字符（乱码、残缺、装饰性符号）→ 删除
|
||
- overlay：水印/网页/播放器等覆盖层文本 → 删除
|
||
- noise：与上下文无关、无实际语义的杂项 → 删除
|
||
- repeat：内容性重复（语气词、呻吟、重复感叹，属于内容本身）→ 保留
|
||
- dialogue：正常对话 → 保留
|
||
未识别输出一律回退 dialogue（宁滥勿缺，避免误删真实对话）。
|
||
|
||
3. **文本去重**：相同文本（忽略全部空白与大小写差异）只调一次 LLM，
|
||
上下文取首次出现位置，结果缓存复用——修复"同一句字幕 5 留 6 删"的
|
||
判定不一致，同时把长视频的 LLM 调用量降到唯一文本数。
|
||
|
||
4. **长文本保护**：长度 ≥ min_keep_len 的文本，仅 garbage/overlay 两个
|
||
明确垃圾类别可删，noise 不构成删除依据，防止长对话被误删。
|
||
|
||
参考真实任务 run_ac7f480a3ccb（2026-08，OCR 1666 条）：旧实现把 394 条
|
||
真实对话当噪声删掉（占删除 35%）、同文本判定不一致；新实现按上述机制
|
||
回归测试已固化在 tests/test_llm_filter.py。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
from nodes.adaptive_pool import AdaptiveThreadPool
|
||
from wov_app.logging import get_logger
|
||
from wov_sdk.models import InvokeRequest, InvokeResponse
|
||
|
||
logger = get_logger("llm-filter")
|
||
|
||
# 匹配 SRT 条目：时间轴行 + 文本（文本可多行），到下一个序号行或文末结束。
|
||
_SRT_BLOCK_RE = re.compile(
|
||
r"(\d{2}:\d{2}:\d{2},\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2},\d{3})\s*\n(.*?)(?=\n\s*\d+\s*\n|\Z)",
|
||
re.DOTALL,
|
||
)
|
||
|
||
# 目标字幕标记：提示词用该标记指明需要判断的那一条字幕。
|
||
TARGET_MARK = "【目标】"
|
||
|
||
# 默认上下文窗口：目标字幕前后各取 10 条。
|
||
DEFAULT_CONTEXT_SIZE = 10
|
||
|
||
# LLM 输出类别：garbage/overlay/noise 删除；repeat/dialogue 保留。
|
||
CATEGORY_GARBAGE = "garbage"
|
||
CATEGORY_OVERLAY = "overlay"
|
||
CATEGORY_NOISE = "noise"
|
||
CATEGORY_REPEAT = "repeat"
|
||
CATEGORY_DIALOGUE = "dialogue"
|
||
_ALL_CATEGORIES = (
|
||
CATEGORY_GARBAGE,
|
||
CATEGORY_OVERLAY,
|
||
CATEGORY_NOISE,
|
||
CATEGORY_REPEAT,
|
||
CATEGORY_DIALOGUE,
|
||
)
|
||
DELETE_CATEGORIES = {CATEGORY_GARBAGE, CATEGORY_OVERLAY, CATEGORY_NOISE}
|
||
|
||
# 规则层正则：横线装饰（含全角/半角横线、下划线、中点、句点等符号组合）。
|
||
_DASH_RE = re.compile(r"^[\s\-—_~=•・。．.、]+$")
|
||
# 规则层正则：URL / 邮箱。
|
||
_URL_OR_MAIL_RE = re.compile(r"^(https?://|www\.)\S+$|^[\w.+-]+@[\w.-]+\.\w+$")
|
||
# 默认水印/覆盖层 token（casefold 后比较，可经 overlay_tokens 参数覆盖）。
|
||
DEFAULT_OVERLAY_TOKENS = frozenset(
|
||
{"html", "background", "___", "cleaning", "buffering", "loading"}
|
||
)
|
||
# 长文本保护阈值：≥ 该长度的文本，noise 类别不构成删除依据。
|
||
DEFAULT_MIN_KEEP_LEN = 12
|
||
|
||
|
||
def _has_cjk(text: str) -> bool:
|
||
"""是否含 CJK 汉字：单字"嗯/好"等可能是内容，规则层不直接删。"""
|
||
return any("\u4e00" <= ch <= "\u9fff" for ch in text)
|
||
|
||
|
||
def parse_srt(text: str) -> list[dict]:
|
||
"""解析 SRT 文本为条目列表：[{"start", "end", "text"}]。"""
|
||
entries: list[dict] = []
|
||
for match in _SRT_BLOCK_RE.finditer(text):
|
||
entries.append(
|
||
{
|
||
"start": match.group(1),
|
||
"end": match.group(2),
|
||
"text": match.group(3).strip(),
|
||
}
|
||
)
|
||
return entries
|
||
|
||
|
||
def serialize_srt(entries: list[dict]) -> str:
|
||
"""把条目列表序列化为标准 SRT 文本（序号重新从 1 编号）。"""
|
||
blocks = [
|
||
f"{index}\n{entry['start']} --> {entry['end']}\n{entry['text']}"
|
||
for index, entry in enumerate(entries, start=1)
|
||
]
|
||
return "\n\n".join(blocks) + "\n"
|
||
|
||
|
||
def _rule_verdict(text: str, overlay_tokens: set[str]) -> bool | None:
|
||
"""确定性规则层：返回 True（删除）/ None（交给 LLM 多维判断）。
|
||
|
||
规则覆盖 OCR 噪声的稳定模式：空文本、纯横线装饰、URL/邮箱、水印 token
|
||
（HTML/background 等）、单双 ASCII 字符。含 CJK 的短文本不算垃圾，
|
||
因为"嗯/好"等可能是内容；其余情况返回 None 交由 LLM 分类。
|
||
"""
|
||
t = text.strip()
|
||
if not t:
|
||
return True
|
||
if _DASH_RE.match(t):
|
||
return True
|
||
if _URL_OR_MAIL_RE.match(t):
|
||
return True
|
||
if t.casefold() in overlay_tokens:
|
||
return True
|
||
if len(t) <= 2 and not _has_cjk(t):
|
||
return True
|
||
return None
|
||
|
||
|
||
def _dedup_key(text: str) -> str:
|
||
"""文本去重键：去掉全部空白并统一大小写。
|
||
|
||
OCR 同一句字幕常带/不带空格（"可没法胜任" vs "可 没法胜任"），
|
||
视为同一文本以保证判定一致；也用于把 LLM 调用量降到唯一文本数。
|
||
"""
|
||
return "".join(text.split()).casefold()
|
||
|
||
|
||
def _should_delete(category: str, text: str, min_keep_len: int) -> bool:
|
||
"""按 LLM 类别与长文本保护决定是否删除。
|
||
|
||
repeat/dialogue 一律保留；garbage/overlay 一律删除（含长文本——
|
||
这两个是明确的垃圾信号）；noise 对短文本删除，但 ≥min_keep_len 的
|
||
长文本不删（noise 太模糊，不足以推翻一句完整台词）。
|
||
"""
|
||
if category not in DELETE_CATEGORIES:
|
||
return False
|
||
if category == CATEGORY_NOISE and len(text) >= min_keep_len:
|
||
return False
|
||
return True
|
||
|
||
|
||
def _judge_category(
|
||
entries: list[dict], index: int, context_size: int, params: dict
|
||
) -> str:
|
||
"""调用 LLM 把目标字幕归入五类之一，返回类别词（未识别回退 dialogue）。
|
||
|
||
请求体只含目标字幕及其前后各 context_size 条字幕的纯文本（无时间戳），
|
||
目标字幕用 TARGET_MARK 标记；模型只输出一个类别英文单词。
|
||
输出无法识别（空/乱码/旧式"保留"）时回退 dialogue，宁滥勿缺。
|
||
"""
|
||
start = max(0, index - context_size)
|
||
end = min(len(entries), index + context_size + 1)
|
||
target_pos = index - start
|
||
lines = [
|
||
f"{TARGET_MARK}{text}" if pos == target_pos else text
|
||
for pos, text in enumerate(entry["text"] for entry in entries[start:end])
|
||
]
|
||
|
||
# LLM 兼容接口配置：地址/Key/模型/超时均可通过环境变量覆盖（默认 SiliconFlow）。
|
||
api_base = os.getenv(
|
||
"LLM_API_BASE",
|
||
"https://api.siliconflow.cn/v1/chat/completions",
|
||
)
|
||
api_key = os.getenv("LLM_API_KEY", "")
|
||
request_timeout = float(os.getenv("LLM_TIMEOUT_SECONDS", "60"))
|
||
model = str(params.get("model") or os.getenv("LLM_MODEL", "Qwen/Qwen3.6-35B-A3B"))
|
||
system_prompt = (
|
||
"你是字幕质量过滤器。用户会提供一段字幕序列（纯文本，不含时间戳），"
|
||
f"其中用{TARGET_MARK}标记的字幕是需要判断的目标。"
|
||
"请把目标字幕归入以下五个类别之一：\n"
|
||
"garbage：垃圾字符（乱码、残缺、装饰性符号、横线）\n"
|
||
"overlay：水印、网页/播放器/字幕组等覆盖层文本，不是视频对白\n"
|
||
"noise：与上下文无关、无实际语义的杂项\n"
|
||
"repeat：内容性重复（如语气词、呻吟、重复的感叹或对话），属于内容本身\n"
|
||
"dialogue：正常对话\n"
|
||
"只输出一个类别英文单词，不要输出其他内容。"
|
||
)
|
||
body = {
|
||
"model": model,
|
||
"messages": [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": "\n".join(lines)},
|
||
],
|
||
# 关闭推理模式：Qwen3 等模型默认会把思考过程写入 reasoning_content，
|
||
# 导致 content 为空或包含多余内容。
|
||
"enable_thinking": False,
|
||
# 类别词很短，输出上限给得很小即可。
|
||
"max_tokens": 16,
|
||
}
|
||
headers = {"Content-Type": "application/json"}
|
||
# 配置了 Key 时附带 Bearer 鉴权头。
|
||
if api_key:
|
||
headers["Authorization"] = f"Bearer {api_key}"
|
||
request = urllib.request.Request(
|
||
api_base,
|
||
data=json.dumps(body).encode("utf-8"),
|
||
headers=headers,
|
||
method="POST",
|
||
)
|
||
with urllib.request.urlopen(request, timeout=request_timeout) as response:
|
||
payload = json.loads(response.read().decode("utf-8"))
|
||
content = str(payload["choices"][0]["message"]["content"]).strip().lower()
|
||
# 精确匹配五个类别词（兼容"garbagexxx"这类多余输出）。
|
||
for category in _ALL_CATEGORIES:
|
||
if content == category or content.startswith(category):
|
||
return category
|
||
# 旧式"删除/保留"回答兼容：含"删除"视为垃圾，其余一律保留。
|
||
if "删除" in content:
|
||
return CATEGORY_GARBAGE
|
||
return CATEGORY_DIALOGUE
|
||
|
||
|
||
def invoke(request: InvokeRequest) -> InvokeResponse:
|
||
"""过滤 SRT 中多余/无意义的字幕，产物为 filtered.srt。
|
||
|
||
流程：规则层（确定性删除）→ LLM 层（去重后按唯一文本多维分类）→
|
||
长文本保护 → 保留条重新编号输出。
|
||
"""
|
||
srt_uri = request.inputs.get("srt_uri")
|
||
if not srt_uri:
|
||
return InvokeResponse(status="failed", error="srt_uri is required")
|
||
srt_path = Path(srt_uri)
|
||
if not srt_path.is_file():
|
||
return InvokeResponse(status="failed", error="srt file not found")
|
||
|
||
entries = parse_srt(srt_path.read_text(encoding="utf-8"))
|
||
context_size = int(request.params.get("context_size", DEFAULT_CONTEXT_SIZE))
|
||
min_keep_len = int(request.params.get("min_keep_len", DEFAULT_MIN_KEEP_LEN))
|
||
# 水印 token 可经参数覆盖（JSON 数组字符串或列表），默认内置常见覆盖层词。
|
||
raw_tokens = request.params.get("overlay_tokens")
|
||
if isinstance(raw_tokens, str) and raw_tokens.strip():
|
||
raw_tokens = json.loads(raw_tokens)
|
||
overlay_tokens = {
|
||
str(t).casefold() for t in (raw_tokens or DEFAULT_OVERLAY_TOKENS)
|
||
}
|
||
# 去重开关：默认开；关掉时每个条目独立调用 LLM（不省调用，判定各自独立）。
|
||
dedupe = str(request.params.get("dedupe", "1")) not in ("0", "false", "False")
|
||
|
||
# 阶段 1：确定性规则层（不调 LLM）。
|
||
rule_verdicts = [_rule_verdict(entry["text"], overlay_tokens) for entry in entries]
|
||
llm_needed = [i for i, verdict in enumerate(rule_verdicts) if verdict is None]
|
||
|
||
# 阶段 2：LLM 分类层（去重：相同文本只判一次，上下文取首次出现）。
|
||
cat_by_index: dict[int, str] = {}
|
||
if llm_needed:
|
||
if dedupe:
|
||
first_of_key: dict[str, int] = {}
|
||
pool_indices: list[int] = []
|
||
for i in llm_needed:
|
||
key = _dedup_key(entries[i]["text"])
|
||
if key not in first_of_key:
|
||
first_of_key[key] = i
|
||
pool_indices.append(i)
|
||
else:
|
||
pool_indices = llm_needed
|
||
|
||
# 单条判断的工作函数：返回类别词。
|
||
def judge_one(index: int) -> str:
|
||
return _judge_category(entries, index, context_size, request.params)
|
||
|
||
# 进度日志：打印已判定条数、总数与平均处理速度（条/s）。
|
||
def log_progress(done: int, total: int, rate: float) -> None:
|
||
logger.info("字幕判定进度 %d/%d 条 (%.1f 条/s)", done, total, rate)
|
||
|
||
# 自适应并发调用 LLM：按实测负载弹性伸缩，避免压垮 LLM 接口。
|
||
pool = AdaptiveThreadPool(
|
||
worker=judge_one,
|
||
on_progress=log_progress,
|
||
min_workers=int(request.params.get("pool_min_workers", 1)),
|
||
max_workers=int(request.params.get("pool_max_workers", 16)),
|
||
window_seconds=float(request.params.get("pool_window_seconds", 10.0)),
|
||
fast_threshold=float(request.params.get("pool_fast_threshold", 0.3)),
|
||
slow_threshold=float(request.params.get("pool_slow_threshold", 1.0)),
|
||
)
|
||
categories = pool.map(pool_indices)
|
||
|
||
for index, category in zip(pool_indices, categories):
|
||
# 并行下 LLM 异常被线程池隔离为异常结果：任一条失败即整体失败，
|
||
# 避免静默输出未过滤结果。
|
||
if isinstance(category, Exception):
|
||
return InvokeResponse(status="failed", error=str(category))
|
||
cat_by_index[index] = category
|
||
if dedupe:
|
||
# 去重填充：与首次出现同键的条目复用同一类别，保证判定一致。
|
||
for i in llm_needed:
|
||
if i not in cat_by_index:
|
||
cat_by_index[i] = cat_by_index[first_of_key[_dedup_key(entries[i]["text"])]]
|
||
|
||
# 阶段 3：合并规则与 LLM 判定，应用长文本保护并输出。
|
||
kept: list[dict] = []
|
||
removed = 0
|
||
for i, (entry, rule_verdict) in enumerate(zip(entries, rule_verdicts)):
|
||
# 规则层命中即删（True）；未命中则按 LLM 类别与长文本保护判定。
|
||
if rule_verdict is True:
|
||
removed += 1
|
||
continue
|
||
if _should_delete(cat_by_index[i], entry["text"], min_keep_len):
|
||
removed += 1
|
||
logger.info("删除无意义字幕 %d: %r", i + 1, entry["text"][:40])
|
||
else:
|
||
kept.append(entry)
|
||
|
||
output_dir = Path(request.output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
output_path = output_dir / "filtered.srt"
|
||
output_path.write_text(serialize_srt(kept), encoding="utf-8")
|
||
logger.info("字幕过滤完成: 保留 %d 条, 删除 %d 条", len(kept), removed)
|
||
return InvokeResponse(
|
||
status="completed",
|
||
outputs={"srt_uri": str(output_path), "kept": len(kept), "removed": removed},
|
||
)
|