从零到一：如何训练简版生成式GPT模型，快速实现创意写作

从零到一：如何训练简版生成式GPT模型，快速实现创意写作

一、从零到一：你的第一个GPT诞生记

震撼对比：使用相同训练数据（全唐诗30万首）

• 传统RNN生成：“春风吹又生，花落知多少”（模板化）
• 简版GPT生成：“墨染江南烟雨楼，一蓑孤舟任水流。青山不语斜阳暮，白鹭惊飞入画轴”（创意性）

 

二、自回归机制

⾃回归模型（Autoregressive Model）是⽣成式模型的⼀种特例，它预测的新⽬标值是基于前⾯若⼲个已⽣成值的。⾃回归模型在时间序列分析、语⾳信号处理、⾃然语⾔处理等领域有⼴泛应⽤。在序列⽣成问题中，⾃回归模型特别重要，⽐如在机器翻译、⽂本⽣成、语⾳合成等任务中，Transformer的解码器、GPT等模型就是基于⾃回归原理的。

⽤⾃回归机制来逐词⽣成翻译结果。

还是使⽤同样的中英翻译数据集，还是使⽤Transformer模型，这⾥我们只是加⼀个⽤贪婪搜索进⾏⽣成式解码的函数，然后在测试过程中调⽤这个函数重新测试。

之前的数据见：Transformer：GPT背后的造脑工程全解析（含手搓过程）-CSDN博客

2.1定义贪婪解码器函数

# 定义贪婪解码器函数
def greedy_decoder(model, enc_input, start_symbol):
    # 对输入数据进行编码，并获得编码器输出以及自注意力权重
    enc_outputs, enc_self_attns = model.encoder(enc_input)    
    # 初始化解码器输入为全零张量，大小为 (1, 5)，数据类型与 enc_input 一致
    dec_input = torch.zeros(1, 5).type_as(enc_input.data)    
    # 设置下一个要解码的符号为开始符号
    next_symbol = start_symbol    
    # 循环 5 次，为解码器输入中的每一个位置填充一个符号
    for i in range(0, 5):
        # 将下一个符号放入解码器输入的当前位置
        dec_input[0][i] = next_symbol        
        # 运行解码器，获得解码器输出、解码器自注意力权重和编码器 - 解码器注意力权重
        dec_output, _, _ = model.decoder(dec_input, enc_input, enc_outputs)        
        # 将解码器输出投影到目标词汇空间
        projected = model.projection(dec_output)        
        # 找到具有最高概率的下一个单词
        prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]
        next_word = prob.data[i]        
        # 将找到的下一个单词作为新的符号
        next_symbol = next_word.item()        
    # 返回解码器输入，它包含了生成的符号序列
    dec_outputs = dec_input
    return dec_outputs

greedy_decoder 实现了贪婪解码算法的过程，其中：

• 每次解码时，模型选择概率最大的单词作为下一个符号。
• 这个过程在目标序列的每个位置进行，直到生成一个完整的目标序列（此处假设目标序列长度为 5）。

贪婪解码是一种简单的解码策略，它总是选择当前时刻最可能的单词，而不考虑全局最优解。虽然这种方法快速，但可能会陷入局部最优解，导致生成的序列不尽如人意。在更复杂的情况下，可能需要使用如 束搜索（Beam Search）等策略来改进生成的质量。

# 用贪婪解码器生成翻译文本
enc_inputs, dec_inputs, target_batch = corpus.make_batch(batch_size=1, test_batch=True) 
# 使用贪婪解码器生成解码器输入
greedy_dec_input = greedy_decoder(model, enc_inputs, start_symbol=corpus.tgt_vocab['<sos>'])
# 将解码器输入转换为单词序列
greedy_dec_output_words = [corpus.tgt_idx2word[n.item()] for n in greedy_dec_input.squeeze()]
# 打印编码器输入和贪婪解码器生成的文本
enc_inputs_words = [corpus.src_idx2word[code.item()] for code in enc_inputs[0]]
print(enc_inputs_words, '->', greedy_dec_output_words)

['我', '爱', '学习', '人工智能', '<pad>'] -> ['<sos>', 'I', 'love', 'studying', 'AI']

 

三、构建GPT模型

3.1 搭建GPT模型（解码器）

GPT只使⽤了Transformer的解码器部分，其关键组件如下图所示。

搭建GPT模型的代码的关键组件如下。

组件1 多头⾃注意⼒：通过ScaledDotProductAttention类实现缩放点积注意⼒机制，然后通过MultiHeadAttention类实现多头⾃注意⼒机制。

组件2 逐位置前馈⽹络：通过PoswiseFeedForwardNet类实现逐位置前馈⽹络。

组件3 正弦位置编码表：通过get_sin_code_table函数⽣成正弦位置编码表。

组件4 填充掩码：通过get_attn_pad_mask函数为填充token⽣成注意⼒掩码，避免注意⼒机制关注⽆⽤的信息。

组件5 后续掩码：通过get_attn_subsequent_mask函数为后续token（当前位置后⾯的信息）⽣成注意⼒掩码，避免解码器中的注意⼒机制“偷窥”未来的⽬标数据。

组件6 解码器层：通过DecoderLayer类定义解码器的单层。

组件7 解码器：通过Decoder类定义Transformer模型的完整解码器部分。

组件8 GPT：在解码器的基础上添加⼀个投影层，将解码器输出的特征向量转换为预测结果，实现⽂本⽣成。

之前在手搓transfomer中已经讲过前五个组件，这次从第六个组件开始讲解。

 

3.2 解码器层类

因为GPT模型没有编码器组件，也不需要来⾃编码器的输出，因此GPT解码器的实现更简洁。GPT模型也省略了编码器-解码器注意⼒机制，因此模型的训练速度更快。

# 定义解码器层类
class DecoderLayer(nn.Module):
    def __init__(self):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention()  # 多头自注意力层
        self.feed_forward = PoswiseFeedForwardNet()  # 逐位置前馈网络层
        self.norm1 = nn.LayerNorm(d_embedding)  # 第一个层归一化
        self.norm2 = nn.LayerNorm(d_embedding)  # 第二个层归一化
    def forward(self, dec_inputs, attn_mask=None):
        # 使用多头自注意力处理输入
        attn_output, _ = self.self_attn(dec_inputs, dec_inputs, dec_inputs, attn_mask)
        # 将注意力输出与输入相加并进行第一个层归一化
        norm1_outputs = self.norm1(dec_inputs + attn_output)
        # 将归一化后的输出输入到位置前馈神经网络
        ff_outputs = self.feed_forward(norm1_outputs)
        # 将前馈神经网络输出与第一次归一化后的输出相加并进行第二个层归一化
        dec_outputs = self.norm2(norm1_outputs + ff_outputs)
        return dec_outputs # 返回解码器层输出

GPT解码器层的构造⽐Transformer的解码器层简单，仅包含⼀个多头⾃注意⼒层MultiHeadAttention和⼀个逐位置前馈⽹络层PosFeedForwardNet，后⾯接了两个层归⼀化nn.LayerNorm。

主要是两个归一化层

解码器层中，两个层归⼀化的作⽤如下。

■ 第⼀个层归⼀化norm1：在多头⾃注意⼒self_attn处理后，将注意⼒输出attn_output与原始输⼊dec_inputs相加。这种加和操作实现了残差连接，可以加速梯度反向传播，有助于训练深层⽹络。将相加后的结果进⾏层归⼀化。层归⼀化对输⼊进⾏标准化处理，使其具有相同的均值和⽅差。这有助于减少梯度消失或梯度爆炸问题，从⽽提⾼模型训练的稳定性。

■ 第⼆个层归⼀化norm2：在逐位置前馈⽹络feed_forward处理后，将前馈神经⽹络输出ff_outputs与第⼀个层归⼀化输出norm1_outputs相加。这⾥同样实现了残差连接。将相加后的结果进⾏层归⼀化。这⼀步骤的⽬的与第⼀个层归⼀化相同，即标准化输⼊数据，以提⾼训练稳定性。

 

3.3 解码器类

#  定义解码器类
n_layers = 6  # 设置 Decoder 的层数
class Decoder(nn.Module):
    def __init__(self, vocab_size, max_seq_len):
        super(Decoder, self).__init__()
        # 词嵌入层（参数为词典维度）
        self.src_emb = nn.Embedding(vocab_size, d_embedding)  
        # 位置编码层（参数为序列长度）
        self.pos_emb = nn.Embedding(max_seq_len, d_embedding)
        # 初始化 N 个解码器层       
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)]) 
    def forward(self, dec_inputs):        
        # 创建位置信息
        positions = torch.arange(len(dec_inputs), device=dec_inputs.device).unsqueeze(-1)
        # 将词嵌入与位置编码相加
        inputs_embedding = self.src_emb(dec_inputs) + self.pos_emb(positions)
        # 生成自注意力掩码
        attn_mask = get_attn_subsequent_mask(inputs_embedding).to(device)
        # 初始化解码器输入，这是第一层解码器层的输入 
        dec_outputs =  inputs_embedding 
        for layer in self.layers:
            # 将输入数据传递给解码器层，并返回解码器层的输出，作为下一层的输入
            dec_outputs = layer(dec_outputs, attn_mask) 
        return dec_outputs # 返回解码器输出

GPT解码器的结构⽐Transformer解码器的结构简单，因为GPT是⼀个单向⽣成式模型，只关注⽣成⽂本⽽不关注源⽂本。GPT不需要实现编码器-解码器注意⼒的部分，仅接收解码器的输⼊，然后进⾏词嵌⼊和位置编码，并将⼆者相加，继⽽⽣成后续⾃注意⼒掩码，来保证每个位置只能看到当前位置之前的信息，以保持⽣成⽂本的⾃回归特性。最后把嵌⼊向量和掩码信息传递给解码器层，并⾏处理，并接收结果向量dec_outputs，然后把它返回给GPT模型。

一句话说的话就是，解码器只管生成文本。

 

3.4 GPT组件

# 定义 GPT 模型
class GPT(nn.Module):
    def __init__(self, vocab_size, max_seq_len):
        super(GPT, self).__init__()
        self.decoder = Decoder(vocab_size, max_seq_len) # 解码器，用于学习文本生成能力
        self.projection = nn.Linear(d_embedding, vocab_size)  # 全连接层，输出预测结果
    def forward(self, dec_inputs):        
        dec_outputs = self.decoder(dec_inputs) # 将输入数据传递给解码器
        logits = self.projection(dec_outputs) # 传递给全连接层以生成预测
        return logits # 返回预测结py
        果

在这个简化版的GPT模型中：解码器类负责学习⽂本⽣成能⼒；⼀个全连接层将解码器输出的特征向量映射到⼀个概率分布，表示⽣成每个单词的概率logits，⽤于将解码器的输出转换为与词汇表⼤⼩相匹配的预测结果。

GPT模型仅包含解码器部分，没有编码器部分。因此，它更适⽤于⽆条件⽂本⽣成任务，⽽不是类似机器翻译或问答等需要编码器-解码器结构的任务。

 

四、完成文本生成任务

4.1 构建文本生成任务的数据集

我们要给GPT准备⼀个训练语料库。这个语料库是由现实中存在的⽂字组成的。当然，⽐起维基百科等⼤型语料库，我们的语料库中的数据⽐较少，你可以把它看成⼈类语料库的⼀个缩影。

把这个语料库存储在⽂件lang.txt中，等待程序读取。

# 构建语料库
from collections import Counter
class LanguageCorpus:
    def __init__(self, sentences):
        self.sentences = sentences
        # 计算语言的最大句子长度，并加 2 以容纳特殊符号 <sos> 和 <eos>
        self.seq_len = max([len(sentence.split()) for sentence in sentences]) + 2
        self.vocab = self.create_vocabulary() # 创建源语言和目标语言的词汇表
        self.idx2word = {v: k for k, v in self.vocab.items()} # 创建索引到单词的映射
    def create_vocabulary(self):
        vocab = {'<pad>': 0, '<sos>': 1, '<eos>': 2}
        counter = Counter()
        # 统计语料库的单词频率
        for sentence in self.sentences:
            words = sentence.split()
            counter.update(words)
        # 创建词汇表，并为每个单词分配一个唯一的索引
        for word in counter:
            if word not in vocab:
                vocab[word] = len(vocab)
        return vocab
    def make_batch(self, batch_size, test_batch=False):
        input_batch, output_batch = [], [] # 初始化批数据
        sentence_indices = torch.randperm(len(self.sentences))[:batch_size] # 随机选择句子索引
        for index in sentence_indices:
            sentence = self.sentences[index]
            # 将句子转换为索引序列
            seq = [self.vocab['<sos>']] + [self.vocab[word] for word in sentence.split()] + [self.vocab['<eos>']]
            seq += [self.vocab['<pad>']] * (self.seq_len - len(seq)) # 对序列进行填充
            # 将处理好的序列添加到批次中
            input_batch.append(seq[:-1])
            output_batch.append(seq[1:])
        return torch.LongTensor(input_batch), torch.LongTensor(output_batch)

这个类的主要功能是创建词汇表、将句⼦转换为索引序列、⽣成批次数据等，其中最重要的是make_batch⽅法中⽣成批次数据时的“向右位移”操作，这是训练⽣成式语⾔模型的关键所在。

查看一些语料库的数据：

with open("lang.txt", "r") as file: # 从文件中读入语料
    sentences = [line.strip() for line in file.readlines()]
corpus = LanguageCorpus(sentences) # 创建语料库
vocab_size = len(corpus.vocab) # 词汇表大小
max_seq_len = corpus.seq_len # 最大句子长度（用于设置位置编码）
print(f" 语料库词汇表大小 : {vocab_size}") # 打印词汇表大小
print(f" 最长句子长度 : {max_seq_len}") # 打印最大序列长

 语料库词汇表大小 : 133

 最长句子长度 : 17

 

4.2 训练过程中的自回归

import torch.optim as optim # 导入优化器
device = "cuda" if torch.cuda.is_available() else "cpu" # 设置设备
model = GPT(vocab_size, max_seq_len).to(device) # 创建 GPT 模型实例
criterion = nn.CrossEntropyLoss() # 损失函数
optimizer = optim.Adam(model.parameters(), lr=0.0001) # 优化器
epochs = 500 # 训练轮次
for epoch in range(epochs):  # 训练 epochs 轮
    optimizer.zero_grad() # 梯度清零
    inputs, targets = corpus.make_batch(batch_size) # 创建训练数据
    inputs, targets = inputs.to(device), targets.to(device)
    outputs = model(inputs) # 获取模型输出 
    loss = criterion(outputs.view(-1, vocab_size), targets.view(-1)) # 计算损失
    if (epoch + 1) % 100 == 0: # 打印损失
        print(f"Epoch: {epoch + 1:04d} cost = {loss:.6f}")
    loss.backward() # 反向传播
    optimizer.step() # 更新参数

 

4.3 文本生成的自回归(贪婪搜索)

# 测试文本生成
def generate_text(model, input_str, max_len=50):
    model.eval()  # 将模型设置为评估（测试）模式，关闭 dropout 和 batch normalization 等训练相关的层
    # 将输入字符串中的每个 token 转换为其在词汇表中的索引
    input_tokens = [corpus.vocab[token] for token in input_str]
    # 创建一个新列表，将输入的 tokens 复制到输出 tokens 中 , 目前只有输入的词
    output_tokens = input_tokens.copy()
    with torch.no_grad():  # 禁用梯度计算，以节省内存并加速测试过程
        for _ in range(max_len):  # 生成最多 max_len 个 tokens
            # 将输出的 token 转换为 PyTorch 张量，并增加一个代表批次的维度 [1, len(output_tokens)]
            inputs = torch.LongTensor(output_tokens).unsqueeze(0).to(device)
            outputs = model(inputs) # 输出 logits 形状为 [1, len(output_tokens), vocab_size]
            # 在最后一个维度上获取 logits 中的最大值，并返回其索引（即下一个 token）
            _, next_token = torch.max(outputs[:, -1, :], dim=-1)            
            next_token = next_token.item() # 将张量转换为 Python 整数            
            if next_token == corpus.vocab["<eos>"]:
                break # 如果生成的 token 是 EOS（结束符），则停止生成过程           
            output_tokens.append(next_token) # 将生成的 tokens 添加到 output_tokens 列表
    # 将输出 tokens 转换回文本字符串
    output_str = " ".join([corpus.idx2word[token] for token in output_tokens])
    return output_str
input_str = ["Python"] # 输入一个词：Python
generated_text = generate_text(model, input_str) # 模型跟着这个词生成后续文本
print(" 生成的文本 :", generated_text) # 打印预测文本

 生成的文本 : Python is a popular programming language.

 

五、使用WikiText2数据集训练Wiki-GPT模型

使⽤⼀个从互联⽹中收集真实语料的数据集WikiText 。

WikiText数据集是从维基百科上经过验证的精选⽂章集中提取的超过1亿个标记的数据集合。让我们⽤这个更真实的语料数据集，来看⼀看现实世界的模型是如何训练出来的。

■ ⽬前NLP领域中常⽤数据集的导⼊（通过Torchtext库）、设计（创建PyTorch Dataset）、加载（使⽤PyTorch Data Loader），以及如何将数据集中的数据转换为

我们的GPT模型可以读取的格式（通过Torchtext的分词⼯具Tokenizer）。

■ 如何对模型的效能进⾏测试。

之前的数据集都很⼩，没有拆分成训练数据集和测试数据集的必要，⽽现在，⽤这个真实的、包含上万条语料的数据集，就可以⽤其中的⼀部分数据来测试模型的效能。这样，在多轮训练的过程中，我们就可以选择测试集上得分最⾼的模型。

 

5.1 用WikiText2构建Dataset和DataLoader

from torchtext.datasets import WikiText2 # 导入WikiText2
from torchtext.data.utils import get_tokenizer # 导入Tokenizer分词工具
from torchtext.vocab import build_vocab_from_iterator # 导入Vocabulary工具
from torch.utils.data import DataLoader, Dataset # 导入Pytorch的DataLoader和Dataset

tokenizer = get_tokenizer("basic_english") # 定义数据预处理所需的tokenizer

train_iter = WikiText2(split='train') # 加载WikiText2数据集的训练部分
valid_iter = WikiText2(split='valid') # 加载WikiText2数据集的验证部分

# 定义一个生成器函数，用于将数据集中的文本转换为tokens
def yield_tokens(data_iter):
    for item in data_iter:
        yield tokenizer(item)

# 创建词汇表，包括特殊tokens："<pad>", "<sos>", "<eos>"
vocab = build_vocab_from_iterator(yield_tokens(train_iter), 
                                  specials=["<pad>", "<sos>", "<eos>"])
vocab.set_default_index(vocab["<pad>"])

# 打印词汇表信息
print("词汇表大小:", len(vocab))
print("词汇示例(word to index):", 
      {word: vocab[word] for word in ["<pad>", "<sos>", "<eos>", "the", "apple"]})

from torch.utils.data import Dataset # 导入Dataset
max_seq_len = 256 # 设置序列的最大长度

# 定义一个处理WikiText2数据集的自定义数据集类
class WikiDataset(Dataset):
    def __init__(self, data_iter, vocab, max_len=max_seq_len):
        self.data = []        
        for sentence in data_iter: # 遍历数据集，将文本转换为tokens
            # 对每个句子进行tokenization，并截取长度为max_len-2，为<sos>和<eos>留出空间
            tokens = tokenizer(sentence)[:max_len - 2]
            tokens = [vocab["<sos>"]] + vocab(tokens) + [vocab["<eos>"]] # 添加<sos>和<eos>            
            self.data.append(tokens) # 将处理好的tokens添加到数据集中
    
    def __len__(self): # 定义数据集的长度
        return len(self.data)    
    
    def __getitem__(self, idx): # 定义数据集的索引方法 (即抽取数据条目)        
        source = self.data[idx][:-1] # 获取当前数据，并将<eos>移除，作为source        
        target = self.data[idx][1:] # 获取当前数据，并将<sos>移除，作为target（右移1位）       
        return torch.tensor(source), torch.tensor(target) # 转换为tensor并返回

train_dataset = WikiDataset(train_iter, vocab) # 创建训练数据集
valid_dataset = WikiDataset(valid_iter, vocab) # 创建验证数据集
print(f"Dataset数据条目: {len(train_dataset)}")
sample_source, sample_target = train_dataset[100]
print(f"输入序列张量样例: {sample_source}")
print(f"目标序列张量样例: {sample_target}")
decoded_source = ' '.join(vocab.lookup_tokens(sample_source.tolist()))
decoded_target = ' '.join(vocab.lookup_tokens(sample_target.tolist()))
print(f"输入序列样例文本: {decoded_source}")
print(f"目标序列样例文本: {decoded_target}")

from torch.utils.data import DataLoader # 导入Dataloader
# 定义pad_sequence函数，用于将一批序列补齐到相同长度
def pad_sequence(sequences, padding_value=0, length=None):
    # 计算最大序列长度，如果length参数未提供，则使用输入序列中的最大长度
    max_length = max(len(seq) for seq in sequences) if length is None else length    
    # 创建一个具有适当形状的全零张量，用于存储补齐后的序列
    result = torch.full((len(sequences), max_length), padding_value, dtype=torch.long)    
    # 遍历序列，将每个序列的内容复制到结果张量中
    for i, seq in enumerate(sequences):
        end = len(seq)
        result[i, :end] = seq[:end]
    return result

# 定义collate_fn函数，用于将一个批次的数据整理成适当的形状
def collate_fn(batch):
    # 从批次中分离源序列和目标序列
    sources, targets = zip(*batch)    
    # 计算批次中的最大序列长度
    max_length = max(max(len(s) for s in sources), max(len(t) for t in targets))    
    # 使用pad_sequence函数补齐源序列和目标序列
    sources = pad_sequence(sources, padding_value=vocab["<pad>"], length=max_length)
    targets = pad_sequence(targets, padding_value=vocab["<pad>"], length=max_length)    
    # 返回补齐后的源序列和目标序列
    return sources, targets

# 创建一个训练数据加载器，使用自定义的collate_fn函数
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, 
                              shuffle=True, collate_fn=collate_fn)
# 创建一个验证数据加载器，使用自定义的collate_fn函数
valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size,
                              shuffle=False, collate_fn=collate_fn)

5.2 训练Wik-GPT

import torch.optim as optim  # 导入优化器
device = "cuda" if torch.cuda.is_available() else "cpu"  # 设置设备
model = GPT(len(vocab), max_seq_len).to(device)  # 创建GPT模型实例
criterion = nn.CrossEntropyLoss(ignore_index=vocab["<pad>"])
optimizer = optim.Adam(model.parameters(), lr=0.0001)  # 优化器
epochs = 2  # 训练轮次

for epoch in range(epochs):
    epoch_loss = 0
    for batch_idx, (source, target) in enumerate(train_dataloader): # 用Dataloader加载数据
        inputs, targets = source.to(device), target.to(device)
        optimizer.zero_grad()  # 梯度清零
        outputs = model(inputs)  # 获取模型输出
        loss = criterion(outputs.view(-1, len(vocab)), targets.view(-1))  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        epoch_loss += loss.item()        
        if (batch_idx + 1) % 500 == 0: # 每500个批次打印一次损失
            print(f"Batch {batch_idx + 1}/{len(train_dataloader)}, Loss: {loss.item()}")    
    epoch_loss /= len(train_dataloader) # 每轮打印一次损失
    print(f"Epoch {epoch + 1}/{epochs}, Average Loss: {epoch_loss}")

5.3 生成文本

# Replace 'model_timestamp.pt' with your saved model's filename
model.load_state_dict(torch.load('trained_model_2023-05-05_14-08-24.pt'))
# 测试文本生成
def generate_text_greedy_search(model, input_str, max_len=50):
    model.eval()  # 将模型设置为评估（测试）模式，关闭dropout和batch normalization等训练相关的层
    # 将输入字符串中的每个Token 转换为其在词汇表中的索引
    input_tokens = [vocab[token] for token in input_str.split()]
    # 创建一个新列表，将输入的Token复制到输出Token中,目前只有输入的词
    output_tokens = input_tokens.copy()
    with torch.no_grad():  # 禁用梯度计算，以节省内存并加速测试过程
        for _ in range(max_len):  # 生成最多max_len个Token
            # 将输出token转换为 PyTorch张量，并增加一个代表批次的维度[1, len(output_tokens)]
            inputs = torch.LongTensor(output_tokens).unsqueeze(0).to(device)
            outputs = model(inputs) # 输出 logits形状为[1, len(output_tokens), vocab_size]
            logits = outputs[:, -1, :] # 只关心最后一个时间步（即最新生成的token）的logits
            # 在最后一个维度上获取logits中的最大值，并返回其索引（即下一个Token）
            _, next_token = torch.max(logits, dim=-1)            
            next_token = next_token.item() # 将张量转换为Python整数            
            if next_token == vocab["<eos>"]:
                break # 如果生成的Token是 EOS（结束符），则停止生成过程           
            output_tokens.append(next_token) # 将生成的Token添加到output_tokens列表
    # 将输出Token转换回文本字符串
    output_str = " ".join([vocab.get_itos()[token] for token in output_tokens
                           if vocab.get_itos()[token] != "<pad>" and vocab.get_itos()[token] != "<unk>" ])
    return output_str

input_str = "how are you" # 输入一个词：Python
generated_text = generate_text_greedy_search(model, input_str) # 模型跟着这个字生成后续文本
print("生成的文本:", generated_text) # 打印预测文本

# 定义集束搜索的函数
def generate_text_beam_search(model, input_str, max_len=50, beam_width=5):
    model.eval()  # 将模型设置为评估（测试）模式，关闭dropout和batch normalization等训练相关的层
    # 将输入字符串中的每个token 转换为其在词汇表中的索引
    input_tokens = [vocab[token] for token in input_str.split()]
    # 创建一个列表，用于存储候选序列
    candidates = [(input_tokens, 0.0)]
    with torch.no_grad():  # 禁用梯度计算，以节省内存并加速测试过程
        for _ in range(max_len):  # 生成最多max_len个tokens
            new_candidates = []
            for candidate, candidate_score in candidates:
                inputs = torch.LongTensor(candidate).unsqueeze(0).to(device)
                outputs = model(inputs) # 输出 logits形状为[1, len(output_tokens), vocab_size]
                logits = outputs[:, -1, :] # 只关心最后一个时间步（即最新生成的token）的logits
                # 找到具有最高分数的前beam_width个tokens
                scores, next_tokens = torch.topk(logits, beam_width, dim=-1)
                final_results = [] # 初始化输出序列
                for score, next_token in zip(scores.squeeze(), next_tokens.squeeze()):
                    new_candidate = candidate + [next_token.item()]
                    new_score = candidate_score - score.item()  # 使用负数，因为我们需要降序排列
                    if next_token.item() == vocab["<eos>"]:
                        # 如果生成的token是EOS（结束符），将其添加到最终结果中
                        final_results.append((new_candidate, new_score))
                    else:
                        # 将新生成的候选序列添加到新候选列表中
                        new_candidates.append((new_candidate, new_score))
            # 从新候选列表中选择得分最高的beam_width个序列
            candidates = sorted(new_candidates, key=lambda x: x[1])[:beam_width]
    # 选择得分最高的候选序列
    best_candidate, _ = sorted(candidates, key=lambda x: x[1])[0]
    # 将输出 token 转换回文本字符串
    output_str = " ".join([vocab.get_itos()[token] for token in best_candidate if vocab.get_itos()[token] != "<pad>"])
    return output_str

model.load_state_dict(torch.load('trained_model_2023-05-05_14-08-24.pt')) # 加载模型
input_str = "my name"  # 输入几个词
generated_text = generate_text_beam_search(model, input_str)  # 模型跟着这些词生成后续文本
print("生成的文本:", generated_text)  # 打印生成的文本

生成的文本: my name was also used in 1897 by lucasfilm games in the common by lucasfilm games in the common by lucasfilm games in the common by lucasfilm games in the common by lucasfilm games in the common by lucasfilm games in the common by lucasfilm games in the common by lucasfilm games