vrsub/scripts/build_segment_eval.py

"""从最终中文字幕反查各类"片段窗口"，供人工挑选评测用例。

背景与目的
----------
评测翻译模型需要"原文 + 该原文对应的真实含义"的成对基准。本仓库同一部视频
（CJOD-255，2 小时）有两条真实产物：

- 日语原文：run_d386ccf124f7（learn-translate）的 whisper ASR，
  1161 条 cue（约每 1-4 句一条）；
- 中文基准：run_479b411f299d（ocr-subtitle）过滤后的烧录字幕 1144 条
  ——烧录字幕是画面的一部分，也可视为官方/人工字幕，但它**是一条条按屏幕
  显示合并的整行**（约 2-8 秒一条），与 whisper 的 cue 切分不同。

因此不能逐条对照。正确口径是**窗口级**：取中文基准一条字幕的时间窗，收集该窗
内覆盖的日语 cue（1-3 条），窗口即为一个评测用例——比较"模型对这几条 cue 的
译文合起来"是否表达了该中文行。本脚本把这类窗口导出，供人工确认与挑选。

用法：
    uv run python scripts/build_segment_eval.py --dump-window 1150  # 看某处的窗口细节
    uv run python scripts/build_segment_eval.py --export scripts/data/translate_eval/windows.jsonl
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent

# 两条真实产物（相对仓库根）。
DEFAULT_JA = PROJECT_ROOT / "data/storage/runs/run_d386ccf124f7/steps/asr/transcript.srt"
DEFAULT_ZH = PROJECT_ROOT / "data/storage/runs/run_479b411f299d/steps/filter/filtered.srt"


def seconds(ts: str) -> float:
    """SRT 时间戳 -> 秒。"""
    hours, minutes, rest = ts.split(":")
    secs, millis = rest.split(",")
    return int(hours) * 3600 + int(minutes) * 60 + int(secs) + int(millis) / 1000


def build_windows(ja_path: Path, zh_path: Path, max_ja: int = 3, max_duration: float = 15.0) -> list[dict]:
    """返回可用评测窗口列表。

    窗口 = 一条中文基准字幕 + 其时间窗内的日语 cue 组（1..max_ja 条）。
    过滤条件：窗口时长 <= max_duration（避免整段合并的长窗口无法逐句核对）、
    cue 数在 1..max_ja 之间（0 条说明 OCR 行无对应日语，>max_ja 说明配对太碎）。
    """
    from nodes.srt import parse_srt

    ja_cues = parse_srt(ja_path.read_text(encoding="utf-8"))
    zh_cues = parse_srt(zh_path.read_text(encoding="utf-8"))

    windows: list[dict] = []
    for zh in zh_cues:
        start, end = seconds(zh.start), seconds(zh.end)
        group = [c for c in ja_cues if seconds(c.start) < end and start < seconds(c.end)]
        if not 1 <= len(group) <= max_ja or (end - start) > max_duration:
            continue
        windows.append(
            {
                "zh_start": zh.start,
                "zh_end": zh.end,
                "zh_ref": zh.text,
                "ja_start": group[0].start,
                "ja_end": group[-1].end,
                "ja_lines": [c.text for c in group],
                "ja_ids": [ja_cues.index(c) + 1 for c in group],
            }
        )
    return windows


def main() -> None:
    parser = argparse.ArgumentParser(description="导出窗口级翻译评测候选")
    parser.add_argument("--ja", type=Path, default=DEFAULT_JA)
    parser.add_argument("--zh", type=Path, default=DEFAULT_ZH)
    parser.add_argument("--export", type=Path, default=None, help="写出 windows.jsonl")
    parser.add_argument("--dump-window", type=float, default=None, help="打印指定秒附近的窗口")
    parser.add_argument("--limit", type=int, default=0)
    args = parser.parse_args()

    windows = build_windows(args.ja, args.zh)
    print(f"可用窗口 {len(windows)} 个（1-3 条日语 cue 且窗口 <=15s）")

    if args.export:
        args.export.parent.mkdir(parents=True, exist_ok=True)
        with args.export.open("w", encoding="utf-8") as fh:
            for index, item in enumerate(windows, 1):
                fh.write(json.dumps({"id": index, **item}, ensure_ascii=False) + "\n")
        print(f"已写入 {args.export}")

    if args.dump_window is not None:
        for item in windows:
            if abs(seconds(item["zh_start"]) - args.dump_window) <= 12:
                print(json.dumps(item, ensure_ascii=False, indent=1))

    for item in (windows if args.limit <= 0 else windows[: args.limit]):
        joined = " | ".join(item["ja_lines"])
        print(f"{item['zh_start']} JA[{len(item['ja_lines'])}]: {joined!r} -> {item['zh_ref']!r}")


if __name__ == "__main__":
    main()