1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3# 加载模型和tokenizer
4model = AutoModelForCausalLM.from_pretrained("your-username/model-name")
5tokenizer = AutoTokenizer.from_pretrained("your-username/model-name")
6
7# 生成文本
8def generate_story(prompt, max_length=500):
9 inputs = tokenizer(prompt, return_tensors="pt")
10 outputs = model.generate(
11 **inputs,
12 max_length=max_length,
13 temperature=0.7,
14 do_sample=True,
15 top_p=0.9,
16 repetition_penalty=1.1
17 )
18 return tokenizer.decode(outputs[0], skip_special_tokens=True)
19
20# 示例
21story = generate_story("请写一个关于友情的故事：")
22print(story)
