机器翻译：从 Seq2Seq 到 Transformer

# BLEU 分数计算详解
from collections import Counter
import math

def compute_bleu(candidate, references, max_n=4):
    """手动实现 BLEU 分数计算"""
    # 计算各阶 n-gram 修正精度
    precisions = []
    for n in range(1, max_n + 1):
        cand_ngrams = Counter()
        ref_ngram_counts = [Counter() for _ in references]
        
        # 提取候选 n-gram
        for i in range(len(candidate) - n + 1):
            ngram = tuple(candidate[i:i + n])
            cand_ngrams[ngram] += 1
        
        # 提取参考 n-gram
        for ref_idx, ref in enumerate(references):
            for i in range(len(ref) - n + 1):
                ngram = tuple(ref[i:i + n])
                ref_ngram_counts[ref_idx][ngram] += 1
        
        # 修正精度：裁剪候选 n-gram 计数
        clipped_count = 0
        total_count = 0
        for ngram, count in cand_ngrams.items():
            max_ref_count = max(ref_counts.get(ngram, 0)
                                for ref_counts in ref_ngram_counts)
            clipped_count += min(count, max_ref_count)
            total_count += count
        
        precision = clipped_count / total_count if total_count > 0 else 0
        precisions.append(precision)
    
    # 简短惩罚 (Brevity Penalty)
    cand_len = len(candidate)
    ref_lens = [len(ref) for ref in references]
    closest_ref_len = min(ref_lens, key=lambda x: abs(x - cand_len))
    bp = 1.0 if cand_len > closest_ref_len else math.exp(1 - closest_ref_len / cand_len)
    
    # 几何平均
    if 0 in precisions:
        return 0.0
    geo_mean = math.exp(sum(math.log(p) for p in precisions) / max_n)
    return bp * geo_mean * 100

# 测试
candidate = ["the", "cat", "is", "on", "the", "mat"]
refs = [
    ["the", "cat", "is", "sitting", "on", "the", "mat"],
    ["there", "is", "a", "cat", "on", "the", "mat"],
]
bleu = compute_bleu(candidate, refs)
print(f"BLEU-4 分数: {bleu:.2f}")

# BLEU vs COMET 对比分析
import math

print("=== 自动评估指标对比 ===\n")

# 场景 1: 同义词替换
print("场景 1: 同义词替换")
candidate = ["the", "child", "is", "playing", "outside"]
ref1 = ["the", "kid", "is", "playing", "outdoors"]
ref2 = ["the", "child", "is", "playing", "outside"]
refs = [ref1, ref2]
print(f"  候选: {' '.join(candidate)}")
print(f"  参考1: {' '.join(ref1)}")
print(f"  参考2: {' '.join(ref2)}")
print(f"  BLEU-4 (vs ref1): {compute_bleu(candidate, [ref1]):.2f}")
print(f"  BLEU-4 (vs ref2): {compute_bleu(candidate, [ref2]):.2f}")
print(f"  BLEU-4 (双参考):  {compute_bleu(candidate, refs):.2f}")
print(f"  → BLEU 因 'child' vs 'kid' 的精确匹配差异而评分不同")

# 场景 2: 词序变化
print("\n场景 2: 词序变化")
c1 = ["i", "eat", "an", "apple"]
c2 = ["an", "apple", "i", "eat"]
ref = ["i", "eat", "an", "apple"]
print(f"  候选1: {' '.join(c1)} (BLEU: {compute_bleu(c1, [ref]):.2f})")
print(f"  候选2: {' '.join(c2)} (BLEU: {compute_bleu(c2, [ref]):.2f})")
print(f"  → BLEU 对词序变化极度敏感")

# BLEU 的优势与局限
print("\n=== BLEU 总结 ===")
print("优势: 计算快速、可重复、与人类判断有一定相关性")
print("局限: 忽略语义等价、对词序敏感、不考虑句法正确性")
print("替代方案: COMET (基于预训练模型), BERTScore (基于 BERT 语义相似度)")

import torch
import torch.nn as nn

class EncoderRNN(nn.Module):
    """RNN 编码器：将源句子编码为上下文向量"""
    
    def __init__(self, vocab_size, embed_dim, hidden_dim, n_layers=2, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.rnn = nn.LSTM(embed_dim, hidden_dim, n_layers,
                           dropout=dropout if n_layers > 1 else 0,
                           batch_first=True)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, src, src_lengths):
        """
        src: (batch, src_len) 源句子索引
        src_lengths: (batch,) 每个源句子的实际长度
        返回: (outputs, hidden, cell) 用于解码器
        """
        embedded = self.dropout(self.embedding(src))  # (batch, src_len, embed)
        
        # pack_padded_sequence 处理变长序列
        packed = nn.utils.rnn.pack_padded_sequence(
            embedded, src_lengths.cpu(), batch_first=True, enforce_sorted=False
        )
        outputs, (hidden, cell) = self.rnn(packed)
        outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs, batch_first=True)
        
        # outputs: (batch, src_len, hidden_dim)
        # hidden: (n_layers, batch, hidden_dim)
        # cell: (n_layers, batch, hidden_dim)
        return outputs, hidden, cell

import torch
import torch.nn as nn
import random

class DecoderRNN(nn.Module):
    """RNN 解码器：根据上下文向量生成目标句子"""
    
    def __init__(self, vocab_size, embed_dim, hidden_dim, n_layers=2, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.rnn = nn.LSTM(embed_dim, hidden_dim, n_layers,
                           dropout=dropout if n_layers > 1 else 0)
        self.fc_out = nn.Linear(hidden_dim, vocab_size)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, input_token, hidden, cell):
        """
        input_token: (batch,) 当前输入词索引
        hidden: (n_layers, batch, hidden_dim)
        cell: (n_layers, batch, hidden_dim)
        返回: output, hidden, cell
        """
        input_token = input_token.unsqueeze(1)  # (batch, 1)
        embedded = self.dropout(self.embedding(input_token))  # (batch, 1, embed)
        embedded = embedded.transpose(0, 1)  # (1, batch, embed)
        
        output, (hidden, cell) = self.rnn(embedded, (hidden, cell))
        prediction = self.fc_out(output.squeeze(0))  # (batch, vocab_size)
        return prediction, hidden, cell

class Seq2Seq(nn.Module):
    """完整的 Seq2Seq 翻译模型"""
    
    def __init__(self, encoder, decoder, device, teacher_forcing_ratio=0.5):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.device = device
        self.teacher_forcing_ratio = teacher_forcing_ratio
    
    def forward(self, src, src_lengths, trg):
        """训练时的前向传播"""
        batch_size = src.shape[0]
        trg_len = trg.shape[1]
        trg_vocab_size = self.decoder.fc_out.out_features
        
        outputs = torch.zeros(batch_size, trg_len, trg_vocab_size).to(self.device)
        encoder_outputs, hidden, cell = self.encoder(src, src_lengths)
        
        # 解码器输入从 <sos> 开始
        input_token = trg[:, 0]  # (batch,)
        
        for t in range(1, trg_len):
            output, hidden, cell = self.decoder(input_token, hidden, cell)
            outputs[:, t] = output
            # Teacher Forcing: 以一定概率使用真实词作为下一步输入
            top1 = output.argmax(1)
            teacher_force = random.random() < self.teacher_forcing_ratio
            input_token = trg[:, t] if teacher_force else top1
        
        return outputs

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

class BahdanauAttention(nn.Module):
    """Bahdanau 加性注意力"""
    
    def __init__(self, hidden_dim):
        super().__init__()
        self.attn = nn.Linear(hidden_dim * 2, hidden_dim)  # W_s · s + W_h · h
        self.v = nn.Linear(hidden_dim, 1, bias=False)  # 注意力权重向量
    
    def forward(self, hidden, encoder_outputs):
        """
        hidden: (batch, hidden_dim) 解码器上一时刻隐藏状态
        encoder_outputs: (batch, src_len, hidden_dim) 编码器输出
        返回: context (batch, hidden_dim), attn_weights (batch, src_len)
        """
        src_len = encoder_outputs.shape[1]
        
        # 重复 hidden 以匹配源序列长度
        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)  # (batch, src_len, hidden)
        
        # 计算对齐分数: e = v^T tanh(W·[s;h])
        energy = torch.tanh(self.attn(torch.cat([hidden, encoder_outputs], dim=2)))
        attention = self.v(energy).squeeze(2)  # (batch, src_len)
        
        # 注意力权重归一化
        attn_weights = F.softmax(attention, dim=1)  # (batch, src_len)
        
        # 上下文向量: 加权和
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs).squeeze(1)
        
        return context, attn_weights

import torch
import torch.nn as nn

class AttentionDecoderRNN(nn.Module):
    """带注意力的解码器"""
    
    def __init__(self, vocab_size, embed_dim, hidden_dim, n_layers=2, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.attention = BahdanauAttention(hidden_dim)
        # 拼接: [Embed(y) + 注意力上下文] → RNN 输入
        self.rnn = nn.LSTM(embed_dim + hidden_dim, hidden_dim, n_layers,
                           dropout=dropout if n_layers > 1 else 0)
        self.fc_out = nn.Linear(embed_dim + hidden_dim * 2, vocab_size)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, input_token, hidden, cell, encoder_outputs):
        """
        input_token: (batch,) 当前输入词
        hidden/cell: (n_layers, batch, hidden_dim)
        encoder_outputs: (batch, src_len, hidden_dim)
        """
        input_token = input_token.unsqueeze(1)  # (batch, 1)
        embedded = self.dropout(self.embedding(input_token))  # (batch, 1, embed)
        
        # 注意力计算
        context, attn_weights = self.attention(hidden[-1], encoder_outputs)
        # context: (batch, hidden_dim)
        # attn_weights: (batch, src_len)
        
        # 拼接嵌入和上下文向量
        rnn_input = torch.cat([embedded, context.unsqueeze(1)], dim=2)
        rnn_input = rnn_input.transpose(0, 1)  # (1, batch, embed+hidden)
        
        output, (hidden, cell) = self.rnn(rnn_input, (hidden, cell))
        
        # 拼接: [嵌入, 上下文, 解码器输出] → 预测
        output = torch.cat([embedded.squeeze(1), context, output.squeeze(0)], dim=1)
        prediction = self.fc_out(output)  # (batch, vocab_size)
        
        return prediction, hidden, cell, attn_weights

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    """多头自注意力机制"""
    
    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 每个头的维度
        
        # 线性投影矩阵
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(dropout)
        self.scale = math.sqrt(self.d_k)
    
    def forward(self, query, key, value, mask=None):
        """
        query/key/value: (batch, seq_len, d_model)
        mask: (batch, 1, 1, seq_len) 或 (batch, seq_len, seq_len)
        """
        batch_size = query.shape[0]
        
        # 线性投影并分头: (batch, n_heads, seq_len, d_k)
        Q = self.W_Q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        # Scaled Dot-Product Attention
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
        
        attn_probs = torch.softmax(attn_scores, dim=-1)
        attn_probs = self.dropout(attn_probs)
        
        # 加权求和
        context = torch.matmul(attn_probs, V)  # (batch, n_heads, seq_len, d_k)
        
        # 拼接头并线性变换
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        output = self.W_O(context)
        
        return output, attn_probs

import torch
import torch.nn as nn

class PositionalEncoding(nn.Module):
    """正弦位置编码"""
    
    def __init__(self, d_model=512, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                            (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数维度用 sin
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维度用 cos
        
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        """x: (batch, seq_len, d_model)"""
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

class TransformerBlock(nn.Module):
    """单个 Transformer 编码器块"""
    
    def __init__(self, d_model=512, n_heads=8, ff_dim=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ff_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ff_dim, d_model),
            nn.Dropout(dropout),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        """
        x: (batch, seq_len, d_model)
        """
        # 自注意力 + 残差 + 层归一化
        attn_output, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout1(attn_output))
        
        # FFN + 残差 + 层归一化
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout2(ffn_output))
        
        return x

from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

# 加载预训练的 mBART-50 模型
model_name = "facebook/mbart-large-50-many-to-many-mmt"
model = MBartForConditionalGeneration.from_pretrained(model_name)
tokenizer = MBart50TokenizerFast.from_pretrained(model_name)

# 翻译: 英文 → 中文
text = "Machine translation has made tremendous progress."
tokenizer.src_lang = "en_XX"  # 设置源语言
encoded = tokenizer(text, return_tensors="pt")

generated_tokens = model.generate(
    **encoded,
    forced_bos_token_id=tokenizer.lang_code_to_id["zh_CN"],  # 目标语言
    max_length=50,
    num_beams=5,
)
translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
print(f"原文: {text}")
print(f"翻译: {translation[0]}")

# 查看支持的语言列表
print(f"\n支持的语言数: {len(tokenizer.lang_code_to_id)}")
langs = list(tokenizer.lang_code_to_id.keys())[:10]
print(f"示例: {langs}")

# NLLB 多语言翻译实战
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 加载 NLLB-200 精简版 (3.3B 参数)
model_name = "facebook/nllb-200-distilled-600M"  # 轻量版，600M 参数
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# 多语言翻译演示
translations = {
    ("eng_Latn", "zho_Hans"): "The quick brown fox jumps over the lazy dog.",
    ("eng_Latn", "jpn_Jpan"): "The quick brown fox jumps over the lazy dog.",
    ("eng_Latn", "spa_Latn"): "The quick brown fox jumps over the lazy dog.",
    ("eng_Latn", "fra_Latn"): "The quick brown fox jumps over the lazy dog.",
    ("eng_Latn", "ara_Arab"): "The quick brown fox jumps over the lazy dog.",
}

for (src_lang, tgt_lang), text in translations.items():
    tokenizer.src_lang = src_lang
    inputs = tokenizer(text, return_tensors="pt", padding=True)
    
    translated = model.generate(
        **inputs,
        forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang],
        max_length=50,
    )
    result = tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
    print(f"{src_lang} → {tgt_lang}: {result}")

# NLLB 架构参数
print("\n=== NLLB-200 架构 ===")
print("参数量: 54B (全尺寸) / 3.3B / 1.3B (蒸馏版)")
print("架构: Mixture of Experts (MoE)")
print("专家数: 64 个，每次激活 3-6 个")
print("支持语言: 200 种")
print("训练数据: 130 亿句平行语料")

# 回译（Back-Translation）生成伪平行语料
from transformers import pipeline
import json

# 步骤 1: 加载现有的反向翻译模型（目标语言 → 源语言）
# 假设我们想训练 藏语→英语 的翻译，但只有 1000 句平行语料
# 我们还有 10 万句英语单语数据

# 模拟反向翻译
def back_translate(monolingual_tgt_sentences, reverse_model_name):
    """
    回译: 将目标语言的单语数据翻译回源语言
    monolingual_tgt_sentences: 目标语言的句子列表
    reverse_model_name: 反向翻译模型 (tgt → src)
    """
    reverse_translator = pipeline(
        "translation",
        model=reverse_model_name,
        device=0  # GPU
    )
    
    pseudo_parallel = []
    for sent in monolingual_tgt_sentences:
        result = reverse_translator(sent, max_length=200)
        pseudo_src = result[0]['translation_text']
        pseudo_parallel.append({
            "src": pseudo_src,
            "tgt": sent
        })
    
    return pseudo_parallel

# 演示
monolingual_english = [
    "The weather is beautiful today.",
    "She enjoys reading books in the library.",
    "Technology is changing our lives rapidly.",
    "Children love to play in the park.",
]

print("=== 回译演示 ===")
print("原始单语数据（目标语言=英语）:")
for sent in monolingual_english:
    print(f"  EN: {sent}")

# 伪平行语料（源语言由模型生成）
# 实际使用中需要调用真实翻译模型
print("\n伪平行语料 (模拟回译结果):")
for i, sent in enumerate(monolingual_english):
    # 这里用简单替换模拟反向翻译
    pseudo_src = f"[反向翻译结果{i+1}]"
    print(f"  SRC: {pseudo_src} → EN: {sent}")

print(f"\n回译生成了 {len(monolingual_english)} 条伪平行语料")
print("可以与真实平行语料混合训练，大幅提升低资源翻译质量")

# 低资源翻译的训练策略对比
import torch
import torch.nn as nn

print("=== 低资源翻译策略对比 ===\n")

# 策略 1: 直接从零训练 (Zero-Shot)
print("策略 1: 从零训练 (Baseline)")
print("  数据: 1000 句平行语料")
print("  模型: Transformer Base (65M 参数)")
print("  预期 BLEU: 8-12")
print("  问题: 严重过拟合，欠拟合同时存在")

# 策略 2: 迁移学习 (Transfer Learning)
print("\n策略 2: 迁移学习")
print("  步骤 1: 英-法 (100M 句) 预训练")
print("  步骤 2: 英-藏 (1K 句) 微调")
print("  预期 BLEU: 15-20")
print("  优势: 翻译能力迁移，只需适配语言对特定部分")

# 策略 3: 多语言联合训练
print("\n策略 3: 多语言联合训练")
print("  数据: 英-法(100M) + 英-中(50M) + 英-藏(1K) 混合")
print("  模型: 共享编码器 + 语言特定嵌入")
print("  预期 BLEU: 18-25")
print("  优势: 正迁移，高资源语言帮助低资源语言学习")

# 策略 4: 回译增强 + 联合训练
print("\n策略 4: 回译 + 联合训练 (最优)")
print("  步骤 1: 英→藏 模型回译 100K 藏语句 → 伪英句")
print("  步骤 2: 伪平行数据(100K) + 真实平行数据(1K) 混合")
print("  步骤 3: 多语言联合训练")
print("  预期 BLEU: 25-32")
print("  优势: 伪数据扩大训练规模，联合训练稳定学习")

# 模拟训练数据量对比
print("\n=== 训练数据量对比 ===")
strategies = {
    "从零训练": 1_000,
    "迁移学习": 1_000 + 100_000_000,  # 预训练 + 微调
    "多语言联合": 1_000 + 50_000_000,  # 低资源 + 高资源混合
    "回译+联合": 1_000 + 100_000 + 50_000_000,  # 真实 + 伪 + 高资源
}
for name, count in strategies.items():
    print(f"  {name}: {count:,} 句")

# HuggingFace 翻译模型微调实战
from datasets import load_dataset
from transformers import (
    AutoTokenizer, AutoModelForSeq2SeqLM,
    DataCollatorForSeq2Seq, Seq2SeqTrainingArguments,
    Seq2SeqTrainer
)
import evaluate
import numpy as np

# 1. 加载预训练模型和分词器
model_name = "Helsinki-NLP/opus-mt-en-zh"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# 2. 加载自定义平行语料 (模拟 TED Talks 格式)
# 实际使用: dataset = load_dataset("ted_talks", "en-zh")
raw_datasets = {
    "train": {
        "translation": [
            {"en": "Hello, how are you?", "zh": "你好，你好吗？"},
            {"en": "The weather is nice today.", "zh": "今天天气很好。"},
        ]
    },
    "validation": {
        "translation": [
            {"en": "Thank you very much.", "zh": "非常感谢。"},
        ]
    }
}

# 3. 预处理函数
max_input_length = 128
max_target_length = 128
source_lang = "en"
target_lang = "zh"

def preprocess_function(examples):
    inputs = [ex[source_lang] for ex in examples["translation"]]
    targets = [ex[target_lang] for ex in examples["translation"]]
    
    model_inputs = tokenizer(
        inputs, max_length=max_input_length,
        truncation=True, padding="max_length"
    )
    
    # 目标序列也需要分词
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(
            targets, max_length=max_target_length,
            truncation=True, padding="max_length"
        )
    
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# 4. 应用预处理
# tokenized_datasets = raw_datasets.map(preprocess_function, batched=True)

print("=== 数据预处理示例 ===")
sample = {"translation": [{"en": "Machine translation is amazing.", "zh": "机器翻译太棒了。"}]}
tokenized = preprocess_function(sample)
print(f"输入 token_ids: {tokenized['input_ids'][:10]}...")
print(f"标签 token_ids: {tokenized['labels'][:10]}...")
print(f"Attention mask:  {tokenized['attention_mask'][:10]}...")

# 训练配置与执行
from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer
import evaluate

# 加载 BLEU 评估指标
bleu = evaluate.load("sacrebleu")

def compute_metrics(eval_preds):
    """评估时计算 BLEU 分数"""
    preds, labels = eval_preds
    
    # 将 token IDs 解码为文本
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
    
    # 替换 labels 中的 -100 (pad token 忽略)
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    # 计算 BLEU
    result = bleu.compute(
        predictions=decoded_preds,
        references=[[label] for label in decoded_labels]
    )
    return {"bleu": result["score"]}

# 训练配置
training_args = Seq2SeqTrainingArguments(
    output_dir="./translation-model",
    eval_strategy="epoch",
    learning_rate=3e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=10,
    predict_with_generate=True,
    fp16=True,  # 混合精度训练加速
    gradient_accumulation_steps=4,  # 模拟 batch_size=64
    warmup_ratio=0.1,  # 学习率预热
    save_strategy="epoch",
    logging_steps=50,
    load_best_model_at_end=True,
    metric_for_best_model="bleu",
    greater_is_better=True,
)

print("=== 训练配置 ===")
print(f"输出目录: {training_args.output_dir}")
print(f"学习率: {training_args.learning_rate}")
print(f"每设备 batch: {training_args.per_device_train_batch_size}")
print(f"梯度累积: {training_args.gradient_accumulation_steps}")
print(f"等效 batch: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}")
print(f"训练轮数: {training_args.num_train_epochs}")
print(f"预热比例: {training_args.warmup_ratio}")
print(f"混合精度: {training_args.fp16}")
print(f"早停: {training_args.load_best_model_at_end}")

# 数据收集器（自动处理 padding）
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

# 初始化 Trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset="tokenized_train",
    eval_dataset="tokenized_val",
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

# 开始训练
# trainer.train()
print("\n调用 trainer.train() 开始训练")
print("训练完成后调用 trainer.save_model('./final-model') 保存")