目录
一、项目准备与环境依赖
二、数据预处理
1. 数据集加载与划分
2. 构建自定义 Tokenizer
3. 词表构建与文本编码
三、构建 DataLoader
四、搭建 Transformer 翻译模型
1. 位置编码层
2. 完整翻译模型
五、模型训练
六、模型预测
七、全部完整代码
Transformer 打破了此前循环神经网络（RNN）等在序列建模任务中的垄断，它以自注意力机制为核心，成为 NLP、计算机视觉等领域的基础架构，像 GPT、BERT、ViT 等知名模型均基于它构建。以下是其具体介绍及核心优点：
- 核心结构
- 编码器：通常由多层相同结构堆叠而成，每层包含多头自注意力机制、前馈神经网络，且配有残差连接与层归一化。先通过输入嵌入将文本等元素转为向量并添加位置编码保留顺序信息，再经多头自注意力捕捉元素间全局依赖，最后由前馈神经网络进一步处理特征。
- 解码器：同样由多层结构组成，除了包含编码器类似的前馈神经网络等模块，还多了掩码多头自注意力和编码器 - 解码器注意力模块。掩码机制能避免生成序列时模型看到后续元素，编码器 - 解码器注意力则可让解码器获取编码器的上下文信息。
- 核心优点
- 并行计算效率高：传统 RNN 需按顺序逐个处理序列元素，当前输出依赖前一时刻结果，无法发挥 GPU 并行计算优势。而 Transformer 借助自注意力机制，可同时计算序列中所有元素的依赖关系，比如处理百词句子时能同步计算所有词的注意力关联，训练速度相比 RNN 可提升 10 - 100 倍，原本 RNN 需数周训练的模型，Transformer 仅需数天就能完成。
- 长距离依赖建模强：RNN 和 LSTM 处理长序列时易出现梯度消失问题，导致难以捕捉远距离元素关联，例如长句中难以关联首尾的指代关系。Transformer 的每个元素能直接与序列中所有其他元素建立联系，通过注意力权重量化关联强度，哪怕是序列两端的元素，也能一次注意力计算就建立关联，完美解决长距离依赖难题。
- 适配多场景且泛化性好：其设计的多头注意力可在多个子空间捕捉不同维度语义关联，比如机器翻译中既能关注语法结构，又能捕捉语义关联。同时还支持掩码注意力、交叉注意力等多种形式，适配翻译、文本生成、摘要等不同任务。而且它的统一架构能跨领域迁移，不仅在 NLP 领域表现出色，在计算机视觉领域，将图像分割为类似文本 token 的图像块后，也能用 Transformer 建模块间关系，在图像分类等任务上超越传统卷积神经网络。
在自然语言处理领域，机器翻译是极具代表性的任务之一，而 Transformer 模型凭借其自注意力机制，成为了机器翻译任务的主流架构。本文将详细介绍如何基于 PyTorch 框架，从零构建一个简单的中英翻译模型，涵盖数据预处理、模型搭建、训练及预测全流程。
一、项目准备与环境依赖
首先需要搭建对应的开发环境，确保安装以下核心依赖库：
import torch import pandas as pd from sklearn.model_selection import train_test_split from typing import List from torch import nn,optim from torch.nn.utils.rnn import pad_sequence from tqdm import tqdm import json from nltk.tokenize.treebank import TreebankWordDetokenizer,TreebankWordTokenizer from torch.utils.data import Dataset, DataLoader import math二、数据预处理
1. 数据集加载与划分
本文使用的是中英双语平行语料库（cmn.txt），包含 29155 条中英对照语句。首先加载数据并划分训练集和测试集：
# 加载数据 data=pd.read_csv('./data/cmn.txt',sep='\t',header=None,usecols=[0,1],names=['en','zh']) # 划分训练集和测试集（8:2） train_df,test_df=train_test_split(data,test_size=0.2)数据集如图
2. 构建自定义 Tokenizer
为了将文本转换为模型可识别的数字序列，我们构建了基础 Tokenizer 类，并分别实现中文和英文的分词逻辑：
- 中文 Tokenizer：按字符级分词，因为中文汉字无天然分隔符；
- 英文 Tokenizer：基于 TreebankWordTokenizer 实现单词级分词，同时支持解码（将数字序列转回文本）。
核心代码如下：
class BaseTokenizer: # 基础Tokenizer类，定义通用逻辑 pad_index=0 unk_index=1 start_index=2 end_index=3 def __init__(self,vocab_list): self.vocab_list=vocab_list self.vocab_size = len(vocab_list) self.world2index={value:index for index,value in enumerate(vocab_list)} self.index2world={index:value for index,value in enumerate(vocab_list)} @staticmethod def tokenize(text:str)->List[str]: pass def encode(self,text:str,is_mark=False)->List[int]: tokens=self.tokenize(text) tokens_index=[self.world2index.get(token,self.unk_index) for token in tokens] if is_mark: tokens_index.insert(0,self.start_index) tokens_index.append(self.end_index) return tokens_index @classmethod def build_vocab(cls,sentences:List[str],unk_token='<unknown>',pad_token='<padding>',start_token='<start>',end_token='<end>',vocab_path='./vocab.json'): vocab_set=set() for sentence in tqdm(sentences,desc='构建词表：'): vocab_set.update(cls.tokenize(sentence)) vocab_list = [pad_token, unk_token,start_token,end_token] + sorted(list(vocab_set)) vocab_dict={index:value for index,value in enumerate(vocab_list)} with open(vocab_path,'w',encoding='utf-8') as f: json.dump(vocab_dict, f, ensure_ascii=False, indent=2) @classmethod def read_vocab(cls,vocab_path='./vocab.json'): with open(vocab_path,'r',encoding='utf-8') as f: json_dict=json.load(f) sentences=[value for key,value in json_dict.items()] return cls(sentences) class ChinseeTokenizer(BaseTokenizer): # 中文Tokenizer @staticmethod def tokenize(text:str)->List[str]: return list(text) class EnglishTokenizer(BaseTokenizer): # 英文Tokenizer tokenizer=TreebankWordTokenizer() detokenizer=TreebankWordDetokenizer() @classmethod def tokenize(cls,text:str)->List[str]: return cls.tokenizer.tokenize(text) def decode(self,indexs:List[str])->str: tokens=[self.index2world.get(index,'<unknown>') for index in indexs] return self.detokenizer.detokenize(tokens)3. 词表构建与文本编码
基于训练集构建中英词表，并将所有文本转换为数字序列，最后保存为 JSONL 格式：
# 构建词表 ChinseeTokenizer.build_vocab(sentences=train_df['zh'].tolist(), vocab_path='./zh_vocab.json') EnglishTokenizer.build_vocab(sentences=train_df['en'].tolist(), vocab_path='./en_vocab.json') # 加载词表 cn_tokenizer = ChinseeTokenizer.read_vocab('./zh_vocab.json') en_tokenizer = EnglishTokenizer.read_vocab('./en_vocab.json') # 文本编码 train_df['en']=train_df['en'].apply(lambda x:en_tokenizer.encode(x,is_mark=True)) train_df['zh']=train_df['zh'].apply(lambda x:cn_tokenizer.encode(x)) test_df['en']=test_df['en'].apply(lambda x:en_tokenizer.encode(x,is_mark=True)) test_df['zh']=test_df['zh'].apply(lambda x:cn_tokenizer.encode(x)) # 保存编码后的数据 train_df.to_json('./train.jsonl',orient='records',lines=True) test_df.to_json('./test.jsonl',orient='records',lines=True)三、构建 DataLoader
自定义 Dataset 类加载编码后的数据，并通过 collate_fn 实现批次内序列的 padding 对齐：
class TranslationDataset(Dataset): def __init__(self,path): self.data=pd.read_json(path,orient='records',lines=True).to_dict(orient='records') def __len__(self): return len(self.data) def __getitem__(self,index): input_tensor=torch.tensor(self.data[index]['zh'],dtype=torch.long) target_tensor=torch.tensor(self.data[index]['en'],dtype=torch.long) return input_tensor,target_tensor # 自定义collate_fn，实现padding def collate_fn(batch): input_tensor=[tensor[0] for tensor in batch] target_tensor=[tensor[1] for tensor in batch] input_tensor=pad_sequence(sequences=input_tensor,batch_first=True,padding_value=0) target_tensor=pad_sequence(sequences=target_tensor,batch_first=True,padding_value=0) return input_tensor,target_tensor # 构建DataLoader train_dataset=TranslationDataset('./train.jsonl') test_dataset=TranslationDataset('./test.jsonl') train_dataloader=DataLoader(train_dataset,batch_size=32,shuffle=True,collate_fn=collate_fn) test_dataloader=DataLoader(test_dataset,batch_size=32,collate_fn=collate_fn)四、搭建 Transformer 翻译模型
Transformer 模型核心由编码器（Encoder）、解码器（Decoder）组成，我们还添加了位置编码层，为序列注入位置信息：
1. 位置编码层
class PositionalEncoding(nn.Module): def __init__(self,max_len,dim_model): super(PositionalEncoding,self).__init__() pe=torch.zeros([max_len,dim_model],dtype=torch.float) for pos in range(max_len): for i in range(0,dim_model,2): pe[pos,i] = math.sin(pos/(10000**(i/dim_model))) pe[pos,i+1] = math.cos(pos/(10000**(i/dim_model))) self.register_buffer('pe',pe) def forward(self,x): seq_len=x.shape[1] part_pe=self.pe[0:seq_len] return x+part_pe2. 完整翻译模型
class TranslationModel(nn.Module): def __init__(self,zh_vocab_size,en_vocab_size,zh_padding_idx,en_padding_idx): super(TranslationModel,self).__init__() # 嵌入层 self.zh_embedding=nn.Embedding(num_embeddings=zh_vocab_size,embedding_dim=128,padding_idx=zh_padding_idx) self.en_embedding=nn.Embedding(num_embeddings=en_vocab_size,embedding_dim=128,padding_idx=en_padding_idx) # 位置编码 self.position_encoding=PositionalEncoding(max_len=500,dim_model=128) # Transformer核心 self.transformer=nn.Transformer( d_model=128, nhead=8, num_encoder_layers=6, num_decoder_layers=6, batch_first=True, dropout=0.1, ) # 输出层 self.linear=nn.Linear(in_features=128,out_features=en_vocab_size) def forward(self,src,tgt,src_pad_mask,tgt_mask): memory=self.encode(src,src_pad_mask) outputs=self.decode(tgt,memory,tgt_mask,src_pad_mask) return outputs def encode(self,src,src_pad_mask): zh_embed=self.zh_embedding(src) zh_embed=self.position_encoding(zh_embed) memory=self.transformer.encoder(src=zh_embed,src_key_padding_mask=src_pad_mask) return memory def decode(self,tgt,memory,tgt_mask,memory_pad_mask): en_embed=self.en_embedding(tgt) en_embed=self.position_encoding(en_embed) output=self.transformer.decoder(tgt=en_embed,memory=memory,tgt_mask=tgt_mask,memory_key_padding_mask=memory_pad_mask) outputs=self.linear(output) return outputs五、模型训练
设置训练超参数，定义损失函数和优化器，分训练、验证阶段迭代训练模型，并保存最优模型：
# 设备选择 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 初始化模型 model=TranslationModel( zh_vocab_size=cn_tokenizer.vocab_size, en_vocab_size=en_tokenizer.vocab_size, zh_padding_idx=0, en_padding_idx=0 ).to(device) # 训练配置 epochs=5 lr=1e-4 loss_fn=nn.CrossEntropyLoss(ignore_index=en_tokenizer.pad_index) # 忽略padding的损失 optimizer=optim.Adam(model.parameters(),lr=lr) # 训练循环 best_loss=float('inf') for epoch in range(epochs): print(f'================第{epoch+1}轮=================') # 训练阶段 model.train() train_total_loss=0.0 for train_x,train_y in tqdm(train_dataloader,desc='训练'): src,tgt=train_x.to(device),train_y.to(device) decoder_inputs=tgt[:,:-1] # 解码器输入（去掉最后一个token） decoder_targets=tgt[:,1:] # 解码器目标（去掉第一个token） src_pad_mask=(src==model.zh_embedding.padding_idx) tgt_mask=model.transformer.generate_square_subsequent_mask(sz=decoder_inputs.shape[1]).to(device) pred_y=model(src,decoder_inputs,src_pad_mask,tgt_mask) loss=loss_fn(pred_y.reshape(-1,pred_y.shape[-1]),decoder_targets.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() train_total_loss+=loss.item() # 验证阶段 model.eval() test_total_loss=0.0 with torch.no_grad(): for test_x, test_y in tqdm(test_dataloader,desc='验证'): src,tgt=test_x.to(device),test_y.to(device) decoder_inputs=tgt[:,:-1] decoder_targets=tgt[:,1:] src_pad_mask=(src==model.zh_embedding.padding_idx) tgt_mask=model.transformer.generate_square_subsequent_mask(sz=decoder_inputs.shape[1]).to(device) pred_y=model(src,decoder_inputs,src_pad_mask,tgt_mask) loss=loss_fn(pred_y.reshape(-1,pred_y.shape[-1]),decoder_targets.reshape(-1)) test_total_loss+=loss.item() # 计算平均损失 avg_train_loss=train_total_loss/len(train_dataloader) avg_test_loss=test_total_loss/len(test_dataloader) print(f'训练平均loss:{avg_train_loss},验证平均loss:{avg_test_loss}') # 保存最优模型 if test_total_loss<best_loss: best_loss=test_total_loss torch.save(model.state_dict(),'./best_model.pt')模型训练效果如下，总体loss下降还是比较明显，效果也是非常不错
六、模型预测
实现预测函数，输入中文文本，输出模型翻译的英文结果：
def predict(model,text,device): # 编码输入文本 text=cn_tokenizer.encode(text=text) model.eval() with torch.no_grad(): src=torch.tensor(text,dtype=torch.long).unsqueeze(0).to(device) src_pad_mask=(src==model.zh_embedding.padding_idx) # 编码器编码 memory=model.encode(src,src_pad_mask) batch_size=src.shape[0] # 解码器初始输入（<start> token） decoder_input=torch.full([batch_size,1],en_tokenizer.start_index,device=device) generated=[] is_finished=torch.full([batch_size],False,device=device) # 逐token生成 for i in range(500): tgt_mask=model.transformer.generate_square_subsequent_mask(sz=decoder_input.shape[1]).to(device) decoder_output=model.decode(decoder_input,memory,tgt_mask,src_pad_mask) # 取概率最大的token next_token_index=torch.argmax(decoder_output[:,-1,:],dim=-1,keepdim=True) generated.append(next_token_index) decoder_input=torch.cat([decoder_input,next_token_index],dim=-1) # 判断是否生成<end> token is_finished |=(next_token_index.squeeze(1)==en_tokenizer.end_index) if is_finished.all(): break # 处理生成结果 generated_tensor=torch.cat(generated,dim=-1) generated_list=generated_tensor.tolist() for index,value in enumerate(generated_list): if en_tokenizer.end_index in value: end_pos=value.index(en_tokenizer.end_index) generated_list[index]=value[:end_pos] # 解码为文本 return en_tokenizer.decode(generated_list[0]) # 加载最优模型并预测 device=torch.device('cuda' if torch.cuda.is_available() else 'cpu') model=TranslationModel( zh_vocab_size=cn_tokenizer.vocab_size, en_vocab_size=en_tokenizer.vocab_size, zh_padding_idx=0, en_padding_idx=0 ).to(device) model.load_state_dict(torch.load('./best_model.pt')) # 测试预测 text='我是你爸爸' result=predict(model,text,device) print(f'输入：{text}') print(f'输出：{result}') # 输出：I'm your father.七、全部完整代码
import torch import pandas as pd from sklearn.model_selection import train_test_split from typing import List from torch import nn,optim from torch.nn.utils.rnn import pad_sequence from tqdm import tqdm import json from nltk.tokenize.treebank import TreebankWordDetokenizer,TreebankWordTokenizer from torch.utils.data import Dataset, DataLoader import math # 加载数据 data=pd.read_csv('./data/cmn.txt',sep='\t',header=None,usecols=[0,1],names=['en','zh']) # 划分训练集和验证集 train_df,test_df=train_test_split(data,test_size=0.2) # 构建分词器 class BaseTokenizer: # 构建tokenizer pad_index=0 unk_index=1 start_index=2 end_index=3 def __init__(self,vocab_list): self.vocab_list=vocab_list self.vocab_size = len(vocab_list) self.world2index={value:index for index,value in enumerate(vocab_list)} self.index2world={index:value for index,value in enumerate(vocab_list)} @staticmethod def tokenize(text:str)->List[str]: pass def encode(self,text:str,is_mark=False)->List[int]: tokens=self.tokenize(text) tokens_index=[self.world2index.get(token,self.unk_index) for token in tokens] if is_mark: tokens_index.insert(0,self.start_index) tokens_index.append(self.end_index) return tokens_index @classmethod def build_vocab( cls,sentences:List[str], unk_token:str='<unknown>', pad_token:str='<padding>', start_token:str='<start>', end_token:str='<end>', vocab_path:str='./vocab.json' ): vocab_set=set() for sentence in tqdm(sentences,desc='构建词表：'): vocab_set.update(cls.tokenize(sentence)) vocab_list = [pad_token, unk_token,start_token,end_token] + sorted(list(vocab_set)) vocab_dict={index:value for index,value in enumerate(vocab_list)} with open(vocab_path,'w',encoding='utf-8') as f: json.dump(vocab_dict, f, ensure_ascii=False, indent=2) @classmethod def read_vocab(cls,vocab_path:str='./vocab.json'): with open(vocab_path,'r',encoding='utf-8') as f: json_dict=json.load(f) sentences=[value for key,value in json_dict.items()] return cls(sentences) # 中文分词器 class ChinseeTokenizer(BaseTokenizer): @staticmethod def tokenize(text:str)->List[str]: return list(text) # 英文分词器 class EnglishTokenizer(BaseTokenizer): tokenizer=TreebankWordTokenizer() detokenizer=TreebankWordDetokenizer() @classmethod def tokenize(cls,text:str)->List[str]: return cls.tokenizer.tokenize(text) def decode(self,indexs:List[str])->str: tokens=[self.index2world.get(index,'<unknown>') for index in indexs] return self.detokenizer.detokenize(tokens) # 创建词表 ChinseeTokenizer.build_vocab(sentences=train_df['zh'].tolist(), vocab_path='./zh_vocab.json') EnglishTokenizer.build_vocab(sentences=train_df['en'].tolist(), vocab_path='./en_vocab.json') cn_tokenizer = ChinseeTokenizer.read_vocab('./zh_vocab.json') en_tokenizer = EnglishTokenizer.read_vocab('./en_vocab.json') train_df['en']=train_df['en'].apply(lambda x:en_tokenizer.encode(x,is_mark=True)) train_df['zh']=train_df['zh'].apply(lambda x:cn_tokenizer.encode(x)) test_df['en']=test_df['en'].apply(lambda x:en_tokenizer.encode(x,is_mark=True)) test_df['zh']=test_df['zh'].apply(lambda x:cn_tokenizer.encode(x)) train_df.to_json('./train.jsonl',orient='records',lines=True) test_df.to_json('./test.jsonl',orient='records',lines=True) # 构建Dataloader class TranslationDataset(Dataset): def __init__(self,path): self.data=pd.read_json(path,orient='records',lines=True).to_dict(orient='records') def __len__(self): return len(self.data) def __getitem__(self,index): input_tensor=torch.tensor(self.data[index]['zh'],dtype=torch.long) target_tensor=torch.tensor(self.data[index]['en'],dtype=torch.long) return input_tensor,target_tensor def collate_fn(batch): input_tensor=[tensor[0] for tensor in batch] target_tensor=[tensor[1] for tensor in batch] input_tensor=pad_sequence(sequences=input_tensor,batch_first=True,padding_value=0) target_tensor=pad_sequence(sequences=target_tensor,batch_first=True,padding_value=0) return input_tensor,target_tensor train_dataset=TranslationDataset('./train.jsonl') test_dataset=TranslationDataset('./test.jsonl') train_dataloader=DataLoader(train_dataset,batch_size=32,shuffle=True,collate_fn=collate_fn) test_dataloader=DataLoader(test_dataset,batch_size=32,collate_fn=collate_fn) # 构建位置编码 class PositionalEncoding(nn.Module): def __init__(self,max_len,dim_model): super(PositionalEncoding,self).__init__() pe=torch.zeros([max_len,dim_model],dtype=torch.float) for pos in range(max_len): for i in range(0,dim_model,2): pe[pos,i] = math.sin(pos/(10000**(i/dim_model))) pe[pos,i+1] = math.cos(pos/(10000**(i/dim_model))) self.register_buffer('pe',pe) def forward(self,x): seq_len=x.shape[1] part_pe=self.pe[0:seq_len] return x+part_pe # 构建模型 class TranslationModel(nn.Module): def __init__(self,zh_vocab_size,en_vocab_size,zh_padding_idx,en_padding_idx): super(TranslationModel,self).__init__() self.zh_embedding=nn.Embedding(num_embeddings=zh_vocab_size,embedding_dim=128,padding_idx=zh_padding_idx) self.en_embedding=nn.Embedding(num_embeddings=en_vocab_size,embedding_dim=128,padding_idx=en_padding_idx) self.position_encoding=PositionalEncoding(max_len=500,dim_model=128) self.transformer=nn.Transformer( d_model=128, nhead=8, num_encoder_layers=6, num_decoder_layers=6, batch_first=True, dropout=0.1, ) self.linear=nn.Linear(in_features=128,out_features=en_vocab_size) def forward(self,src,tgt,src_pad_mask,tgt_mask): memory=self.encode(src,src_pad_mask) outputs=self.decode(tgt,memory,tgt_mask,src_pad_mask) return outputs def encode(self,src,src_pad_mask): zh_embed=self.zh_embedding(src) zh_embed=self.position_encoding(zh_embed) memory=self.transformer.encoder(src=zh_embed,src_key_padding_mask=src_pad_mask) return memory def decode(self,tgt,memory,tgt_mask,memory_pad_mask): en_embed=self.en_embedding(tgt) en_embed=self.position_encoding(en_embed) output=self.transformer.decoder(tgt=en_embed,memory=memory,tgt_mask=tgt_mask,memory_key_padding_mask=memory_pad_mask) outputs=self.linear(output) return outputs device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model=TranslationModel( zh_vocab_size=cn_tokenizer.vocab_size, en_vocab_size=en_tokenizer.vocab_size, zh_padding_idx=0, en_padding_idx=0 ).to(device) # 定义模型超参数 epochs=5 lr=1e-4 loss_fn=nn.CrossEntropyLoss(ignore_index=en_tokenizer.pad_index) optimizer=optim.Adam(model.parameters(),lr=lr) # 模型训练和保存 best_loss=float('inf') for epoch in range(epochs): print(f'================第{epoch+1}轮=================') model.train() train_total_loss=0.0 for train_x,train_y in tqdm(train_dataloader,desc='训练'): src,tgt=train_x.to(device),train_y.to(device) decoder_inputs=tgt[:,:-1] decoder_targets=tgt[:,1:] src_pad_mask=(src==model.zh_embedding.padding_idx) tgt_mask=model.transformer.generate_square_subsequent_mask(sz=decoder_inputs.shape[1]).to(device) pred_y=model(src,decoder_inputs,src_pad_mask,tgt_mask) loss=loss_fn(pred_y.reshape(-1,pred_y.shape[-1]),decoder_targets.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() train_total_loss+=loss.item() model.eval() test_total_loss=0.0 with torch.no_grad(): for test_x, test_y in tqdm(test_dataloader,desc='验证'): src,tgt=test_x.to(device),test_y.to(device) decoder_inputs=tgt[:,:-1] decoder_targets=tgt[:,1:] src_pad_mask=(src==model.zh_embedding.padding_idx) tgt_mask=model.transformer.generate_square_subsequent_mask(sz=decoder_inputs.shape[1]).to(device) pred_y=model(src,decoder_inputs,src_pad_mask,tgt_mask) loss=loss_fn(pred_y.reshape(-1,pred_y.shape[-1]),decoder_targets.reshape(-1)) test_total_loss+=loss.item() avg_train_loss=train_total_loss/len(train_dataloader) avg_test_loss=test_total_loss/len(test_dataloader) print(f'训练平均loss:{avg_train_loss},验证平均loss:{avg_test_loss}') if test_total_loss<best_loss: best_loss=test_total_loss torch.save(model.state_dict(),'./best_model.pt') # 模型测试 def predict(model,text,device): text=cn_tokenizer.encode(text=text) model.eval() with torch.no_grad(): src=torch.tensor(text,dtype=torch.long).unsqueeze(0).to(device) src_pad_mask=(src==model.zh_embedding.padding_idx) memory=model.encode(src,src_pad_mask) batch_size=src.shape[0] decoder_input=torch.full([batch_size,1],en_tokenizer.start_index,device=device) generated=[] is_finished=torch.full([batch_size],False,device=device) for i in range(500): tgt_mask=model.transformer.generate_square_subsequent_mask(sz=decoder_input.shape[1]).to(device) decoder_output=model.decode(decoder_input,memory,tgt_mask,src_pad_mask) next_token_index=torch.argmax(decoder_output[:,-1,:],dim=-1,keepdim=True) generated.append(next_token_index) decoder_input=torch.cat([decoder_input,next_token_index],dim=-1) is_finished |=(next_token_index.squeeze(1)==en_tokenizer.end_index) if is_finished.all(): break generated_tensor=torch.cat(generated,dim=-1) generated_list=generated_tensor.tolist() for index,value in enumerate(generated_list): if en_tokenizer.end_index in value: end_pos=value.index(en_tokenizer.end_index) generated_list[index]=value[:end_pos] return en_tokenizer.decode(generated_list[0]) device=torch.device('cuda' if torch.cuda.is_available() else 'cpu') model=TranslationModel( zh_vocab_size=cn_tokenizer.vocab_size, en_vocab_size=en_tokenizer.vocab_size, zh_padding_idx=0, en_padding_idx=0 ).to(device) model.load_state_dict(torch.load('./best_model.pt')) text='我是你爸爸' result=predict(model,text,device) result
