古诗生成代码

以下是一个基于GPT-2模型的古诗生成代码示例：

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

def generate_poem(input_text):
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    output = model.generate(input_ids, max_length=100, num_return_sequences=1, do_sample=True)
    generated_poem = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_poem

input_text = '寒江独钓'
generated_poem = generate_poem(input_text)
print(generated_poem)

这个代码使用了Hugging Face的transformers库，首先加载了GPT-2的tokenizer和预训练模型，然后定义了一个生成古诗的函数generate_poem，它接受一个输入文本作为生成古诗的前缀，调用模型的generate方法生成一首100字长的古诗（num_return_sequences=1表示只生成一首），最后返回生成的古诗字符串。在这个示例代码中，我们以“寒江独钓”作为前缀生成了一首古诗。