Transformer复现学习

class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        """
            input_size: 输入维度
            hidden_size: 隐藏状态维度
            output_size: 输出维度
        """
        super().__init__()
        self.hidden_size = hidden_size
        # 权重矩阵
        self.W_xh = nn.Linear(input_size, hidden_size) # U
        self.W_hh = nn.Linear(hidden_size, hidden_size) # W
        self.W_hy = nn.Linear(hidden_size, output_size) # V
        
    def forward(self, x, hidden=None):
        """
            x: [batch_size, seq_len, embd_size]
            batch_size 个样本, 每个样本序列 seq_len 个词元, 每个词元映射为 embd_size 长的嵌入向量, 也称作特征维度
        """
        batch_size = x.size(0)
        seq_len = x.size(1)
        
        # 初始化隐藏状态为全 0 [batch_size, hidden_size]
        if hidden is None:
            hidden = torch.zeros(batch_size, self.hidden_size, device=x.device)
        
        # 存储每个时间步的输出
        outputs = []
        
        for t in range(seq_len):
            x_t = x[:, t, :] # 当前时间步的输入 [batch_size, embd_size]
            hidden_t = torch.tanh(
                self.W_xh(x_t) +  # 当前输入影响 [batch_size, hidden_size]
                self.W_hh(hidden)  # 历史信息影响 [batch_size, hidden_size]
            )
            y_t = self.W_hy(hidden_t) # 输出 [batch_size, output_size]
            hidden = hidden_t # 更新隐藏状态，用于下一个时间步
            outputs.append(y_t.unsqueeze(1))  # 增加时间步维度 [batch_size, seq_len, output_size]
        
        # 将所有时间步的输出拼接起来
        outputs = torch.cat(outputs, dim=1)  # [batch_size, seq_len, output_size]
        return outputs, hidden

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        """
            input_size: 输入维度
            hidden_size: 隐藏状态维度
            output_size: 输出维度
        """
        super().__init__()
        self.hidden_size = hidden_size
        
        self.W_xf = nn.Linear(input_size, hidden_size)  # 遗忘门输入权重
        self.W_hf = nn.Linear(hidden_size, hidden_size)  # 遗忘门隐藏权重
        self.b_f = nn.Parameter(torch.zeros(hidden_size))  # 遗忘门偏置

        self.W_xi = nn.Linear(input_size, hidden_size)  # 输入门输入权重
        self.W_hi = nn.Linear(hidden_size, hidden_size)  # 输入门隐藏权重
        self.b_i = nn.Parameter(torch.zeros(hidden_size))  # 输入门偏置
        
        self.W_xo = nn.Linear(input_size, hidden_size)  # 输出门输入权重
        self.W_ho = nn.Linear(hidden_size, hidden_size)  # 输出门隐藏权重
        self.b_o = nn.Parameter(torch.zeros(hidden_size))  # 输出门偏置
        
        self.W_xc = nn.Linear(input_size, hidden_size)  # 候选记忆单元输入权重
        self.W_hc = nn.Linear(hidden_size, hidden_size)  # 候选记忆单元隐藏权重
        self.b_c = nn.Parameter(torch.zeros(hidden_size))  # 候选记忆单元偏置
        
        # 输出层
        self.W_hy = nn.Linear(hidden_size, output_size)
        
    def forward(self, x, hidden_state=None):
        """
            x: [batch_size, seq_len, embd_size]
            hidden_state: 元组 (hidden, cell_state)
                hidden: [batch_size, hidden_size] 隐藏状态
                cell_state: [batch_size, hidden_size] 细胞状态
        """
        batch_size = x.size(0)
        seq_len = x.size(1)
        
        # 初始化隐藏状态和细胞状态为全0
        if hidden_state is None:
            hidden = torch.zeros(batch_size, self.hidden_size, device=x.device)
            cell_state = torch.zeros(batch_size, self.hidden_size, device=x.device)
        else:
            hidden, cell_state = hidden_state
        
        outputs = []
        for t in range(seq_len):
            x_t = x[:, t, :]  # 当前时间步的输入 [batch_size, embd_size]
            
            # [batch_size, hidden_size]
            f_t = torch.sigmoid(self.W_xf(x_t) + self.W_hf(hidden) + self.b_f)  # 遗忘门
            i_t = torch.sigmoid(self.W_xi(x_t) + self.W_hi(hidden) + self.b_i)  # 输入门
            o_t = torch.sigmoid(self.W_xo(x_t) + self.W_ho(hidden) + self.b_o)  # 输出门
            c_tilde = torch.tanh(self.W_xc(x_t) + self.W_hc(hidden) + self.b_c)  # 候选记忆细胞
            cell_state = f_t * cell_state + i_t * c_tilde # 细胞状态 每一层
            hidden = o_t * torch.tanh(cell_state) # 隐藏状态 每一层
            
            y_t = self.W_hy(hidden)  # [batch_size, output_size]
            outputs.append(y_t.unsqueeze(1))
        
        outputs = torch.cat(outputs, dim=1)  # 所有: [batch_size, seq_len, output_size]
        return outputs, (hidden, cell_state)

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
        """
            input_dim: 输入词汇表大小 vocab_size
            emb_dim: 词嵌入维度
            hid_dim: 隐藏层维度: 即 h_t, c_t 第 t 时间步的隐藏状态和细胞状态的维度
            n_layers: LSTM层数
            dropout: 丢弃概率
        """
        super(Encoder, self).__init__()
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        # input_dim: 输入特征维度, emb_dim: 词向量维度
        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.rnn = nn.LSTM(emb_dim, hid_dim, n_layers, dropout=dropout)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        embedded = self.dropout(self.embedding(src))
        outputs, (hidden, cell) = self.rnn(embedded)
        return hidden, cell

class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
        super(Decoder, self).__init__()
        self.output_dim = output_dim # 输出词汇表大小
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.LSTM(emb_dim, hid_dim, n_layers, dropout=dropout)
        self.fc_out = nn.Linear(hid_dim, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, input, hidden, cell):
		    """
			    input: [batch_size]
		      hidden: [n_layers * n_directions, batch_size, hid_dim] -> [n_layers, batch_size, hid_dim]
		      cell: [n_layers * n_directions, batch_size, hid_dim] -> [n_layers, batch_size, hid_dim]
		    """
        # input = [1, batch size]
        input = input.unsqueeze(0) # [1, batch_size]
        embedded = self.dropout(self.embedding(input)) # [1, batc_size, emb_dim]
        
        # output: [seq_len, batch_size, hid_dim * n_directions] -> [1, batch_size, hid_dim]
        # hidden: [n_layers * n_directions, batch_size, hid_dim] -> [n_layers, batch_size, hid_dim]
        # cell: [n_layers * n_directions, batch_size, hid_dim] -> [n_layers, batch_size, hid_dim]
        output, (hidden, cell) = self.rnn(embedded, (hidden, cell))
        
        prediction = self.fc_out(output.squeeze(0)) # [batch_size, output_dim]
        return prediction, hidden, cell

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super(Seq2Seq, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.device = device
        
        assert encoder.hid_dim == decoder.hid_dim, "编码器和解码器的隐藏维度必须相等!"
        assert encoder.n_layers == decoder.n_layers, "编码器和解码器必须有相同层数!"

    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        """
            src = [src_len, batch_size] 源语言序列
            trg = [trg_len, batch_size] 目标语言序列
            teacher_forcing_ratio: 训练过程中的每个时刻，有一定概率使用上一时刻的输出作为输入，也有一定概率使用正确的 target 作为输入
        """
        batch_size = trg.shape[1]
        trg_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim # 输出词汇表大小
        
        # 初始化输出张量，用于存储每个时间步的预测结果
        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)
        
        # encoder 模块
        hidden, cell = self.encoder(src)
        # decoder 模块的第一个输入 token 是序列开始标记: <bos>
        input = trg[0, :]
        
        for t in range(1, trg_len):

            output, hidden, cell = self.decoder(input, hidden, cell)
            outputs[t] = output # 预测结果
            teacher_force = random.random() < teacher_forcing_ratio
            
            # 获取预测中概率最高的 token
            top1 = output.argmax(1)
            # 决定输入是根据预测得到的（预测可能是错误的），还是真实值（起纠偏作用）
            input = trg[t] if teacher_force else top1
        
        return outputs

def make_batch(seq_data, n_step, num_dic):
    """
        创建训练批次数据 seq_data: [batch_size, 2]
    """
    input_batch, output_batch, target_batch = [], [], []

    for seq in seq_data:
        src_word = seq[0] + 'P' * (n_step - len(seq[0]))
        trg_word = seq[1] + 'P' * (n_step - len(seq[1]))
        
        # 实际此处命名为 output 不合适
        input_idx = [num_dic[n] for n in src_word]  # Encoder 输入源序列索引
        output_idx = [num_dic[n] for n in ('S' + trg_word)]  # Decoder 输入序列索引（加S）
        target_idx = [num_dic[n] for n in (trg_word + 'E')]  # Decoder 目标序列索引（加E）
        
        input_batch.append(input_idx)
        output_batch.append(output_idx)
        target_batch.append(target_idx)
    
    input_tensor = torch.LongTensor(input_batch).t()  # [seq_len, batch_size]
    output_tensor = torch.LongTensor(output_batch).t()  # [seq_len+1, batch_size]
    target_tensor = torch.LongTensor(target_batch).t()  # [seq_len+1, batch_size]
    
    return input_tensor, output_tensor, target_tensor

def make_testbatch(input_word, n_step, num_dic):
    """创建测试批次数据"""
    # 填充输入
    input_w = input_word + 'P' * (n_step - len(input_word))
    input_idx = [num_dic[n] for n in input_w]
    output_idx = [num_dic[n] for n in ('S' + 'P' * n_step)] # 初始化为 SPPPP, 只需要第一个为 S
    
    # 转换为索引
    input_batch = torch.LongTensor([input_idx]).t()  # [seq_len, 1]
    output_batch = torch.LongTensor([output_idx]).t()  # [seq_len+1, 1]
    
    return input_batch, output_batch

if __name__ == '__main__':

    n_step = 8  # 序列最大长度, 不够则使用 padding, 单词数
    hidden_dim = 128  # 隐藏层维度
    emb_dim = 32  # 词向量维度
    dropout = 0.0 # LSTM 一层不使用 dropout
    char_arr = [c for c in 'SEPabcdefghijklmnopqrstuvwxyz']
    num_dic = {n: i for i, n in enumerate(char_arr)}
    n_class = len(num_dic)  # 词汇表大小
    # 训练数据
    seq_data = [
        ['man', 'woman'], 
        ['black', 'white'], 
        ['king', 'queen'], 
        ['girl', 'boy'], 
        ['up', 'down'], 
        ['high', 'low'],
        ['men', 'women'],
        ['you', 'me']
    ]
    batch_size = len(seq_data) # 样本数量
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    encoder = Encoder(n_class, emb_dim, hidden_dim, n_layers=1, dropout=dropout)
    decoder = Decoder(n_class, emb_dim, hidden_dim, n_layers=1, dropout=dropout)
    model = Seq2Seq(encoder, decoder, device).to(device)
    
    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss(ignore_index=num_dic['P'])
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # 准备训练数据
    input_batch, output_batch, target_batch = make_batch(seq_data, n_step, num_dic)
    input_batch = input_batch.to(device)
    output_batch = output_batch.to(device)
    target_batch = target_batch.to(device)
    
    for epoch in range(5000):
        optimizer.zero_grad()
        outputs = model(input_batch, output_batch) # [trg_len, batch_size, vocab_size]
        
        # 计算损失: 跳过第一个时间步
        outputs = outputs[1:].reshape(-1, n_class)  # [trg_len-1 * batch_size, vocab_size]
        targets = target_batch[1:].reshape(-1)  # [trg_len-1 * batch_size]
        
        loss = criterion(outputs, targets)
        
        loss.backward()
        optimizer.step()
        
        if (epoch + 1) % 500 == 0:
            print(f'Epoch: {epoch+1:04d}, Loss: {loss.item():.6f}')

def translate(word):
    input_batch, output_batch = make_testbatch(word, n_step, num_dic)
    input_batch = input_batch.to(device)
    output_batch = output_batch.to(device)
    
    model.eval()
    with torch.no_grad():
        outputs = model(input_batch, output_batch, teacher_forcing_ratio=0.0)
    
    # 获取预测结果
    # outputs形状: [seq_len+1, 1, n_class]
    predictions = outputs.argmax(2)  # [seq_len+1, 1]
    
    # 转换为字符
    decoded = []
    for idx in predictions[:, 0].cpu().numpy():
        decoded.append(char_arr[idx])
    
    # 找到结束标记
    try:
        end = decoded.index('E')
        translated = ''.join(decoded[:end])
    except ValueError:
        translated = ''.join(decoded)
    
    return translated.replace('P', '')

test_words = ['man', 'mans', 'king', 'black', 'ups']
for word in test_words:
		result = translate(word)
		print(f'{word} -> {result}'

class MHABlock(nn.Module):
    def __init__(self, d_model: int, h: int, dropout: float):
        super().__init__()
        self.d_model = d_model # 嵌入向量维度/特征维度
        self.h = h # 头数
        assert d_model % h == 0, "h 无法整除 d_model"

        self.d_k = d_model // h # 每个头的特征维度
        self.WQ = nn.Linear(d_model, d_model, bias=False)
        self.WK = nn.Linear(d_model, d_model, bias=False)
        self.WV = nn.Linear(d_model, d_model, bias=False)
        self.WO = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    @staticmethod
    def attention(Q: torch.tensor, K: torch.tensor, V: torch.tensor, mask, dropout: nn.Dropout):
        # Q, K, V: (B, h, S, d_k)
        d_k = Q.shape[-1]
        attn_scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            attn_scores.masked_fill_(mask == 0, -1e9)
        attn_scores = attn_scores.softmax(dim=-1)
        if dropout is not None:
            attn_scores = dropout(attn_scores)
        return attn_scores @ V
    
    def forward(self, XQ: torch.tensor, XK: torch.tensor, XV: torch.tensor, mask):
        # 多个 X 用于注意力变体: 交叉注意力, Q 来自解码器, K、V 来自编码器
        Q = self.WQ(XQ)
        K = self.WK(XK)
        V = self.WV(XV)
        
        # 多头拆分 (B, S, d_model) ==> (B, S, h, d_k) ==> (B, h, S, d_k)
        Q = Q.view(Q.shape[0], Q.shape[1], self.h, self.d_k).transpose(1, 2)
        K = K.view(K.shape[0], K.shape[1], self.h, self.d_k).transpose(1, 2)
        V = V.view(V.shape[0], V.shape[1], self.h, self.d_k).transpose(1, 2)

        attn = MHABlock.attention(Q, K, V, mask, self.dropout)
        # 多头拼接 (B, h, S, d_k) ==> (B, S, h, d_k) ==> (B, S, d_model)
        attn = attn.transpose(1, 2).contiguous().view(attn.shape[0], -1, self.d_model)
        return self.WO(attn)

class LayerNorm(nn.Module):
    def __init__(self, features: int, eps: float=1e-6):
        super().__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta = nn.Parameter(torch.zeros(features))

    def forward(self, x):
        # x: (B, S, d_model)
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True) # (B, S, 1)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, seq_len: int, dropout: float):
        super().__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.dropout = nn.Dropout(dropout)

        # pe: (seq_len, d_model)
        pe = torch.zeros(seq_len, d_model)
        # position: (seq_len, 1) 数值从 0 到 seq_len
        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
        # div_term: (d_model / 2)
        div_term = torch.pow(10000.0, -torch.arange(0, d_model, 2, dtype=torch.float) / d_model)
        
        # (seq_len, d_model/2)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # (1, seq_len, d_model) 在 batch 维度上广播, 方便 x[B, S, d] + pe[1, S, d]
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe) # 注册为模块缓冲区, 不参与梯度更新, 位置编码固定, 不学习, 但随模型一起保存在 GPU
    
    def forward(self, x):
        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
        return self.dropout(x)

class FFNBlock(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout: float):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )

    def forward(self, x):
        return self.net(x)

class EncoderBlock(nn.Module):
    def __init__(self, features: int, num_heads: int, d_ff: int, dropout: float):
        super().__init__()
        self.mha_block = MHABlock(d_model=features, h=num_heads, dropout=dropout)
        self.ffn = FFNBlock(d_model=features, d_ff=d_ff, dropout=dropout)
        self.norm1 = LayerNorm(features=features)
        self.norm2 = LayerNorm(features=features)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, src_mask):
        # Pre-Norm
        _x = self.norm1(x)
        x = x + self.dropout(self.mha_block(_x, _x, _x, src_mask))

        _x = self.norm2(x)
        x = x + self.dropout(self.ffn(_x))
        return x

class DecoderBlock(nn.Module):
    def __init__(self, features: int, num_heads: int, d_ff: int, dropout: float):
        super().__init__()
        self.mha_block = MHABlock(d_model=features, h=num_heads, dropout=dropout)
        self.masked_mha_block = MHABlock(d_model=features, h=num_heads, dropout=dropout)
        self.ffn_block = FFNBlock(d_model=features, d_ff=d_ff, dropout=dropout)
        self.norm1 = LayerNorm(features)
        self.norm2 = LayerNorm(features)
        self.norm3 = LayerNorm(features)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, encoder_output, src_mask, tgt_mask):
        # Pre-Norm
        _x = self.norm1(x)
        x = x + self.dropout(self.masked_mha_block(_x, _x, _x, tgt_mask))
 
        _x = self.norm2(x)
        x = x + self.dropout(self.mha_block(_x, encoder_output, encoder_output, src_mask))
 
        _x = self.norm3(x)
        x = x + self.dropout(self.ffn_block(_x))
        return x

class Transformer(nn.Module):
    def __init__( 
            self,
            src_vocab_size: int,
            tgt_vocab_size: int,
            d_model: int=512,
            nhead: int=8,
            num_encoder_layers: int=6,
            num_decoder_layers: int=6,
            dim_feedforward: int=2048,
            dropout: float=0.1,
            src_seq_len: int=512,
            tgt_seq_len: int=512
    ):
        super().__init__()
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.input_pe = PositionalEncoding(d_model, src_seq_len, dropout)
        self.output_pe = PositionalEncoding(d_model, tgt_seq_len, dropout)
 
        self.encoder = Encoder(num_encoder_layers, d_model, nhead, dim_feedforward, dropout)
        self.decoder = Decoder(num_decoder_layers, d_model, nhead, dim_feedforward, dropout)
 
        self.fc_out = nn.Linear(d_model, tgt_vocab_size)
        self.dropout = nn.Dropout(dropout)
        self.d_model = d_model
 
    def forward(self, src, tgt, src_mask, tgt_mask):
        src = self.input_pe(self.src_embedding(src) * math.sqrt(self.d_model))
        tgt = self.output_pe(self.tgt_embedding(tgt) * math.sqrt(self.d_model))
 
        encoder_output = self.encoder(src, src_mask)
        decoder_output = self.decoder(tgt, encoder_output, src_mask, tgt_mask)
 
        output = self.fc_out(decoder_output)
        return output

from torch.utils.data import Dataset, DataLoader
import json
import os
 
class TranslationDataset(Dataset):
    def __init__(self, data_path, src_tokenizer, tgt_tokenizer, max_len=128):
        self.data = self.load_data(data_path)
        self.src_tokenizer = src_tokenizer # 源语言分词器 chinese
        self.tgt_tokenizer = tgt_tokenizer # 目标语言分词器 english
        self.max_len = max_len # 最大序列长度
 
    def load_data(self, path):
        """数据加载"""
        data = []
        if not os.path.exists(path):
            print(f"ERROR：文件 {path} 不存在")
            return None
 
        with open(path, 'r', encoding='utf-8') as f:
            content = f.read().strip()
            lines = content.split('\n')
            for line in lines:
                data.append(json.loads(line))
        return data
 
    def __len__(self):
        return len(self.data)
 
    def __getitem__(self, idx):
        item = self.data[idx]
        src_text = item.get('chinese', '')
        tgt_text = item.get('english', '')
 
        # BertTokenizer 会自动添加[CLS]和[SEP]来表示句子的起始和中止，在这里手动添加
        src_encoding = self.src_tokenizer(src_text, truncation=True, max_length=self.max_len, add_special_tokens=False)
        tgt_encoding = self.tgt_tokenizer(tgt_text, truncation=True, max_length=self.max_len, add_special_tokens=False)
 
        # 获取 ID 列表
        src_ids = src_encoding['input_ids']
        tgt_ids = tgt_encoding['input_ids']
 
        # 添加特殊 token (借用BERT的[CLS]作为SOS, [SEP]作为EOS)
        sos_id = self.src_tokenizer.cls_token_id
        eos_id = self.src_tokenizer.sep_token_id
 
        return torch.tensor([sos_id] + src_ids + [eos_id], dtype=torch.long), \
               torch.tensor([sos_id] + tgt_ids + [eos_id], dtype=torch.long)

def create_masks(src, tgt, pad_idx=0):
    src_seq_len = src.shape[1]
    tgt_seq_len = tgt.shape[1]
 
    # (src != pad_idx) -> (B, S) -> unsqueeze -> (B, 1, 1, S)
    src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2)
 
    # 屏蔽pad: (B, 1, 1, T)
    tgt_pad_mask = (tgt != pad_idx).unsqueeze(1).unsqueeze(2)
    # 屏蔽未来 (Look-ahead): (1, 1, T, T)
    # 下三角矩阵为1，其余为0
    tgt_sub_mask = torch.tril(torch.ones((tgt_seq_len, tgt_seq_len), device=src.device)).bool()
    # 两个矩阵结合
    tgt_mask = tgt_pad_mask & tgt_sub_mask.unsqueeze(0).unsqueeze(0)
 
    return src_mask, tgt_mask

class Config:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
    train_data_path = "./translation2019zh_train.json"
    valid_data_path = "./translation2019zh_valid.json"
 
    # 训练超参数
    batch_size = 32
    num_epochs = 1
    learning_rate = 0.0001
    max_len = 128  # 限制最大句子长度
 
    # 模型参数
    d_model = 512
    nhead = 8
    num_encoder_layers = 3
    num_decoder_layers = 3
    dim_feedforward = 2048
    dropout = 0.1
 
    # 使用 Hugging Face 的 BERT 分词器
    src_tokenizer_name = 'bert-base-chinese'
    tgt_tokenizer_name = 'bert-base-uncased'

from tqdm import tqdm
 
def train(model, iterator, optimizer, criterion, clip, device):
    model.train()
    epoch_loss = 0
    proc_bar = tqdm(iterator, desc="Training")
    for i, (src, tgt) in enumerate(proc_bar):
        src = src.to(device)
        tgt = tgt.to(device)
        # 原始目标序列 [SOS] I love China [EOS] -> [1, 10, 20, 30, 2]
        # tgt_input: [SOS] I love China       -> [1, 10, 20, 30]
        # tgt_output: I love China [EOS]      -> [10, 20, 30, 2]
        tgt_input = tgt[:, :-1]		# tgt输入: 去掉最后一个 token <eos>
        tgt_output = tgt[:, 1:]		# tgt输出(标签): 去掉第一个 token <sos>
 
        src_mask, tgt_mask = create_masks(src, tgt_input, pad_idx=0)
 
        optimizer.zero_grad()
        # 前向传播 每个位置表示某样本在某时间步在词汇表上的概率分布
        output = model(src, tgt_input, src_mask, tgt_mask)	# (batch_size, seq_len, tgt_vocab_size)
        output_dim = output.shape[-1]
 
        # 拉平以计算损失
        output = output.contiguous().view(-1, output_dim) # [batch_size * seq_len, tgt_vocab_size]
        tgt_output = tgt_output.contiguous().view(-1) # [batch_size * seq_len]
        loss = criterion(output, tgt_output)
 
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip) # 梯度裁剪: 防止梯度爆炸，稳定训练
        optimizer.step()
 
        epoch_loss += loss.item()
        proc_bar.set_postfix(loss=loss.item())
 
    return epoch_loss / len(iterator)

def evaluate(model, iterator, criterion, device):
    model.eval()
    epoch_loss = 0
 
    with torch.no_grad(): # 禁用梯度计算
        for i, (src, tgt) in enumerate(iterator):
            src = src.to(device)
            tgt = tgt.to(device)
 
            tgt_input = tgt[:, :-1]
            tgt_output = tgt[:, 1:]
 
            src_mask, tgt_mask = create_masks(src, tgt_input, pad_idx=0)
 
            output = model(src, tgt_input, src_mask, tgt_mask)
 
            output_dim = output.shape[-1]
            output = output.contiguous().view(-1, output_dim)
            tgt_output = tgt_output.contiguous().view(-1)
 
            loss = criterion(output, tgt_output)
            epoch_loss += loss.item()
 
    return epoch_loss / len(iterator)

import torch.optim as optim
import os

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HF_HOME"] = "./model/"
from transformers import AutoTokenizer  # 缓存和下载地址的配置都必须在导入这个包之前完成
 
src_tokenizer = AutoTokenizer.from_pretrained(Config.src_tokenizer_name)
tgt_tokenizer = AutoTokenizer.from_pretrained(Config.tgt_tokenizer_name)
 
train_dataset = TranslationDataset(Config.train_data_path, src_tokenizer, tgt_tokenizer, max_len=Config.max_len)
val_dataset = TranslationDataset(Config.valid_data_path, src_tokenizer, tgt_tokenizer, max_len=Config.max_len)
# 传入数据的同时填充序列到相同长度（按照批次中最长的那个序列）
train_loader = DataLoader(train_dataset, batch_size=Config.batch_size, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(val_dataset, batch_size=Config.batch_size, shuffle=False, collate_fn=collate_fn)
 
model = Transformer(
    src_vocab_size=src_tokenizer.vocab_size,
    tgt_vocab_size=tgt_tokenizer.vocab_size,
    d_model=Config.d_model,
    nhead=Config.nhead,
    num_encoder_layers=Config.num_encoder_layers,
    num_decoder_layers=Config.num_decoder_layers,
    dim_feedforward=Config.dim_feedforward,
    dropout=Config.dropout,
    src_seq_len=Config.max_len + 5,
    tgt_seq_len=Config.max_len + 5
).to(Config.device)
 
# 初始化权重Xavier init通常对Transformer有效, 只对维度大于1的权重进行初始化（排除偏置项）
def initialize_weights(m):
    if hasattr(m, 'weight') and m.weight.dim() > 1:
        nn.init.xavier_uniform_(m.weight.data)
model.apply(initialize_weights)
 
optimizer = optim.Adam(model.parameters(), lr=Config.learning_rate)
# 忽略<pad>的loss
criterion = nn.CrossEntropyLoss(ignore_index=0)
best_valid_loss = float('inf')
start_epoch = 0
 
for epoch in range(start_epoch, Config.num_epochs):
    train_loss = train(model, train_loader, optimizer, criterion, 1.0, Config.device)
    valid_loss = evaluate(model, val_loader, criterion, Config.device)
    # 保存最佳模型
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'transformer_model.pt')

def translate_sentence(sentence, src_tokenizer, tgt_tokenizer, model, device, max_len=128):
    model.eval()
 
    # 手动添加特殊 token，因此 add_special_tokens=False, 获取 token_id
    tokens = src_tokenizer(sentence, truncation=True, max_length=max_len, add_special_tokens=False)['input_ids']
 
    sos_idx = src_tokenizer.cls_token_id
    eos_idx = src_tokenizer.sep_token_id
 
    src_indices = [sos_idx] + tokens + [eos_idx]
    src_tensor = torch.LongTensor(src_indices).unsqueeze(0).to(device) # (1, src_len)
 
    # 屏蔽 <pad>
    src_mask = (src_tensor != 0).unsqueeze(1).unsqueeze(2) # (1, 1, 1, src_len)
 
    with torch.no_grad():
        # 对应Transformer.forward中的: src = self.input_pe(self.src_embedding(src) * math.sqrt(self.d_model))
        src_emb = model.src_embedding(src_tensor) * math.sqrt(model.d_model)
        src_emb = model.input_pe(src_emb)
        encoder_output = model.encoder(src_emb, src_mask)
 
    # 自回归生成
    tgt_indices = [sos_idx] # 初始输入只有<sos>
 
    for i in range(max_len):
        tgt_tensor = torch.LongTensor(tgt_indices).unsqueeze(0).to(device) # (1, current_tgt_len)
 
        # Look-ahead mask (注意, 在推理的过程中实际上不需要掩码了, 因为Decoder的过程本来就是预测未来)
        tgt_len = tgt_tensor.shape[1]
        tgt_sub_mask = torch.tril(torch.ones((tgt_len, tgt_len), device=device)).bool() # (T, T)
        tgt_mask = tgt_sub_mask.unsqueeze(0).unsqueeze(0) # (1, 1, T, T)
 
        with torch.no_grad():
            # 对应Transformer.forward中的: tgt = self.output_pe(self.tgt_embedding(tgt) * math.sqrt(self.d_model))
            tgt_emb = model.tgt_embedding(tgt_tensor) * math.sqrt(model.d_model)
            tgt_emb = model.output_pe(tgt_emb)
            output = model.decoder(tgt_emb, encoder_output, src_mask, tgt_mask)
            # 通过全连接层映射到词表
            output = model.fc_out(output)
 
        # 获取最后一个时间步的预测结果 output shape: (1, T, tgt_vocab_size)
        # argmax(2): 沿着第二个维度取最大值 -> (1, T) 每个位置保存概率最高的 token id
        # [:, -1]: -> (1,)
        # .item(): -> 1 转换为标量
        pred_token = output.argmax(2)[:, -1].item() # 取概率最高的 token
        tgt_indices.append(pred_token)
 
        # 预测出结束符<eos>则停止生成
        if pred_token == eos_idx:
            break
    # token id 转换回 token, skip_special_tokens=True 会自动去掉特殊token
    translated_sentence = tgt_tokenizer.decode(tgt_indices, skip_special_tokens=True)
    return translated_sentence

src_tokenizer = AutoTokenizer.from_pretrained(Config.src_tokenizer_name)
tgt_tokenizer = AutoTokenizer.from_pretrained(Config.tgt_tokenizer_name)
 
model = Transformer(
    src_vocab_size=src_tokenizer.vocab_size,
    tgt_vocab_size=tgt_tokenizer.vocab_size,
    d_model=Config.d_model,
    nhead=Config.nhead,
    num_encoder_layers=Config.num_encoder_layers,
    num_decoder_layers=Config.num_decoder_layers,
    dim_feedforward=Config.dim_feedforward,
    dropout=Config.dropout,
    src_seq_len=Config.max_len + 5,
    tgt_seq_len=Config.max_len + 5
).to(Config.device)

model_path = 'transformer_model.pt'
model.load_state_dict(torch.load(model_path, map_location=Config.device))
 
while True:
    sentence = input("Chinese: ")
    if sentence.lower() == 'q':
        break
 
    translation = translate_sentence(sentence, src_tokenizer, tgt_tokenizer, model, Config.device)
    print(f"==> English: {translation}")