audio2text/app/services/translate_service.py

"""翻译服务：NLLB-200，英译中。

通过 model_manager 加载，确保 ASR 已卸载、翻译器独占显存，从而可用大 batch_size。

GPU 优化：按长度排序后分批翻译。
- 同一批内句子长度相近 → padding 浪费最小化 → GPU 有效计算占比提升
- 翻译完按原始下标散回，保证 zh_texts[i] 对应 subs[i]（时间戳对齐不变）
- 批切分用 token 预算 + 条数上限双重约束：短句自动攒大批，长句自动拆小批

单条翻译失败时该位置回退为原英文。
"""

from __future__ import annotations

import logging
import os
from typing import Callable

from ..config import get_settings
from .model_manager import get_model_manager
from .types import Subtitle

logger = logging.getLogger("audio2text.translate")

# 估算每条字幕的 token 数：英文约 1 token/词，留 20% 余量覆盖标点/子词拆分
_TOKENS_PER_WORD = 1.2


def translate(
    subtitles: list[Subtitle],
    on_progress: Callable[[int, int], None] | None = None,
) -> list[str]:
    """批量翻译英文字幕为中文。

    Args:
        subtitles: 断句后的英文字幕条目（按时间顺序）
        on_progress: 可选进度回调 (done_count, total_count)，每批完成时调一次。

    Returns:
        list[str]，与 subtitles 等长、顺序对应的中文译文。
        单条翻译失败时该位置回退为原英文。
    """
    if not subtitles:
        return []

    s = get_settings().translation
    pipe = get_model_manager().get_translator()
    batch_size = s.batch_size
    max_len = s.max_length

    # 取纯文本（去掉折行），避免翻译把换行符当语义
    texts = [sub.text.replace("\n", " ").strip() for sub in subtitles]
    sort_by_length = s.sort_by_length

    # 环境变量覆盖：便于 A/B 基准对比（test/bench_translate.py 用）
    if os.environ.get("TRANSLATE_NO_SORT") == "1":
        sort_by_length = False

    if sort_by_length:
        results = _translate_sorted(pipe, texts, batch_size, max_len, on_progress)
    else:
        results = _translate_sequential(pipe, texts, batch_size, max_len, on_progress)

    logger.debug("翻译完成：%d 条。", len(results))
    return results


# ---------------- 长度排序批处理（默认）----------------

def _translate_sorted(
    pipe, texts: list[str], batch_size: int, max_len: int,
    on_progress: Callable[[int, int], None] | None = None,
) -> list[str]:
    """按长度排序后分批翻译，翻译完按原序散回。

    1. 记录 (orig_idx, text, est_tokens)
    2. 按 est_tokens 升序排序 → 相近长度的聚到同一批
    3. token 预算 + 条数上限双重约束切批：短句攒大批，长句拆小批
    4. 逐批翻译，按 orig_idx 把译文放回 results[orig_idx]
    """
    n = len(texts)
    # 估算每条 token 数（用词数 × 1.2，至少 1 避免除零）
    items = [
        (i, texts[i], max(1, int(len(texts[i].split()) * _TOKENS_PER_WORD)))
        for i in range(n)
    ]
    # 按 token 长度升序：短句在前，长句在后
    items.sort(key=lambda x: x[2])

    # token 预算上限：一批的总 token 不超过 batch_size * max_len
    # 短句（每条 ~10 token）可攒到 batch_size 条；长句（~200 token）自动拆成更小批
    token_budget = batch_size * max_len

    batches: list[list[tuple[int, str, int]]] = []  # [(orig_idx, text, tok), ...]
    cur_batch: list[tuple[int, str, int]] = []
    cur_max = 0  # 当前批内最长句的 token 数

    for orig_idx, text, tok in items:
        new_max = max(cur_max, tok)
        new_tokens = (len(cur_batch) + 1) * new_max  # 批内所有句都 pad 到 new_max
        if cur_batch and (len(cur_batch) >= batch_size or new_tokens > token_budget):
            batches.append(cur_batch)
            cur_batch = []
            cur_max = 0
            new_max = tok
        cur_batch.append((orig_idx, text, tok))
        cur_max = new_max
    if cur_batch:
        batches.append(cur_batch)

    # padding 浪费对比（DEBUG 日志量化收益）
    pad_sorted = sum(len(b) * max(t for _, _, t in b) - sum(t for _, _, t in b) for b in batches)
    pad_seq = _estimate_sequential_padding(texts, batch_size)
    saving = (1 - pad_sorted / pad_seq) * 100 if pad_seq else 0
    logger.debug(
        "翻译分批：%d 条 → %d 批（长度排序）。padding 浪费：顺序 %d → 排序 %d token（节省 %.0f%%）",
        n, len(batches), pad_seq, pad_sorted, saving,
    )

    results: list[str | None] = [None] * n
    done = 0
    for batch in batches:
        orig_indices = [b[0] for b in batch]
        batch_texts = [b[1] for b in batch]
        translated = _translate_batch(pipe, batch_texts, max_len)
        for idx, zh in zip(orig_indices, translated):
            results[idx] = zh
        done += len(batch)
        if on_progress is not None:
            on_progress(done, n)
        elif (done // batch_size + 1) % 5 == 0:
            logger.debug("已翻译 %d/%d 条。", done, n)

    # None（理论不会发生，_translate_batch 保证返回等长）→ 回退原文
    return [results[i] or texts[i] for i in range(n)]


def _estimate_sequential_padding(texts: list[str], batch_size: int) -> int:
    """估算按原序分批的 padding 浪费（token 数）。"""
    total = 0
    for i in range(0, len(texts), batch_size):
        chunk = texts[i:i + batch_size]
        toks = [max(1, int(len(t.split()) * _TOKENS_PER_WORD)) for t in chunk]
        batch_max = max(toks)
        total += batch_max * len(chunk) - sum(toks)
    return total


# ---------------- 顺序批处理（A/B 对比用 / sort_by_length=false）----------------

def _translate_sequential(
    pipe, texts: list[str], batch_size: int, max_len: int,
    on_progress: Callable[[int, int], None] | None = None,
) -> list[str]:
    """按原序分批翻译（旧行为，便于 A/B 对比）。"""
    results: list[str] = []
    n = len(texts)
    for i in range(0, n, batch_size):
        chunk = texts[i:i + batch_size]
        translated = _translate_batch(pipe, chunk, max_len)
        results.extend(translated)
        done = min(i + len(chunk), n)
        if on_progress is not None:
            on_progress(done, n)
        elif (i // batch_size + 1) % 5 == 0:
            logger.debug("已翻译 %d/%d 条。", done, n)
    return results


# ---------------- 单批翻译 + 逐条重试回退 ----------------

def _translate_batch(pipe, chunk: list[str], max_len: int) -> list[str]:
    """翻译一个批次，失败时降级到逐条重试。

    Args:
        chunk: 本批的文本列表
        max_len: 单条最大生成长度
    """
    try:
        out = pipe(chunk, max_length=max_len, truncation=True)
        return [item.get("translation_text", "").strip() for item in out]
    except Exception as exc:  # pragma: no cover
        logger.warning("批次翻译失败（%d 条），逐条重试：%s", len(chunk), exc)
        results: list[str] = []
        for t in chunk:
            try:
                out = pipe([t], max_length=max_len, truncation=True)
                results.append(out[0].get("translation_text", "").strip())
            except Exception as exc2:
                logger.warning("单条翻译失败，回退原文：%s", exc2)
                results.append(t)  # 回退原文
        return results