Views
No views yet
Mario-RC/aif-emotional-model.google/gemma-2-9b-itmario-rc/emotional-rlaif-dpo-gemma-2-9b-itgemmadialoguesaif_annotations preference pairs1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from peft import PeftModel
4
5base_model_id = "google/gemma-2-9b-it"
6adapter_id = "mario-rc/emotional-rlaif-dpo-gemma-2-9b-it"
7
8tokenizer = AutoTokenizer.from_pretrained(base_model_id)
9model = AutoModelForCausalLM.from_pretrained(
10 base_model_id,
11 device_map="auto",
12 torch_dtype=torch.bfloat16,
13)
14model = PeftModel.from_pretrained(model, adapter_id)
15model.eval()
16
17messages = [
18 {"role": "user", "content": "I feel overwhelmed today. Can you respond with empathy?"}
19]
20
21inputs = tokenizer.apply_chat_template(
22 messages,
23 add_generation_prompt=True,
24 return_tensors="pt",
25).to(model.device)
26
27with torch.no_grad():
28 outputs = model.generate(
29 inputs,
30 max_new_tokens=256,
31 do_sample=True,
32 temperature=0.7,
33 top_p=0.9,
34 )
35
36print(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-dpo-gemma-2-9b-it"
9
10
11def get_turn_markers():
12 return {
13 'bos': '<bos>',
14 'user_start': '<start_of_turn>user\n',
15 'user_end': '<end_of_turn>\n',
16 'assistant_start': '<start_of_turn>model\n',
17 'assistant_end': '<end_of_turn>\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)
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 )
83 if not torch.cuda.is_available():
84 self.model = self.model.to(self.device)
85 self.model.eval()
86
87 def split_emo_chatbot(sentence):
88 """Extract emotions and utterances from a chatbot response."""
89 response_pos_ini = [i for i, c in enumerate(sentence) if c == "("]
90 response_pos_end = [i for i, c in enumerate(sentence) if c == ")"]
91 response_r1_utt = sentence[response_pos_end[0] + 2:response_pos_ini[1]].strip()
92 response_r2_utt = sentence[response_pos_end[1] + 2:response_pos_ini[2]].strip()
93 response_r3_utt = sentence[response_pos_end[2] + 2:].lstrip()
94 return response_r1_utt, response_r2_utt, response_r3_utt
95
96 def select_dialogue(self, dialogue_language):
97 """Return a list of example dialogues for the given language."""
98 if dialogue_language == "en":
99 dialogue_base = [
100 [["HAPPINESS", "Hi, who are you?"],
101 ["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?"]],
102 [["HAPPINESS", "I'm interested in talking about you, tell me more."],
103 ["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?"]],
104 ]
105 dialogue = [
106 [["HAPPINESS", "Nice to meet you, Ray. I'd like to know more about you."],
107 ["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?"]],
108 [["HAPPINESS", "I love talking to you, you're very interesting."],
109 ["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?"]],
110 ]
111 else:
112 dialogue_base = [
113 [["HAPPINESS", "Hola, ¿quién eres?"],
114 ["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?"]],
115 [["HAPPINESS", "Me interesa hablar sobre ti, cuéntame más detalles."],
116 ["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?"]],
117 ]
118 dialogue = [
119 [["HAPPINESS", "Mucho gusto, Ray. Me gustaría saber más sobre ti."],
120 ["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?"]],
121 [["HAPPINESS", "Me encanta hablar contigo, eres muy interesante."],
122 ["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?"]],
123 ]
124 return dialogue_base + dialogue
125
126 def chat_with_model(self, dialogues, max_new_tokens=96):
127 prompt_text = update_prompt(dialogues)
128 inputs = self.tokenizer(prompt_text, return_tensors="pt").to(self.model.device)
129 with torch.no_grad():
130 outputs = self.model.generate(
131 **inputs,
132 max_new_tokens=max_new_tokens,
133 do_sample=True,
134 temperature=0.7,
135 top_p=0.9,
136 eos_token_id=self.tokenizer.eos_token_id,
137 )
138 generated = outputs[0][inputs["input_ids"].shape[-1]:]
139 response = self.tokenizer.decode(generated, skip_special_tokens=True).splitlines()[0].strip()
140 print("Response:", response, "\n")
141 return response
142
143 def main(self):
144 emotions = ["ANGER", "FEAR", "SADNESS", "DISGUST", "HAPPINESS", "SURPRISE", "NEUTRAL"]
145 dialogue = self.select_dialogue(self.dialogue_language)
146
147 while True:
148 if len(dialogue) > 7:
149 dialogue.pop(2)
150
151 p_emo = random.choice(emotions)
152 user_sentence = input(f"Enter your sentence: ({p_emo}) ")
153 if user_sentence.strip().lower() == "exit":
154 break
155
156 r2_emo = random.choice(emotions)
157 dialogue.append([[p_emo, user_sentence], [p_emo, "", r2_emo, "", "NEUTRAL", ""]])
158 response = self.chat_with_model(dialogue)
159
160 try:
161 r1_utt, r2_utt, r3_utt = self.split_emo_chatbot(response)
162 except Exception:
163 if self.dialogue_language == "en":
164 r1_utt, r2_utt, r3_utt = "I'm sorry.", "I didn't understand you.", "Could you repeat?"
165 else:
166 r1_utt, r2_utt, r3_utt = "Lo siento.", "No te he entendido.", "¿Podrías repetirme?"
167
168 dialogue[-1][1] = [p_emo, r1_utt, r2_emo, r2_utt, "NEUTRAL", r3_utt]
169
170
171if __name__ == "__main__":
172 language = sys.argv[1] if len(sys.argv) > 1 else "en"
173 chatbot = Chatbot(dialogue_language=language)
174 chatbot.main()