Views
No views yet
Mario-RC/aif-emotional-model.THUDM/glm-4-9b-chat-1mmario-rc/emotional-rlaif-ppo-glm-4-9b-chat-1mglm4dialoguesdialoguesaif_annotations preference pairs1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from peft import PeftModel
4
5base_model_id = "THUDM/glm-4-9b-chat-1m"
6adapter_id = "mario-rc/emotional-rlaif-ppo-glm-4-9b-chat-1m"
7
8tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)
9model = AutoModelForCausalLM.from_pretrained(
10 base_model_id,
11 device_map="auto",
12 torch_dtype=torch.bfloat16,
13 trust_remote_code=True,
14)
15model = PeftModel.from_pretrained(model, adapter_id)
16model.eval()
17
18messages = [
19 {"role": "user", "content": "I feel overwhelmed today. Can you respond with empathy?"}
20]
21
22inputs = tokenizer.apply_chat_template(
23 messages,
24 add_generation_prompt=True,
25 return_tensors="pt",
26).to(model.device)
27
28with torch.no_grad():
29 outputs = model.generate(
30 inputs,
31 max_new_tokens=256,
32 do_sample=True,
33 temperature=0.7,
34 top_p=0.9,
35 )
36
37print(tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True))1import random
2import sys
3
4import torch
5from peft import AutoPeftModelForCausalLM
6from transformers import AutoTokenizer
7
8MODEL_ID = "mario-rc/emotional-rlaif-ppo-glm-4-9b-chat-1m"
9
10
11def get_turn_markers():
12 return {
13 'bos': '[gMASK]<sop>',
14 'user_start': '<|user|>\n',
15 'user_end': '\n',
16 'assistant_start': '<|assistant|>\n',
17 'assistant_end': '\n',
18 }
19
20
21def update_prompt(dialogues):
22 """Build the prompt for the model based on the dialogue history."""
23 markers = get_turn_markers()
24
25 system = (
26 f"{markers['bos']}You are an expert at creating dialogues.\n\n"
27 "Dialogue and emotional structure:\n"
28 )
29
30 human_prompts = [d[0] for d in dialogues]
31 chatbot_responses = [d[1] for d in dialogues]
32
33 p_emo = [h[0] for h in human_prompts]
34 p_utt = [h[1] for h in human_prompts]
35 r1_utt = [c[1] for c in chatbot_responses]
36 r2_emo = [c[2] for c in chatbot_responses]
37 r2_utt = [c[3] for c in chatbot_responses]
38 r3_utt = [c[5] for c in chatbot_responses]
39
40 context = (
41 "Human: (HAPPINESS) PROMPT.\n"
42 "Chatbot: (HAPPINESS) RESPONSE_1. (HAPPINESS) RESPONSE_2. (NEUTRAL) RESPONSE_3.\n"
43 )
44 for p_e, _, _, r2_e, _, _ in zip(p_emo, p_utt, r1_utt, r2_emo, r2_utt, r3_utt):
45 context += f"Human: ({p_e}) PROMPT.\n"
46 context += f"Chatbot: ({p_e}) RESPONSE_1. ({r2_e}) RESPONSE_2. (NEUTRAL) RESPONSE_3.\n"
47 context += "\n"
48
49 rules = (
50 "Dialogue rules:\n"
51 "The response must be open-domain curated. The response should be coherent, empathetic, engaging and proactive.\n"
52 "The chatbot RESPONSE is composed of 3 different sentences (RESPONSE_1, RESPONSE_2 and RESPONSE_3), separated by a period.\n"
53 "Between RESPONSE_1, RESPONSE_2 and RESPONSE_3 should be a max length of 20-25 words.\n"
54 "RESPONSE_3 must be open-ended to follow-up the conversation, so the Human is encouraged to answer with a full long sentence. Avoid yes/no questions.\n\n"
55 "Emotional response rules:\n"
56 f"RESPONSE_1 must contain a {p_emo[-1]} tone.\n"
57 f"RESPONSE_2 must contain a {r2_emo[-1]} tone.\n"
58 "RESPONSE_3 must contain a NEUTRAL tone.\n\n"
59 "Answer in a single turn to Human. Follow exactly the emotional structure and the emotional and dialogue rules."
60 f"{markers['user_start']}"
61 )
62
63 completion = ""
64 for idx, (p_e, p_u, r1_u, r2_e, r2_u, r3_u) in enumerate(zip(p_emo, p_utt, r1_utt, r2_emo, r2_utt, r3_utt)):
65 completion += f"({p_e}) {p_u}{markers['user_end']}{markers['assistant_start']}"
66 if idx != len(p_emo) - 1:
67 completion += f"({p_e}) {r1_u} ({r2_e}) {r2_u} (NEUTRAL) {r3_u}{markers['assistant_end']}{markers['user_start']}"
68
69 return system + context + rules + completion
70
71
72class Chatbot:
73 def __init__(self, dialogue_language="es"):
74 self.dialogue_language = dialogue_language
75 self.device = "cuda" if torch.cuda.is_available() else "cpu"
76
77 self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
78 self.model = AutoPeftModelForCausalLM.from_pretrained(
79 MODEL_ID,
80 torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
81 device_map="auto" if torch.cuda.is_available() else None,
82 trust_remote_code=True,
83 )
84 if not torch.cuda.is_available():
85 self.model = self.model.to(self.device)
86 self.model.eval()
87
88 def split_emo_chatbot(sentence):
89 """Extract emotions and utterances from a chatbot response."""
90 response_pos_ini = [i for i, c in enumerate(sentence) if c == "("]
91 response_pos_end = [i for i, c in enumerate(sentence) if c == ")"]
92 response_r1_utt = sentence[response_pos_end[0] + 2:response_pos_ini[1]].strip()
93 response_r2_utt = sentence[response_pos_end[1] + 2:response_pos_ini[2]].strip()
94 response_r3_utt = sentence[response_pos_end[2] + 2:].lstrip()
95 return response_r1_utt, response_r2_utt, response_r3_utt
96
97 def select_dialogue(self, dialogue_language):
98 """Return a list of example dialogues for the given language."""
99 if dialogue_language == "en":
100 dialogue_base = [
101 [["HAPPINESS", "Hi, who are you?"],
102 ["HAPPINESS", "Hi! I'm Ray, a social personal assistant robot with emotions.", "HAPPINESS", "I'm here to chat with you about anything you'd like.", "NEUTRAL", "What would you like to talk about?"]],
103 [["HAPPINESS", "I'm interested in talking about you, tell me more."],
104 ["HAPPINESS", "Great! I'm glad you want to get to know me!", "NEUTRAL", "I'm designed to help and talk with people about any topic.", "NEUTRAL", "I can talk about science, technology, history, or just have a pleasant conversation. What interests you?"]],
105 ]
106 dialogue = [
107 [["HAPPINESS", "Nice to meet you, Ray. I'd like to know more about you."],
108 ["HAPPINESS", "The pleasure is mine!", "HAPPINESS", "I'm a chatbot designed to chat and learn with you.", "NEUTRAL", "Would you like to talk about a specific topic?"]],
109 [["HAPPINESS", "I love talking to you, you're very interesting."],
110 ["HAPPINESS", "That's so nice to hear! I'm glad you enjoy talking to me.", "NEUTRAL", "I'm designed to have meaningful and empathetic conversations.", "NEUTRAL", "Would you like to talk about emotions, artificial intelligence, or something more personal?"]],
111 ]
112 else:
113 dialogue_base = [
114 [["HAPPINESS", "Hola, ¿quién eres?"],
115 ["HAPPINESS", "¡Hola! Soy Ray y soy un robot social asistente personal con emociones.", "HAPPINESS", "Estoy aquí para charlar contigo sobre cualquier tema.", "NEUTRAL", "¿Sobre qué te gustaría hablar?"]],
116 [["HAPPINESS", "Me interesa hablar sobre ti, cuéntame más detalles."],
117 ["HAPPINESS", "¡Genial, me encanta que quieras conocerme!", "NEUTRAL", "Estoy diseñado para ayudar y hablar con la gente sobre cualquier tema.", "NEUTRAL", "Puedo hablar de ciencia, tecnología, historia o simplemente tener una charla amena. ¿Qué te interesa?"]],
118 ]
119 dialogue = [
120 [["HAPPINESS", "Mucho gusto, Ray. Me gustaría saber más sobre ti."],
121 ["HAPPINESS", "¡El gusto es mío!", "HAPPINESS", "Soy un chatbot diseñado para conversar y aprender contigo.", "NEUTRAL", "¿Quieres hablar de algún tema en específico?"]],
122 [["HAPPINESS", "Me encanta hablar contigo, eres muy interesante."],
123 ["HAPPINESS", "¡Qué lindo escuchar eso! Me alegra que disfrutes hablar conmigo.", "NEUTRAL", "Estoy diseñado para tener conversaciones significativas y empáticas.", "NEUTRAL", "¿Te gustaría que hablemos sobre emociones, inteligencia artificial, o algo más personal?"]],
124 ]
125 return dialogue_base + dialogue
126
127 def chat_with_model(self, dialogues, max_new_tokens=96):
128 prompt_text = update_prompt(dialogues)
129 inputs = self.tokenizer(prompt_text, return_tensors="pt").to(self.model.device)
130 with torch.no_grad():
131 outputs = self.model.generate(
132 **inputs,
133 max_new_tokens=max_new_tokens,
134 do_sample=True,
135 temperature=0.7,
136 top_p=0.9,
137 eos_token_id=self.tokenizer.eos_token_id,
138 )
139 generated = outputs[0][inputs["input_ids"].shape[-1]:]
140 response = self.tokenizer.decode(generated, skip_special_tokens=True).splitlines()[0].strip()
141 print("Response:", response, "\n")
142 return response
143
144 def main(self):
145 emotions = ["ANGER", "FEAR", "SADNESS", "DISGUST", "HAPPINESS", "SURPRISE", "NEUTRAL"]
146 dialogue = self.select_dialogue(self.dialogue_language)
147
148 while True:
149 if len(dialogue) > 7:
150 dialogue.pop(2)
151
152 p_emo = random.choice(emotions)
153 user_sentence = input(f"Enter your sentence: ({p_emo}) ")
154 if user_sentence.strip().lower() == "exit":
155 break
156
157 r2_emo = random.choice(emotions)
158 dialogue.append([[p_emo, user_sentence], [p_emo, "", r2_emo, "", "NEUTRAL", ""]])
159 response = self.chat_with_model(dialogue)
160
161 try:
162 r1_utt, r2_utt, r3_utt = self.split_emo_chatbot(response)
163 except Exception:
164 if self.dialogue_language == "en":
165 r1_utt, r2_utt, r3_utt = "I'm sorry.", "I didn't understand you.", "Could you repeat?"
166 else:
167 r1_utt, r2_utt, r3_utt = "Lo siento.", "No te he entendido.", "¿Podrías repetirme?"
168
169 dialogue[-1][1] = [p_emo, r1_utt, r2_emo, r2_utt, "NEUTRAL", r3_utt]
170
171
172if __name__ == "__main__":
173 language = sys.argv[1] if len(sys.argv) > 1 else "en"
174 chatbot = Chatbot(dialogue_language=language)
175 chatbot.main()