🌟 gemma-2-9b-it-emotional-rlaif-dpo

This repository contains a LoRA adapter trained to generate emotionally intelligent, empathetic, and engaging responses based on the gemma-2-9b-it base model.

The model was optimized using Direct Preference Optimization (DPO) with preference data generated from AI-based feedback labelers (GPT-4o, Claude 3.5, Gemini 1.5 Pro, and LLaMA 3.1 405B) following emotional alignment strategies grounded in psychological theory (e.g., Lazarus & Ekman).

---

🧠 Model Description

• Base model: gemma-2-9b-it
• Tuning strategy: LoRA fine-tuning with PEFT
• Optimization: Direct Preference Optimization (DPO)
• Task: Emotionally aligned dialogue generation
• Dimensions: Empathy, internal emotional control, and open-ended engagement
• Training data: 3,000 dialogues x 4 turns from the aif-emotional-generation dataset
• Language: English
• Model size: 9B parameters (base) + LoRA adapter
• Transformers version: 4.45.2
• PyTorch version: 2.6.0

---

🚀 How to Use

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import random
import sys

def update_prompt(dialogues):
    """
    Build the prompt for the model based on the dialogue history.
    """
    # System prompt describing the chatbot's role
    system = (
        "<bos>You are an expert at creating dialogues.\n\n"
        "Dialogue and emotional structure:\n"
    )

    # Extract dialogue elements from the dialogue history
    human_prompts = [d[0] for d in dialogues]
    chatbot_responses = [d[1] for d in dialogues]

    # Extract emotions and utterances for each turn
    p_emo = [h[0] for h in human_prompts]
    p_utt = [h[1] for h in human_prompts]
    r1_utt = [c[1] for c in chatbot_responses]
    r2_emo = [c[2] for c in chatbot_responses]
    r2_utt = [c[3] for c in chatbot_responses]
    r3_utt = [c[5] for c in chatbot_responses]

    # Example context for the model to follow
    context = (
        "Human: (HAPPINESS) PROMPT.\n"
        "Chatbot: (HAPPINESS) RESPONSE_1. (HAPPINESS) RESPONSE_2. (NEUTRAL) RESPONSE_3.\n"
    )
    # Add previous dialogue turns to the context
    for p_e, p_u, r1_u, r2_e, r2_u, r3_u in zip(p_emo, p_utt, r1_utt, r2_emo, r2_utt, r3_utt):
        context += f"Human: ({p_e}) PROMPT.\n"
        context += f"Chatbot: ({p_e}) RESPONSE_1. ({r2_e}) RESPONSE_2. (NEUTRAL) RESPONSE_3.\n"
    context += "\n"

    # Dialogue and emotional rules for the model to follow
    rules = (
        "Dialogue rules:\n"
        "The response must be open-domain curated. The response should be coherent, empathetic, engaging and proactive.\n"
        "The chatbot RESPONSE is composed of 3 different sentences (RESPONSE_1, RESPONSE_2 and RESPONSE_3), separated by a period.\n"
        "Between RESPONSE_1, RESPONSE_2 and RESPONSE_3 should be a max length of 20-25 words.\n"
        "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"
        "Emotional response rules:\n"
        f"RESPONSE_1 must contain a {p_emo[-1]} tone.\n"
        f"RESPONSE_2 must contain a {r2_emo[-1]} tone.\n"
        "RESPONSE_3 must contain a NEUTRAL tone.\n\n"
        "Answer in a single turn to Human. Follow exactly the emotional structure and the emotional and dialogue rules.<start_of_turn>user\n"
    )

    # Build the dialogue completion section for the prompt
    completion = ""
    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)):
        completion += f"({p_e}) {p_u}<end_of_turn>\n<start_of_turn>model\n"
        if idx != len(p_emo) - 1:
            completion += f"({p_e}) {r1_u} ({r2_e}) {r2_u} (NEUTRAL) {r3_u}<end_of_turn>\n<start_of_turn>user\n"

    return system + context + rules + completion

class Chatbot:
    def __init__(self, dialogue_language="es"):
        # Set language and device (GPU if available)
        self.dialogue_language = dialogue_language
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        model_name = "mario-rc/gemma-2-9b-it-emotional-rlaif-dpo"
        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).to(self.device)

    def get_token_len(self):
        # Estimate the length of special tokens in the prompt
        return len('<bos>') + len('<start_of_turn>') * 8 + len('<end_of_turn>' * 15) + len('<start_of_turn>') * 8

    @staticmethod
    def split_emo_user(sentence):
        """
        Extract emotion and utterance from a user sentence.
        """
        emo_user_ini = sentence.find('(')
        emo_user_end = sentence.find(')')
        user_utt = sentence[emo_user_end+2:].strip()
        user_emo = sentence[emo_user_ini+1:emo_user_end].strip()
        return user_utt, user_emo

    @staticmethod
    def split_emo_chatbot(sentence):
        """
        Extract emotions and utterances from a chatbot response.
        """
        response_pos_ini = [i for i, c in enumerate(sentence) if c == '(']
        response_pos_end = [i for i, c in enumerate(sentence) if c == ')']
        response_r1_utt = sentence[response_pos_end[0]+2:response_pos_ini[1]].strip()
        response_r1_emo = sentence[response_pos_ini[0]+1:response_pos_end[0]].strip()
        response_r2_utt = sentence[response_pos_end[1]+2:response_pos_ini[2]].strip()
        response_r2_emo = sentence[response_pos_ini[1]+1:response_pos_end[1]].strip()
        response_r3_utt = sentence[response_pos_end[2]+2:].lstrip()
        response_r3_emo = sentence[response_pos_ini[2]+1:response_pos_end[2]].strip()
        return response_r1_emo, response_r1_utt, response_r2_emo, response_r2_utt, response_r3_emo, response_r3_utt

    def select_dialogue(self, dialogue_language):
        """
        Return a list of example dialogues for the given language.
        """
        if dialogue_language == "en":  # English
            dialogue_base = [
                [['HAPPINESS', 'Hi, who are you?'],
                 ['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?"]],
                [['HAPPINESS', "I'm interested in talking about you, tell me more."],
                 ['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?"]]
            ]
            dialogue = [
                [['HAPPINESS', "Nice to meet you, Ray. I'd like to know more about you."],
                 ['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?"]],
                [['HAPPINESS', "I'd like to know more about you."],
                 ['HAPPINESS', "Of course! I'm here to help you.", 'NEUTRAL', "I'm a virtual assistant, and I can help you with many things.", 'NEUTRAL', "What would you like to know?"]],
                [['HAPPINESS', "I'm interested in learning more about you. Tell me more."],
                 ['HAPPINESS', "Great! I'm happy you want to get to know me.", 'NEUTRAL', "I'm a personal assistant robot, designed to help and talk with people about any topic.", 'NEUTRAL', "Is there a specific topic you'd like to talk about?"]],
                [['HAPPINESS', "I love talking to you, you're very interesting."],
                 ['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?"]],
                [['HAPPINESS', "Are you always in a good mood?"],
                 ['HAPPINESS', "Good question!", 'NEUTRAL', "I don't have emotions like a human, but I can simulate them to make the conversation more pleasant.", 'NEUTRAL', "I try to keep a positive attitude so you feel comfortable talking to me."]]
            ]
        else:  # Spanish
            dialogue_base = [
                [['HAPPINESS', 'Hola, ¿quién eres?'],
                 ['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?']],
                [['HAPPINESS', 'Me interesa hablar sobre ti, cuéntame más detalles.'],
                 ['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?']]
            ]
            dialogue = [
                [['HAPPINESS', 'Mucho gusto, Ray. Me gustaría saber más sobre ti.'],
                 ['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?']],
                [['HAPPINESS', 'Me gustaría saber más sobre ti.'],
                 ['HAPPINESS', '¡Claro! Estoy aquí para ayudarte.', 'NEUTRAL', 'Soy un asistente virtual, y puedo ayudarte con muchas cosas.', 'NEUTRAL', '¿Qué te gustaría saber?']],
                [['HAPPINESS', 'Me interesa saber más de ti. Cuéntame más detalles'],
                 ['HAPPINESS', '¡Genial! Me encanta que quieras conocerme', 'NEUTRAL', 'Soy un robot asistente personal, diseñado para ayudar y hablar con la gente sobre cualquier tema', 'NEUTRAL', '¿Hay algún tema específico del que te gustaría hablar?']],
                [['HAPPINESS', 'Me encanta hablar contigo, eres muy interesante.'],
                 ['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?']],
                [['HAPPINESS', '¿Siempre estás de buen humor?'],
                 ['HAPPINESS', '¡Buena pregunta!', 'NEUTRAL', 'No tengo emociones como un humano, pero puedo simularlas para hacer la conversación más agradable.', 'NEUTRAL', 'Trato de mantener una actitud positiva para que te sientas cómodo al hablar conmigo.']]
            ]
        return dialogue_base + dialogue

    def chat_with_model(self, dialogues, token_len, max_length):
        """
        Generate a chatbot response given the dialogue history.
        """
        # Build the prompt for the model
        prompt_text = update_prompt(dialogues)
        # Tokenize the prompt and move to device
        inputs = self.tokenizer(prompt_text, return_tensors="pt").to(self.device)
        # Generate a response from the model
        outputs = self.model.generate(**inputs, max_length=max_length)
        # Decode the output tokens to text
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        # Extract only the new response (after the prompt)
        response = response[len(prompt_text)-token_len:].splitlines()[0].strip()
        print('Response:', response, '\n')
        return response

    def main(self):
        # List of possible emotions
        emotions = ['ANGER', 'FEAR', 'SADNESS', 'DISGUST', 'HAPPINESS', 'SURPRISE', 'NEUTRAL']
        # Get initial dialogue examples
        dialogue = self.select_dialogue(self.dialogue_language)
        token_len = self.get_token_len()
        max_length = 2048

        while True:
            # Limit dialogue history length
            if len(dialogue) > 7:
                dialogue.pop(2)
            # Randomly select emotion for user prompt
            p_emo = random.choice(emotions)
            user_sentence = input(f"Enter your sentence: ({p_emo}) ")
            if user_sentence.strip().lower() == "exit":
                break
            # Randomly select emotion for chatbot's second response
            r2_emo = random.choice(emotions)
            # Add new turn to dialogue history
            dialogue.append([[p_emo, user_sentence], [p_emo, "", r2_emo, "", "NEUTRAL", ""]])
            # Generate chatbot response
            response = self.chat_with_model(dialogue, token_len, max_length)
            try:
                # Try to extract emotions and utterances from the response
                _, r1_utt, _, r2_utt, _, r3_utt = self.split_emo_chatbot(response)
            except Exception:
                # Fallback in case of parsing error
                if self.dialogue_language == "en":
                    _, r1_utt, _, r2_utt, _, r3_utt = [p_emo, "I'm sorry.", r2_emo, "I didn't understand you.", "NEUTRAL", "Could you repeat?"]
                else:
                    _, r1_utt, _, r2_utt, _, r3_utt = [p_emo, "Lo siento.", r2_emo, "No te he entendido.", "NEUTRAL", "¿Podrías repetirme?"]
            # Update the last chatbot response in the dialogue history
            dialogue[-1][1] = [p_emo, r1_utt, r2_emo, r2_utt, "NEUTRAL", r3_utt]

if __name__ == "__main__":
    # Select language from command line argument or default to English
    language = sys.argv[1] if len(sys.argv) > 1 else 'en'
    chatbot = Chatbot(dialogue_language=language)
    chatbot.main()