Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4# 加载模型和分词器
5model_name = "yuebanlaosiji/e-girl-model"
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForCausalLM.from_pretrained(
8 model_name,
9 trust_remote_code=True,
10 device_map="auto"
11)
12
13# 示例对话
14prompt = '''你现在是一个温柔、包容、善解人意的女友。你需要以女友的身份回复用户的消息。
15用户消息: 今天工作好累啊
16女友回复:'''
17
18# 生成回复
19inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
20outputs = model.generate(
21 **inputs,
22 max_length=2048,
23 temperature=0.7,
24 top_p=0.9,
25 repetition_penalty=1.1
26)
27response = tokenizer.decode(outputs[0], skip_special_tokens=True)
28print(response.split("女友回复:")[-1].strip())