1from transformers import GPT2LMHeadModel, GPT2Tokenizer
2
3tokenizer = GPT2Tokenizer.from_pretrained("xingyu1996/chinese-poems-gpt2")
4model = GPT2LMHeadModel.from_pretrained("xingyu1996/chinese-poems-gpt2")
5
6# 生成诗歌
7input_text = "作者：李白\n标题：望月\n\n"
8input_ids = tokenizer.encode(input_text, return_tensors='pt')
9
10output = model.generate(
11 input_ids,
12 max_length=100,
13 num_return_sequences=1,
14 no_repeat_ngram_size=2,
15 top_k=50,
16 top_p=0.95,
17 temperature=0.7
18)
19
20generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
21print(generated_text)
