File size: 4,000 Bytes
c1a46f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """
解码策略模块 — 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:
"""
贪心解码。
TODO [Person D]: 实现以下逻辑:
1. encoder_output = model.encode(src_ids, src_padding_mask)
2. 初始化 decoder input: [B, 1] 全为 bos_id
3. for step in range(max_len):
a. logits = model.decode_step(decoder_input, encoder_output, src_padding_mask)
b. next_token = logits.argmax(dim=-1)
c. decoder_input = concat(decoder_input, next_token)
d. 如果所有序列都生成了 eos_id，则提前终止
4. 返回生成的 token ids [B, T]
"""
raise NotImplementedError("TODO: Person D 实现 greedy_decode")
@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:
"""
束搜索解码。
TODO [Person D]: 实现以下逻辑:
1. encoder_output = model.encode(src_ids, src_padding_mask)
2. 将 encoder_output 扩展为 beam_size 份: [B*beam, S, D]
3. 初始化 beam:
- beam_scores: [B, beam_size] 初始为 0
- beam_tokens: [B, beam_size, 1] 初始为 bos_id
4. for step in range(max_len):
a. 对每个 beam 计算 logits
b. log_probs = log_softmax(logits)
c. (可选) 应用 no_repeat_ngram 约束
d. scores = beam_scores + log_probs
e. 选择 top-k candidates (k = beam_size)
f. 更新 beam_tokens 和 beam_scores
g. 将已完成的 beam 移到 finished pool
5. 对 finished beams 应用 length_penalty:
score = score / (length ^ length_penalty)
6. 选择得分最高的序列
7. 返回最佳翻译 [B, T]
这是翻译任务最关键的解码算法，请仔细实现。
"""
raise NotImplementedError("TODO: Person D 实现 beam_search_decode")
@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)。
TODO [Person D]: 实现以下逻辑:
1. 与贪心解码类似，但每步采样而非取 argmax
2. 应用 temperature: logits = logits / temperature
3. 应用 top-k: 只保留概率最高的 k 个 token
4. 应用 top-p (nucleus): 只保留累积概率达到 p 的 token
5. 从过滤后的分布中采样: torch.multinomial
"""
raise NotImplementedError("TODO: Person D 实现 sample_decode")
def _apply_no_repeat_ngram(
logits: torch.Tensor,
generated_tokens: torch.Tensor,
ngram_size: int,
) -> torch.Tensor:
"""
防止生成重复的 n-gram。
TODO [Person D]:
1. 从 generated_tokens 中提取所有已出现的 (ngram_size-1)-gram
2. 对于每个可能导致重复 ngram 的 next token，将其 logits 设为 -inf
"""
raise NotImplementedError("TODO: Person D 实现 _apply_no_repeat_ngram")
|