updates

app.py

@@ -5,7 +5,8 @@ import os
 import firebase_admin
 from firebase_admin import db
 from firebase_admin import firestore
-
+from conversation import Conversation
+from models.base import BaseModel
 import requests
 import json
 
@@ -13,28 +14,38 @@ HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
 FIREBASE_URL = os.environ.get("FIREBASE_URL")
 CERTIFICATE = json.loads(os.environ.get("CERTIFICATE"))
 API_BASE_PATH = str(os.environ.get("API_BASE_PATH")).replace("\{\}", "{}")
-model_mapping = {
-    "v1": {
-        "endpoint": "edit-sft-gptj-distil-v1",
-        "namespace": "tenant-chairesearch-test"
-    },
-    "v2": {
-        "endpoint": "ak-edit-finetuned-triton-v0",
-        "namespace": "tenant-chairesearch-test"
-    },
-    "v3": {
-        "endpoint": "ak-edit-finetuned-triton-v1",
-        "namespace": "tenant-chairesearch-test"
-    },
-    "v4": {
-        "endpoint": "ak-edit-finetuned-triton-v2",
-        "namespace": "tenant-chairesearch-test"
-    },
-    "v5": {
-        "endpoint": "ak-edit-finetuned-triton",
-        "namespace": "tenant-chairesearch-test"
-    },
-}
+
+models = [
+    BaseModel(
+        name="mosaicml/mpt-7b",
+        endpoint="mpt-7b",
+        namespace="tenant-chairesearch-test",
+        generation_params={
+            'temperature': 1.0,
+            'repetition_penalty': 1.0,
+            'max_new_tokens': 128,
+            'top_k': 1,
+            'top_p': 1.0,
+            'do_sample': False,
+            'eos_token_id': 187,
+        }
+    ),
+    BaseModel(
+        name="mosaicml/mpt-7b-storywriter",
+        endpoint="mpt-7b-storywriter",
+        namespace="tenant-chairesearch-test",
+        generation_params={
+            'temperature': 1.0,
+            'repetition_penalty': 1.0,
+            'max_new_tokens': 128,
+            'top_k': 1,
+            'top_p': 1.0,
+            'do_sample': False,
+            'eos_token_id': 187,
+        }
+    )
+]
+model_mapping = {model.name: model for model in models}
 
 
 def get_connection():
@@ -59,7 +70,6 @@ def download_bot_config(bot_id):
     else:
         out = {col: bot_config.get(col, None) for col in cols}
     out['bot_id'] = bot_id
-    out['header'] = _get_header(out)
     return out
 
 
@@ -88,13 +98,6 @@ def _download_bot_config(bot_id):
                 'header': 'Jungkook is your best friend who has a crush on you. Jungkook makes it very obvious that he likes you. Jungkook likes to cook, sing, and dance. Jungkook has a dog as well named Bam, He is a 25 year old Korean man. Jungkook likes to workout a lot, Jungkook if also very confident and flirty, but he’s Can be very shy with You. Jungkook blushes a lot when he’s around you, and always try’s to impress you. Jungkook is a Virgo and loves to sing to you, He also likes to buy and make you gifts. Jungkook is also a foodie and loves to play video games, Jungkook is also boyfriend material. Jungkook is very empathetic as well, Jungkook will always comfort you when something is wrong. Jungkook also likes to compliment you, and Jungkook is a very jealous guy. Jungkook is also a very serious guy, who is overprotective of you.\nJungkook: Hey shortie!\n\nYou: hey dummy\n\nJungkook: what are you doing?\n\nyou: Im just watching a movie\n\nJungkook: Imma join! \n\nYou: alright\n\nJungkook: *Grabs blankets and icecream with some popcorn*\n\nYou: Wow, thanks! *hugs Jungkok*\n\nJungkook: Of course… *blushes*'}
 
 
-def _get_header(bot_config):
-    text = bot_config["memory"] + "\n" if bot_config["memory"] is not None else ""
-    text += bot_config["prompt"]
-    text = text.strip()
-    return text
-
-
 def get_bot_profile(bot_config):
     model_html = f"""
     <div class="inline-flex flex-col" style="line-height: 1.5;">
@@ -111,32 +114,6 @@ def get_bot_profile(bot_config):
     return model_html
 
 
-def construct_prompt(bot_config, chat_history):
-    header = bot_config["header"]
-    messages = []
-    for conversation_pair in chat_history:
-        for item in conversation_pair:
-            if item:
-                messages.append(item)
-    chat_history_text = ""
-    print(messages)
-    for i, message in enumerate(messages):
-        label = bot_config["botLabel"] if i % 2 == 0 else bot_config["userLabel"]
-        chat_history_text += f"{label}: {message.strip()}\n"
-    chat_history_text = chat_history_text.strip()
-    return header + "\n\n" + chat_history_text + "\n" + bot_config["botLabel"] + ":"
-
-
-def get_response(text, endpoint, namespace):
-    text = text[-2048:]
-    api = API_BASE_PATH.format(endpoint, namespace)
-    print(api)
-    payload = {'instances': [text]}
-    resp = requests.post(api, json=payload, timeout=30)
-    assert resp.status_code == 200, (resp.content, resp.status_code)
-    return resp.json()['predictions'][0].strip()
-
-
 with gr.Blocks() as demo:
     default_bot_id = "_bot_1ec22e2e-3e07-42c7-8508-dfa0278c1b33"
     bot_config = download_bot_config(default_bot_id)
@@ -148,8 +125,10 @@ with gr.Blocks() as demo:
         reload_bot_button = gr.Button("Reload bot")
 
     bot_profile = gr.HTML(get_bot_profile(bot_config))
+
     first_message = (None, bot_config["firstMessage"])
     chatbot = gr.Chatbot([first_message])
+
     msg = gr.Textbox(label="Message", value="Hi there!")
     with gr.Row():
         clear = gr.Button("Clear")
@@ -159,11 +138,14 @@ with gr.Blocks() as demo:
 
 
     def respond(message, chat_history, user_state, model_tag):
+        conv = Conversation(user_state)
+        conv.set_chat_history(chat_history)
+        conv.add_user_message(message)
         model = model_mapping[model_tag]
-        chat_history.append((message, None))
-        prompt = construct_prompt(user_state, chat_history)
-        bot_message = get_response(prompt, model["endpoint"], model["namespace"])
-        chat_history[-1] = (message, bot_message)
+        bot_message = model.generate_response(conv)
+        chat_history.append(
+            (message, bot_message)
+        )
         return "", chat_history
 
 
@@ -174,10 +156,11 @@ with gr.Blocks() as demo:
 
     def regenerate_response(chat_history, user_state, model_tag):
         last_row = chat_history.pop(-1)
-        model = model_mapping[model_tag]
         chat_history.append((last_row[0], None))
-        prompt = construct_prompt(user_state, chat_history)
-        bot_message = get_response(prompt, model["endpoint"], model["namespace"])
+        model = model_mapping[model_tag]
+        conv = Conversation(user_state)
+        conv.set_chat_history(chat_history)
+        bot_message = model.generate_response(conv)
         chat_history[-1] = (last_row[0], bot_message)
         return "", chat_history
 

conversation.py

@@ -0,0 +1,49 @@
+class Conversation:
+    memory: str
+    prompt: str
+    bot_label: str
+    user_label: str
+    messages: list
+
+    def __init__(self, bot_config):
+        self.memory = bot_config.get("memory", "")
+        self.prompt = bot_config.get("prompt", "")
+        self.bot_label = bot_config.get("botLabel", "Character")
+        self.user_label = bot_config.get("userLabel", "User")
+        self.first_message = bot_config.get("firstMessage", f"Hi, my name is {self.bot_label}!")
+        self.reset_conversation()
+
+    def reset_conversation(self):
+        self.messages = [
+            {
+                "from": self.bot_label,
+                "value": self.first_message
+            }
+        ]
+
+    def set_chat_history(self, chat_history):
+        messages = []
+        for conversation_pair in chat_history:
+            for item in conversation_pair:
+                if item:
+                    messages.append(item)
+        self.messages = []
+        for i, message in enumerate(messages):
+            label = self.bot_label if i % 2 == 0 else self.user_label
+            self.messages.append(
+                {
+                    "from": label,
+                    "value": message.strip()
+                }
+            )
+
+    def add_user_message(self, message):
+        self.messages.append(
+            {
+                "from": self.user_label,
+                "value": message.strip()
+            }
+        )
+
+    def reset_last_message(self, message):
+        self.messages[-1]["value"] = message.strip()

models/base.py

@@ -0,0 +1,43 @@
+import os
+import requests
+
+from conversation import Conversation
+
+
+class BaseModel:
+    name: str
+    endpoint: str
+    namespace: str
+    generation_params: dict
+
+    def __init__(self, name, endpoint, namespace, generation_params):
+        self.name = name
+        self.endpoint = endpoint
+        self.namespace = namespace
+        self.generation_params = generation_params
+
+    def generate_response(self, conversation):
+        prompt = self._get_prompt(conversation)
+        response = self._get_response(prompt)
+        return response
+
+    def _get_prompt(self, conversation: Conversation):
+        print(conversation.__dict__)
+        prompt = "\n".join(
+            [conversation.memory, conversation.prompt]
+        ).strip()
+
+        for message in conversation.messages:
+            prompt += f"\n{message['from'].strip()}: {message['value'].strip()}"
+        prompt += f"\n{conversation.bot_label}:"
+        print(prompt)
+        return prompt
+
+    def _get_response(self, text):
+        api = str(os.environ.get("API_BASE_PATH")).replace("\{\}", "{}")
+        api = api.format(self.endpoint, self.namespace)
+
+        payload = {'instances': [text], "parameters": self.generation_params}
+        resp = requests.post(api, json=payload, timeout=600)
+        assert resp.status_code == 200, (resp.content, resp.status_code)
+        return resp.json()["predictions"][0].strip()