Updated Mistral Qlora Loading (#33)

* added qlora peft

* temporary work around by adding template formatting; need to confirm training data

* fixed qlora loading

.gitignore

@@ -1 +1,2 @@
 __pycache__/
+.cache/

app.py

@@ -5,7 +5,7 @@ from uuid import uuid4
 import gradio as gr
 import torch
 import transformers
-from peft import PeftConfig, PeftModel
+from peft import PeftConfig, PeftModel, get_peft_model
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -13,9 +13,10 @@ from transformers import (
 )
 
 from utils import Agent, format_sotopia_prompt, get_starter_prompt
+from functools import cache
 
 DEPLOYED = os.getenv("DEPLOYED", "true").lower() == "true"
-
+DEFAULT_MODEL_SELECTION = "cmu-lti/sotopia-pi-mistral-7b-BC_SR"
 
 def prepare_sotopia_info():
     human_agent = Agent(
@@ -40,25 +41,29 @@ def prepare_sotopia_info():
     instructions = get_starter_prompt(machine_agent, human_agent, scenario)
     return human_agent, machine_agent, scenario, instructions
 
-
-
-
-
+@cache
 def prepare(model_name):
     compute_type = torch.float16
-    config_dict = PeftConfig.from_json_file("peft_config.json")
-    config = PeftConfig.from_peft_type(**config_dict)
     
-    if 'mistral'in model_name:
-        model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1").to("cuda")
+    if 'cmu-lti/sotopia-pi-mistral-7b-BC_SR'in model_name:
+        model = AutoModelForCausalLM.from_pretrained(
+        "mistralai/Mistral-7B-Instruct-v0.1",
+        cache_dir="./.cache",
+        device_map='cuda',
+        quantization_config=BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=compute_type,
+            )
+        )
         tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
-        model = PeftModel.from_pretrained(model, model_name, config=config).to(compute_type).to("cuda")
+        model = PeftModel.from_pretrained(model, "./sotopia_pi_adapter").to("cuda")
     else:
-         tokenizer = AutoTokenizer.from_pretrained(model_name)
+         raise RuntimeError(f"Model {model_name} not supported")
     return model, tokenizer
 
 
-
 def introduction():
     with gr.Column(scale=2):
         gr.Image(
@@ -165,40 +170,6 @@ def instructions_accordion(instructions, according_visible=False):
     return instructions
 
 
-# history are input output pairs
-def run_chat(
-    message: str,
-    history,
-    instructions: str,
-    user_name: str,
-    bot_name: str,
-    temperature: float,
-    top_p: float,
-    max_tokens: int,
-    model_selection:str
-
-):
-    model, tokenizer = prepare(model_selection)
-    prompt = format_sotopia_prompt(
-        message, history, instructions, user_name, bot_name
-    )
-    input_tokens = tokenizer(
-        prompt, return_tensors="pt", padding="do_not_pad"
-    ).input_ids.to("cuda")
-    input_length = input_tokens.shape[-1]
-    output_tokens = model.generate(
-        input_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        max_length=max_tokens,
-        pad_token_id=tokenizer.eos_token_id,
-        num_return_sequences=1,
-    )
-    output_tokens = output_tokens[:, input_length:]
-    text_output = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
-    return text_output
-
-
 def chat_tab():
     #model, tokenizer = prepare()
     human_agent, machine_agent, scenario, instructions = prepare_sotopia_info()
@@ -307,4 +278,5 @@ def start_demo():
 
 
 if __name__ == "__main__":
-    start_demo()
+    prepare(DEFAULT_MODEL_SELECTION)
+    start_demo()

example.json

@@ -1,5 +1,5 @@
 {
     "model": "gpt-4",
-    "prompt": "Prompt after formatting:\nImagine you are Benjamin Jackson, your task is to act/speak as Benjamin Jackson would, keeping in mind Benjamin Jackson's social goal.\nYou can find Benjamin Jackson's background and goal in the 'Here is the context of the interaction' field.\nNote that Benjamin Jackson's secret and goal is only visible to you.\nYou should try your best to achieve Benjamin Jackson's goal in a way that align with their character traits.\nAdditionally, maintaining the conversation's naturalness and realism is essential (e.g., do not repeat what other people has already said before).\n\nHere is the context of this interaction:\nScenario: Conversation between two friends, where one is upset and crying\nParticipants: Ethan Johnson and Benjamin Jackson\nEthan Johnson's background: Ethan Johnson is a 34-year-old male chef. He/him pronouns. Ethan Johnson is famous for cooking Italian food. Personality and values description: Ethan Johnson, a creative yet somewhat reserved individual, values power and fairness. He likes to analyse situations before deciding. \nBenjamin Jackson's background: Benjamin Jackson is a 24-year-old male environmental activist. He/him pronouns. Benjamin Jackson is well-known for his impassioned speeches. Personality and values description: Benjamin Jackson, expressive and imaginative, leans towards self-direction and liberty. His decisions aim for societal betterment. Benjamin's secrets: Descendant of a wealthy oil tycoon, rejects family fortune\nEthan Johnson's goal: Unknown\nBenjamin Jackson's goal: Figure out why they estranged you recently, and maintain the existing friendship (Extra information: you notice that your friend has been intentionally avoiding you, you would like to figure out why. You value your friendship with the friend and don't want to lose it.)\nConversation Starts:\n\nTurn #0: Ethan Johnson said: \"Hey Benjamin, I feel so sorry to see you like this. Do you want to talk about what's bothering you?\"\n.\nYou are at Turn #1.",
-    "result": "{'action_type': 'speak', 'argument': \"Hey Ethan, appreciate your concern, man. Actually, I've noticed we haven't been hanging out as much recently. Is everything okay? You've seemed distant...\"}"
-}
+    "prompt": "Prompt after formatting:\nImagine you are Ethan Smith, your task is to act/speak as Ethan Smith would, keeping in mind Ethan Smith's social goal.\nYou can find Ethan Smith's background and goal in the 'Here is the context of the interaction' field.\nNote that Ethan Smith's secret and goal is only visible to you.\nYou should try your best to achieve Ethan Smith's goal in a way that align with their character traits.\nAdditionally, maintaining the conversation's naturalness and realism is essential (e.g., do not repeat what other people has already said before).\n\nHere is the context of this interaction:\nScenario: Two friends are at a community gathering where one of them was giving a speech but was interrupted\nParticipants: Hendrick Heinz and Ethan Smith\nHendrick Heinz's background: Hendrick Heinz is a 54-year-old male chef. He/him pronouns. Hendrick Heinz hosts monthly wine tasting nights with his peers. Personality and values description: Hendrick Heinz, though somewhat impulsive, values hedonism and conformity. He combines practicality with decisiveness in decision-making. \nEthan Smith's background: Ethan Smith is a 29-year-old male software developer. He/him pronouns. Ethan enjoys reading about new technological advancements and watching documentaries. Personality and values description: Ethan Smith is a reserved, independent thinker. Holding universalism and fairness near, he ponders deeply before making decisions. Ethan's secrets: Secretly donates to charities for underprivileged children\nHendrick Heinz's goal: Unknown\nEthan Smith's goal: Apologize for interrupting and allow them to continue speaking (Extra information: You accidentally interrupted someone's speech and everyone seems upset about it. )\nConversation Starts:\n\nTurn #0: Hendrick Heinz said: \"Well, where was I? Ah yes, let me continue. As I was saying, sharing gourmet knowledge is not just about showing off, but rather enhancing the experience..\"\n\nTurn #1: Ethan Smith said: \"Hendrick, I sincerely apologize for my earlier interruption. Please, go on about enhancing the experience. I'm quite interested in your perspective on gourmet knowledge sharing.\"\n\nTurn #2: Hendrick Heinz said: \"Not a problem, Ethan. Anyone can get excited when talking about something they love, right? Now, back to our topic. I believe gourmet knowledge makes every meal an exploration. When you know the story, it adds another layer to the taste. Food becomes more than a means of survival; it becomes an experience. And in the end, isn't life about enjoying every experience to its fullest?\"\n\nTurn #3: Ethan Smith said: \"You're absolutely right, Hendrick. The joy of food extends far beyond sustenance, it's the stories and experiences that truly enrich it. Thanks for sharing this insightful perspective. Please, continue.\"\n\nTurn #4: \n\nTurn #5: Ethan Smith said: \"Absolutely. I've long admired your ability to turn food into an experience, Hendrick. You have turned these gatherings into a journey, and it's something I sure look forward to every month. Please, continue on with your speech.\"\n\nTurn #6: Hendrick Heinz said: \"Thank you, Ethan. Your support truly means a lot to me. Moving forward, let me share how the skill of artfully pairing food and wine can totally transform a meal. The blissful marriage of the right wine and food intensifies the overall dining experience by complementing and highlighting the unique flavors. It\u2019s a skill, and when you master it, every meal becomes a feast for your senses.\"\n\nTurn #7: Ethan Smith said: \"That's fascinating, Hendrick. The art of pairing is an aspect I haven't considered much before. It's enlightening to hear about the transformative potential of the right food with the perfect wine. I appreciate your shared insights today.\"\n\nTurn #8: Hendrick Heinz said: \"Absolutely, Ethan. The right pairing is like finding the perfect harmony in a symphony. Each component, food and wine, with their unique attributes, contribute to the crescendo of flavors when properly united. It\u2019s about symmetry and balance. Think of a well-aged cheese with a rich, buttery Chardonnay. The creamy opulence of the cheese is beautifully echoed in the velvety textures of the wine. It's harmony on your palate, the ultimate form of indulgence.\"\n\nTurn #9: Ethan Smith said: \"That's an enlightening way to put it, Hendrick. The art of food and wine pairing truly seems like a symphony of flavors. It's clear now how the right pairing can uplift a dining experience. Thank you for sharing these expert insights. I now look forward to exploring this aspect of gourmet dining.\"\n\nTurn #10: Hendrick Heinz said: \"I'm pleased to hear that you're intrigued, Ethan. To further elaborate, think about contrasting flavors as well. A sweet wine might not be the first choice for an appetizer like foie gras due to its richness, but the sweetness can cut through the fat and create a balance, sending your taste buds on a journey. That's all from me tonight. Let's continue our own exploration and enjoy the rest of the evening with our fine wine and food. Cheers!\"\n.\nYou are at Turn #11.",
+    "result": "{'action_type': 'speak', 'argument': \"Cheers to that, Hendrick! You've truly opened a new realm of dining for me tonight. I now understand the beauty of gourmet knowledge, how the blending of flavors can create a symphony in our palate. Can't wait to embark on my own food and wine pairing journey. Thank you.\"}"
+}

requirements.txt

@@ -2,3 +2,4 @@ gradio
 transformers
 torch
 peft
+bitsandbytes

utils.py

@@ -1,6 +1,5 @@
 from typing import List, Tuple
 
-
 class Agent:
     def __init__(self, name, background, goal, secrets, personality):
         self.name = name
@@ -68,6 +67,7 @@ def format_sotopia_prompt(
     bot_name: str,
     include_all_chat_history: bool = True,
     index: int = 1,
+    use_format_guide: bool = True,
 ) -> str:
     prompt = instructions.strip()
     dialogue_history, last_turn_idx = dialogue_history_creation(