updates

app.py

@@ -1,15 +1,14 @@
 import gradio as gr
-import random
-import time
 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
 
+from models.chatml import ChatML
+
 HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
 FIREBASE_URL = os.environ.get("FIREBASE_URL")
 CERTIFICATE = json.loads(os.environ.get("CERTIFICATE"))
@@ -21,12 +20,12 @@ models = [
         endpoint="mpt-7b",
         namespace="tenant-chairesearch-test",
         generation_params={
-            'temperature': 1.0,
+            'temperature': 0.72,
             'repetition_penalty': 1.0,
             'max_new_tokens': 128,
-            'top_k': 1,
-            'top_p': 1.0,
-            'do_sample': False,
+            'top_k': 10,
+            'top_p': 0.9,
+            'do_sample': True,
             'eos_token_id': 187,
         }
     ),
@@ -35,15 +34,57 @@ models = [
         endpoint="mpt-7b-storywriter",
         namespace="tenant-chairesearch-test",
         generation_params={
-            'temperature': 1.0,
+            'temperature': 0.72,
+            'repetition_penalty': 1.0,
+            'max_new_tokens': 128,
+            'top_k': 10,
+            'top_p': 0.9,
+            'do_sample': True,
+            'eos_token_id': 187,
+        }
+    ),
+    ChatML(
+        name="mosaicml/mpt-7b-chat",
+        endpoint="mpt-7b-chat",
+        namespace="tenant-chairesearch-test",
+        generation_params={
+            'temperature': 0.72,
+            'repetition_penalty': 1.0,
+            'max_new_tokens': 128,
+            'top_k': 10,
+            'top_p': 0.9,
+            'do_sample': True,
+            'eos_token_id': 50278,
+        }
+    ),
+    BaseModel(
+        name="togethercomputer/RedPajama-INCITE-Base-7B-v0.1",
+        endpoint="redpajama-base-7b",
+        namespace="tenant-chairesearch-test",
+        generation_params={
+            'temperature': 0.72,
+            'repetition_penalty': 1.0,
+            'max_new_tokens': 128,
+            'top_k': 10,
+            'top_p': 0.9,
+            'do_sample': True,
+            'eos_token_id': 187,
+        }
+    ),
+    BaseModel(
+        name="togethercomputer/RedPajama-INCITE-Chat-7B-v0.1",
+        endpoint="redpajama-chat-7b",
+        namespace="tenant-chairesearch-test",
+        generation_params={
+            'temperature': 0.72,
             'repetition_penalty': 1.0,
             'max_new_tokens': 128,
-            'top_k': 1,
-            'top_p': 1.0,
-            'do_sample': False,
+            'top_k': 10,
+            'top_p': 0.9,
+            'do_sample': True,
             'eos_token_id': 187,
         }
-    )
+    ),
 ]
 model_mapping = {model.name: model for model in models}
 
@@ -115,65 +156,70 @@ def get_bot_profile(bot_config):
 
 
 with gr.Blocks() as demo:
-    default_bot_id = "_bot_1ec22e2e-3e07-42c7-8508-dfa0278c1b33"
-    bot_config = download_bot_config(default_bot_id)
-    user_state = gr.State(
-        bot_config
-    )
-    with gr.Row():
-        bot_id = gr.Textbox(label="Chai bot ID", value=default_bot_id, interactive=True)
-        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")
-        regenerate = gr.Button("Regenerate")
-    values = list(model_mapping.keys())
-    model_tag = gr.Dropdown(values, value=values[0], label="Model version")
-
-
-    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]
-        bot_message = model.generate_response(conv)
-        chat_history.append(
-            (message, bot_message)
-        )
-        return "", chat_history
-
-
-    def clear_chat(chat_history, user_state):
-        chat_history = [(None, user_state["firstMessage"])]
-        return "", chat_history
-
-
-    def regenerate_response(chat_history, user_state, model_tag):
-        last_row = chat_history.pop(-1)
-        chat_history.append((last_row[0], None))
-        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
-
-
-    def reload_bot(bot_id, bot_profile, chat_history):
-        bot_config = download_bot_config(bot_id)
-        bot_profile = get_bot_profile(bot_config)
-        return bot_profile, [(None, bot_config["firstMessage"])], bot_config
-
-
-    msg.submit(respond, [msg, chatbot, user_state, model_tag], [msg, chatbot], queue=False)
-    clear.click(clear_chat, [chatbot, user_state], [msg, chatbot], queue=False)
-    regenerate.click(regenerate_response, [chatbot, user_state, model_tag], [msg, chatbot], queue=False)
-    reload_bot_button.click(reload_bot, [bot_id, bot_profile, chatbot], [bot_profile, chatbot, user_state], queue=False)
+    with gr.Tabs():
+        with gr.TabItem("Playground"):
+            default_bot_id = "_bot_e21de304-6151-4a04-b025-4c553ae8cbca"
+            bot_config = download_bot_config(default_bot_id)
+            user_state = gr.State(
+                bot_config
+            )
+            with gr.Row():
+                bot_id = gr.Textbox(label="Chai bot ID", value=default_bot_id, interactive=True)
+                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():
+                send = gr.Button("Send")
+                regenerate = gr.Button("Regenerate")
+                clear = gr.Button("Clear")
+            values = list(model_mapping.keys())
+            model_tag = gr.Dropdown(values, value=values[0], label="Model version")
+
+
+            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]
+                bot_message = model.generate_response(conv)
+                chat_history.append(
+                    (message, bot_message)
+                )
+                return "", chat_history
+
+
+            def clear_chat(chat_history, user_state):
+                chat_history = [(None, user_state["firstMessage"])]
+                return "", chat_history
+
+
+            def regenerate_response(chat_history, user_state, model_tag):
+                last_row = chat_history.pop(-1)
+                chat_history.append((last_row[0], None))
+                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
+
+
+            def reload_bot(bot_id, bot_profile, chat_history):
+                bot_config = download_bot_config(bot_id)
+                bot_profile = get_bot_profile(bot_config)
+                return bot_profile, [(None, bot_config["firstMessage"])], bot_config
+
+
+            send.click(respond, [msg, chatbot, user_state, model_tag], [msg, chatbot], queue=False)
+            msg.submit(respond, [msg, chatbot, user_state, model_tag], [msg, chatbot], queue=False)
+            clear.click(clear_chat, [chatbot, user_state], [msg, chatbot], queue=False)
+            regenerate.click(regenerate_response, [chatbot, user_state, model_tag], [msg, chatbot], queue=False)
+            reload_bot_button.click(reload_bot, [bot_id, bot_profile, chatbot], [bot_profile, chatbot, user_state],
+                                    queue=False)
 
 demo.launch(enable_queue=False)

models/chatml.py

@@ -0,0 +1,17 @@
+from conversation import Conversation
+from models.base import BaseModel
+
+
+class ChatML(BaseModel):
+
+    def _get_prompt(self, conversation: Conversation):
+        system_message = "\n".join(
+            [conversation.memory, conversation.prompt]
+        ).strip()
+        prompt = f"<|im_start|>system\n{system_message}<|im_end|>"
+
+        for message in conversation.messages:
+            prompt += f"\n<|im_start|>{message['from']}\n{message['value']}<|im_end|>"
+        prompt += f"\n<|im_start|>{conversation.bot_label}\n"
+        print(prompt)
+        return prompt