vrsub/nodes/subtitle_cleanup.py

"""Subtitle cleanup: remove long-duration closing/greeting hallucinations.

Background (real run 20260905115050): after fixing the timing alignment,
subtitles still contain "closing/greeting hallucination words" - fixed
phrases like 'wan an / gan xie guan kan / gan xie nin de guan kan'
(good night / thanks for watching) that the ASR/LLM repeatedly emits on
empty segments, filling a full 30s block, unrelated to video content.
Some 2s 'good night' might be real dialogue, so it must be kept.

处理策略（2026-09 用户确认改为"整条剔除"）：
幻觉识别出后（展示时长 ≥ 阈值 且 文本命中套话词表），应**连带时间戳把整条
字幕 cue 删除**（剩余条目重新编号），而不是把文本替换成 '-' 占位留给下游——
占位会一路流到 ASS 渲染成可见的"减号"、在翻译/过滤阶段都要额外特殊处理，
处理位置绕且不彻底。因此清洗统一为：**在幻觉产生处（whisper 转录后 / LLM
翻译后）整条删除**，时间轴随之消失，字幕序号重新连续编号。

两类词表：
- HALLUCINATION_TOKENS：中文（LLM 翻译产物中的寒暄，如"晚安/感谢观看"）；
- JAPANESE_HALLUCINATION_TOKENS：日文（whisper decode_full 无 VAD 解码在
  无语音段直接输出的套话，如"おやすみなさい/ご視聴ありがとうございました"）。

阈值从真实运行实测数据判定：30s 幻觉占位 vs 2s 真实词，分界明显，
默认 threshold=15s（≥15s 才删除；<15s 的相同词可能是剧情真实道晚安，保留）。

Pure functions, unit-testable (tests/test_hallucination_mask.py).
"""

from __future__ import annotations

import re

# Greeting/closing hallucination tokens that LLM repeats on empty/end segments.
HALLUCINATION_TOKENS = (
    "谢谢观看", "感谢观看", "感谢收看", "谢谢收看", "感谢您的观看", "感谢您的收看",
    "晚安", "下次再见", "再会", "敬请期待", "感谢您的光临", "欢迎光临",
    "再见", "多谢观看", "观看愉快",
)

# Display-duration threshold (seconds): only remove entries longer than this.
DEFAULT_THRESHOLD_SECONDS = 15.0

# 日文 ASR 直出（whisper 节点 decode_full 无 VAD 整段解码）的收尾/寒暄幻觉词。
# 无 VAD 解码会把无语音/音乐/呻吟段当语音，whisper 常在这些段重复输出套话
# （实测 savr-1054 全片 119-149s/600-630s 等出现 30s 长"おやすみなさい"
# "ご視聴ありがとうございました"）。短时（≤阈值）的相同词可能是剧情里真实
# 互道晚安，须保留；仅删除展示时长 ≥ 阈值的条目（与中文表同一机制）。
JAPANESE_HALLUCINATION_TOKENS = (
    "おやすみなさい",        # 晚安
    "ご視聴ありがとうございました",  # 感谢观看
    "ありがとうございました",        # 感谢
    "ご視聴ありがとうございます",    # 感谢观看(现在时)
    "また見てね",            # 下次再见
    "お楽しみに",            # 敬请期待
    "チャンネル登録よろしくお願いします",  # 求订阅
    "さようなら",            # 再见
    "音楽",                  # 音乐(自述)
    "Goodbye",
    "Thank you for watching",
)

# 解析 SRT：每个 cue 由 序号行 + 时间轴行 + 文本行(可能多行) + 空行 组成。
# 采用逐行解析（不依赖可能粘连的跨 cue 正则），兼容文本多行。
_TS_RE = re.compile(r"^(\d{2}:\d{2}:\d{2},\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2},\d{3})\s*$")


def _ts_to_seconds(ts: str) -> float:
    """Convert an SRT timestamp HH:MM:SS,mmm to seconds (float)."""
    hours, minutes, rest = ts.split(":")
    seconds, millis = rest.split(",")
    return int(hours) * 3600 + int(minutes) * 60 + int(seconds) + int(millis) / 1000


def remove_hallucination_entries(
    srt_text: str,
    tokens: tuple[str, ...],
    threshold_seconds: float = DEFAULT_THRESHOLD_SECONDS,
) -> str:
    """删除 SRT 中展示时长 ≥ 阈值且文本命中 tokens 的**整条 cue**（连带时间戳）。

    被删除的 cue 不再输出（序号、时间轴、文本全部消失），剩余 cue 按原顺序
    重新从 1 编号，保证产物是合法连续的 SRT。输入不被修改（纯函数）。

    用途：幻觉在产生处直接剔除——whisper decode_full（日语词表）与 LLM 翻译后
    （中文词表）均调用本函数，避免 '-' 占位一路流到 ASS 渲染成可见减号。
    """
    lines = srt_text.splitlines()
    kept: list[str] = []
    number = 1
    index = 0
    while index < len(lines):
        line = lines[index]
        if line.strip() and line.strip().isdigit() and index + 1 < len(lines):
            ts_match = _TS_RE.match(lines[index + 1].strip())
            if ts_match:
                # 收集本 cue 文本：时间轴后直到空行前的所有非空行（可多行）。
                text_lines: list[str] = []
                cursor = index + 2
                while cursor < len(lines) and lines[cursor].strip():
                    text_lines.append(lines[cursor].strip())
                    cursor += 1
                text = "\n".join(text_lines)
                start_sec = _ts_to_seconds(ts_match.group(1))
                end_sec = _ts_to_seconds(ts_match.group(2))
                duration = end_sec - start_sec
                is_hallucination = duration >= threshold_seconds and any(
                    t in text for t in tokens
                )
                if not is_hallucination:
                    # 非幻觉：输出 新序号+时间轴+文本+空行（重建标准 SRT）。
                    kept.append(
                        f"{number}\n{lines[index + 1].strip()}\n{text}\n"
                    )
                    number += 1
                # 幻觉 cue：整条跳过（序号/时间轴/文本都不输出）。
                index = cursor
                continue
        # 非 cue 行（文件头/尾部噪声）跳过，避免序号/空行残留。
        index += 1
    return "\n".join(kept).rstrip() + "\n"


def clean_japanese_hallucinations(
    srt_text: str,
    threshold_seconds: float = DEFAULT_THRESHOLD_SECONDS,
) -> str:
    """从 SRT 中整条删除日文收尾/寒暄长时幻觉（ASR 无 VAD 解码兜底）。

    whisper 节点 decode_full=true（无 VAD 整段解码）在无语音段会输出长时
    套话占位；本函数把展示时长 ≥ 阈值且文本含 JAPANESE_HALLUCINATION_TOKENS
    的**整条 cue 连带时间戳删除**、剩余重编号。短时相同词（剧情真实道晚安）
    保留。纯函数，不修改输入。
    """
    return remove_hallucination_entries(srt_text, JAPANESE_HALLUCINATION_TOKENS, threshold_seconds)


def clean_srt_text(
    srt_text: str,
    threshold_seconds: float = DEFAULT_THRESHOLD_SECONDS,
) -> str:
    """从 SRT 中整条删除中文长时寒暄幻觉（LLM 翻译产物清洗）。

    对展示时长 ≥ 阈值且文本含 HALLUCINATION_TOKENS 的 cue 连带时间戳整条
    删除、剩余重编号；短时相同词（剧情真实互道晚安）保留。与
    clean_japanese_hallucinations 同机制，词表不同。纯函数，不修改输入。
    """
    return remove_hallucination_entries(srt_text, HALLUCINATION_TOKENS, threshold_seconds)