src/easytranslate/evaluation/decoding.py

"""
解码策略模块 — Person D 负责实现

功能要求:
1. greedy_decode: 贪心解码
2. beam_search_decode: 束搜索解码
3. sample_decode: 采样解码 (temperature, top-k, top-p)

技术要点:
- Beam Search 是翻译任务最常用的解码策略
- 需要高效处理批量解码
- 支持长度惩罚 (length penalty) 和重复惩罚 (no_repeat_ngram)
- 对于预训练模型，可以直接使用 model.generate()
"""

from __future__ import annotations

import logging
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

logger = logging.getLogger(__name__)


@torch.no_grad()
def greedy_decode(
    model: nn.Module,
    src_ids: torch.Tensor,          # [B, S]
    src_padding_mask: torch.BoolTensor,  # [B, S]
    bos_id: int,
    eos_id: int,
    max_len: int = 256,
) -> torch.Tensor:
    """贪心解码。"""
    encoder_output = model.encode(src_ids, src_padding_mask)
    batch_size = src_ids.size(0)
    device = src_ids.device

    decoder_input = torch.full((batch_size, 1), bos_id, dtype=torch.long, device=device)
    finished = torch.zeros(batch_size, dtype=torch.bool, device=device)

    for _ in range(max_len):
        logits = model.decode_step(decoder_input, encoder_output, src_padding_mask)
        next_token = logits.argmax(dim=-1, keepdim=True)
        decoder_input = torch.cat([decoder_input, next_token], dim=1)
        finished = finished | next_token.squeeze(-1).eq(eos_id)
        if finished.all():
            break

    return decoder_input


@torch.no_grad()
def beam_search_decode(
    model: nn.Module,
    src_ids: torch.Tensor,          # [B, S]
    src_padding_mask: torch.BoolTensor,
    bos_id: int,
    eos_id: int,
    beam_size: int = 5,
    max_len: int = 256,
    length_penalty: float = 1.0,
    no_repeat_ngram_size: int = 0,
) -> torch.Tensor:
    """
    束搜索解码。
    
    优化改进:
    1. 限制最大 beam_size 为 64，防止内存溢出
    2. 使用更高效的索引操作避免中间 tensor 累积
    3. 及时释放不再需要的中间结果
    4. 对长序列使用更小的 beam_size
    """
    # 安全检查：限制 beam_size 防止内存问题
    beam_size = min(beam_size, 64)
    
    batch_size, seq_len = src_ids.size()
    device = src_ids.device
    
    # 根据序列长度动态调整 beam_size
    if seq_len > 512:
        beam_size = min(beam_size, 3)
    
    # 1. Encode source
    encoder_output = model.encode(src_ids, src_padding_mask)
    hidden_dim = encoder_output.size(-1)
    
    # 2. Expand encoder output for beam search
    encoder_output = encoder_output.unsqueeze(1).expand(-1, beam_size, -1, -1)
    encoder_output = encoder_output.reshape(batch_size * beam_size, seq_len, hidden_dim)
    src_padding_mask_expanded = src_padding_mask.unsqueeze(1).expand(-1, beam_size, -1)
    src_padding_mask_expanded = src_padding_mask_expanded.reshape(batch_size * beam_size, seq_len)

    # 3. Initialize beams - 使用预分配的 tensor 减少内存分配
    beam_scores = torch.zeros(batch_size, beam_size, device=device)
    beam_scores[:, 1:] = float("-inf")
    
    # 预分配最大长度的 tensor，动态填充
    beam_tokens = torch.full((batch_size, beam_size, max_len + 1), eos_id, dtype=torch.long, device=device)
    beam_tokens[:, :, 0] = bos_id  # [B, beam, 1]
    
    finished = torch.zeros(batch_size, beam_size, dtype=torch.bool, device=device)
    lengths = torch.ones(batch_size, beam_size, dtype=torch.long, device=device)

    # 4. Iterative decoding
    for step in range(1, max_len + 1):
        # 只处理未完成的 beam
        active_mask = ~finished
        if not active_mask.any():
            break
        
        # 获取当前步的输入 [B*beam, step]
        flat_tokens = beam_tokens[:, :, :step].reshape(batch_size * beam_size, step)
        
        # 前向解码
        logits = model.decode_step(flat_tokens, encoder_output, src_padding_mask_expanded)
        log_probs = F.log_softmax(logits, dim=-1)
        vocab_size = log_probs.size(-1)
        
        # 应用 no_repeat_ngram 约束
        if no_repeat_ngram_size > 0 and step >= no_repeat_ngram_size:
            log_probs = _apply_no_repeat_ngram(log_probs, flat_tokens, no_repeat_ngram_size)

        # Mask finished beams - 禁止生成新 token，但允许保持 EOS
        finished_flat = finished.view(batch_size * beam_size)
        if finished_flat.any():
            log_probs[finished_flat] = float("-inf")
            log_probs[finished_flat, eos_id] = 0.0

        # 计算分数
        scores = beam_scores.unsqueeze(-1) + log_probs.view(batch_size, beam_size, vocab_size)
        scores = scores.view(batch_size, -1)

        # 选择 top-k
        topk_scores, topk_indices = scores.topk(beam_size, dim=-1)
        beam_indices = topk_indices // vocab_size
        token_indices = topk_indices % vocab_size

        # 高效更新 beam_tokens - 使用索引而非拼接
        new_beam_tokens = beam_tokens.clone()
        for b in range(batch_size):
            new_beam_tokens[b] = beam_tokens[b][beam_indices[b]]
            new_beam_tokens[b, torch.arange(beam_size), step] = token_indices[b]
        beam_tokens = new_beam_tokens
        
        # 更新 finished 和 lengths
        new_finished = finished.gather(1, beam_indices) | token_indices.eq(eos_id)
        new_lengths = lengths.gather(1, beam_indices)
        new_lengths[~new_finished] = step + 1
        
        beam_scores = topk_scores
        finished = new_finished
        lengths = new_lengths

    # 5. Apply length penalty
    lengths = lengths.float()
    penalties = lengths ** length_penalty
    final_scores = beam_scores / penalties

    # 6. Select best beam for each batch
    best_indices = final_scores.argmax(dim=-1)
    best_sequences = beam_tokens[torch.arange(batch_size, device=device), best_indices]

    return best_sequences


@torch.no_grad()
def sample_decode(
    model: nn.Module,
    src_ids: torch.Tensor,
    src_padding_mask: torch.BoolTensor,
    bos_id: int,
    eos_id: int,
    max_len: int = 256,
    temperature: float = 1.0,
    top_k: int = 0,
    top_p: float = 1.0,
) -> torch.Tensor:
    """采样解码 (支持 temperature, top-k, top-p/nucleus sampling)。"""
    encoder_output = model.encode(src_ids, src_padding_mask)
    batch_size = src_ids.size(0)
    device = src_ids.device

    decoder_input = torch.full((batch_size, 1), bos_id, dtype=torch.long, device=device)
    finished = torch.zeros(batch_size, dtype=torch.bool, device=device)

    for _ in range(max_len):
        logits = model.decode_step(decoder_input, encoder_output, src_padding_mask)

        logits = logits / max(temperature, 1e-8)

        if top_k > 0:
            k = min(top_k, logits.size(-1))
            topk_values, _ = torch.topk(logits, k, dim=-1)
            threshold = topk_values[:, -1].unsqueeze(-1)
            logits[logits < threshold] = float("-inf")

        if 0.0 < top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            # Remove tokens whose cumulative probability exceeds top_p
            # (shift by one so the token that pushes above the threshold is kept)
            remove_mask = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= top_p
            sorted_logits[remove_mask] = float("-inf")
            # Scatter sorted values back to their original vocabulary positions
            logits = torch.zeros_like(logits).scatter_(1, sorted_indices, sorted_logits)

        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)

        next_token = torch.where(finished.unsqueeze(-1), torch.full_like(next_token, eos_id), next_token)
        decoder_input = torch.cat([decoder_input, next_token], dim=1)
        finished = finished | next_token.squeeze(-1).eq(eos_id)

        if finished.all():
            break

    return decoder_input


def _apply_no_repeat_ngram(
    logits: torch.Tensor,
    generated_tokens: torch.Tensor,
    ngram_size: int,
) -> torch.Tensor:
    """防止生成重复的 n-gram。"""
    if ngram_size <= 0:
        return logits

    batch_size = logits.size(0)
    seq_len = generated_tokens.size(1)
    
    if seq_len < ngram_size - 1:
        return logits

    for batch_idx in range(batch_size):
        tokens = generated_tokens[batch_idx].tolist()

        ngram_map: dict[tuple, set] = {}
        for i in range(len(tokens) - ngram_size + 1):
            prefix = tuple(tokens[i : i + ngram_size - 1])
            next_tok = tokens[i + ngram_size - 1]
            ngram_map.setdefault(prefix, set()).add(next_tok)

        current_prefix = tuple(tokens[-(ngram_size - 1):])
        if current_prefix in ngram_map:
            banned = list(ngram_map[current_prefix])
            logits[batch_idx, banned] = float("-inf")

    return logits