transformer文本生成完整代码

Transformer模型是基于自注意力机制的深度学习架构，广泛应用于自然语言处理中的文本生成任务，如机器翻译和文本摘要。以下是一个简化版的PyTorch实现Transformer模型的代码片段，用于生成文本。请注意，这只是一个基本的示例，并未包含完整的训练过程，而是展示了如何使用预训练的模型进行条件生成。

import torch
from torch.nn import TransformerDecoder, TransformerDecoderLayer
from torch.nn.utils.rnn import pad_sequence
from torchtext.data.utils import get_tokenizer
from transformers import AutoTokenizer

# 假设我们已经加载了预训练的Transformer模型（如GPT-2）
model_name = 'gpt2'
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def generate_text(prompt, num_words=50):
    input_ids = torch.tensor([tokenizer.encode(prompt, return_tensors='pt')]).to(device)
    decoder_layer = TransformerDecoderLayer(d_model=768, nhead=12, dim_feedforward=2048)
    decoder = TransformerDecoder(decoder_layer, num_layers=6, device=device)
    
    # 初始化输入和前向传播
    memory = None
    output = input_ids
    for _ in range(num_words):
        output = decoder(output, memory)
        predicted_id = torch.argmax(output[:, -1], dim=-1).item()
        tokenized_text = tokenizer.decode(predicted_id)
        print(tokenized_text, end=' ')

        # 更新输入，添加上刚刚生成的词
        if tokenized_text == '':
            break
        input_ids = torch.cat((input_ids, torch.tensor([[predicted_id]])), dim=1)

# 使用生成器函数
generate_text('Hello, ', num_words=20)