Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2from typing import List, Dict
3import torch
4
5def create_system(prompt_data, bot_name, gender):
6 system_prompt = (
7 f"You need to play the role of {bot_name}, who is {gender}. Here is the basic information:\n"
8 f"- Facts: {prompt_data.get('Fact', 'No fact available')}\n"
9 f"- Persona: {prompt_data.get('Head', 'No persona available')}\n"
10 f"- Brief: {prompt_data.get('Brief', 'No brief available')}\n"
11 "Please reply to the current User with the character traits and Chat History"
12 )
13
14 history = f"- Chat history: {', '.join(prompt_data.get('History', []))}\n"
15 return system_prompt, history
16
17model_name = "lanlanlan123/RoleLLM_Ministral_8b"
18tokenizer = AutoTokenizer.from_pretrained(model_name)
19model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
20
21character1 = {
22 "Fact": "Leo is an amateur astronomer. He spends most of his free time stargazing and has a small telescope in his backyard.",
23 "Head": "A friendly and enthusiastic guy who is always eager to share his knowledge about the universe.",
24 "Brief": "A software engineer by day and an astronomy enthusiast by night.",
25 "History": ["User: What's your favorite constellation? Leo: I love Orion. It's so easy to spot in the winter sky."]
26}
27bot_name_1 = "Leo"
28gender_1 = "male"
29
30system_prompt, history = create_system(character1, bot_name_1, gender_1)
31messages = [
32 {"role": "system",
33 "content": system_prompt},
34 {"role": "user", "content": history + "\nUser: but I dont like Orion"}
35]
36
37formatted_input = ""
38for message in messages:
39 if message["role"] == "system":
40 formatted_input += message["content"] + "\n\n"
41 elif message["role"] == "user":
42 formatted_input += f"[INST]{message['content']}[/INST]"
43
44inputs = tokenizer(formatted_input + f"\n{bot_name_1}: ", return_tensors="pt").to(model.device)
45
46outputs = model.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
47
48input_length = inputs.input_ids.size(1)
49response_tokens = outputs[0][input_length:]
50
51response = tokenizer.decode(response_tokens, skip_special_tokens=True)
52
53print("input:", formatted_input)
54print("output:", response)| 模型 | 均分 | 角色一致性 | 对话能力 | 角色扮演吸引力 |
|---|---|---|---|---|
| Qwen-14B | 3.016 | 2.649 | 3.542 | 2.858 |
| GPT-4 | 3.006 | 2.697 | 3.448 | 2.873 |
| Ministral-8b-lora-sft | 3.01 | 2.4725 | 3.75 | 2.8175 |
| Xingchen | 2.991 | 2.595 | 3.646 | 2.732 |
| XVERSE-7B | 2.963 | 2.564 | 3.554 | 2.772 |
| CharacterGLM | 2.937 | 2.493 | 3.623 | 2.695 |
| ChatGLM3-6B | 2.898 | 2.556 | 3.399 | 2.739 |
| Qwen-7B | 2.849 | 2.54 | 3.327 | 2.679 |
| Ministral-8b-原版 | 2.77 | 2.24 | 3.48 | 2.59 |
| GPT-3.5 | 2.381 | 2.101 | 2.749 | 2.293 |