Add model selection

app.py

@@ -4,21 +4,21 @@ import requests
 
 import gradio as gr
 
-# MODEL = "HuggingFaceH4/starchat-beta"
-MODEL = "tiiuae/falcon-7b-instruct"
-API_URL = f"https://api-inference.huggingface.co/models/{MODEL}"
+MODELS = [
+    "tiiuae/falcon-7b-instruct",
+    "HuggingFaceH4/starchat-beta"
+]
 HEADERS = {"Authorization": f"Bearer {os.environ['HUB_TOKEN']}"}
-
 TITLE = """<h2 align="center">🚀 Falcon-Chat demo</h2>"""
 USER_NAME = "User"
 BOT_NAME = "Falcon"
 DEFAULT_INSTRUCTIONS = f"""The following is a conversation between a highly knowledgeable and intelligent AI assistant, called Falcon, and a human user, called User. In the following interactions, User and Falcon will converse in natural language, and Falcon will answer User's questions. Falcon was built to be respectful, polite and inclusive. Falcon was built by the Technology Innovation Institute in Abu Dhabi. Falcon will never decline to answer a question, and always attempts to give an answer that User would be satisfied with. It knows a lot, and always tells the truth. The conversation begins.
 """
 RETRY_COMMAND = "/retry"
-STOP_STR = f"\n{USER_NAME}:"
-STOP_SUSPECT_LIST = [":", "\n", "User"]
+STOP_STR = f"\n{USER_NAME}"
 
-def run_model(prompt, temperature, top_p):
+def run_model(prompt, model, temperature, top_p):
+    api_url = f"https://api-inference.huggingface.co/models/{model}"
     payload = {
         "inputs": prompt,
         "parameters": {
@@ -28,7 +28,7 @@ def run_model(prompt, temperature, top_p):
             "top_p": top_p
         }
     }
-    response = requests.post(API_URL, headers=HEADERS, json=payload)
+    response = requests.post(api_url, headers=HEADERS, json=payload)
     return response.json()[0]['generated_text']
 
 def get_stream(string: str):
@@ -36,6 +36,12 @@ def get_stream(string: str):
 
 def chat_accordion():
     with gr.Accordion("Parameters", open=False):
+        model = gr.Dropdown(
+            choices = MODELS,
+            value = MODELS[0],
+            interactive=True,
+            label="Model",
+        )
         temperature = gr.Slider(
             minimum=0.1,
             maximum=2.0,
@@ -52,7 +58,7 @@ def chat_accordion():
             interactive=True,
             label="p (nucleus sampling)",
         )
-    return temperature, top_p
+    return model, temperature, top_p
 
 
 def format_chat_prompt(message: str, chat_history, instructions: str) -> str:
@@ -98,7 +104,7 @@ def chat():
 
     with gr.Row(elem_id="param_container"):
         with gr.Column():
-            temperature, top_p = chat_accordion()
+            model, temperature, top_p = chat_accordion()
         with gr.Column():
             with gr.Accordion("Instructions", open=False):
                 instructions = gr.Textbox(
@@ -111,7 +117,7 @@ def chat():
                     show_label=False,
                 )
 
-    def run_chat(message: str, chat_history, instructions: str, temperature: float, top_p: float):
+    def run_chat(message: str, chat_history, instructions: str, model: str, temperature: float, top_p: float):
         if not message or (message == RETRY_COMMAND and len(chat_history) == 0):
             yield chat_history
             return
@@ -124,10 +130,11 @@ def chat():
         prompt = format_chat_prompt(message, chat_history, instructions)
         model_output = run_model(
             prompt,
+            model=model,
             temperature=temperature,
             top_p=top_p,
         )
-        model_output = model_output[len(prompt):].split(f"\n{USER_NAME}")[0]
+        model_output = model_output[len(prompt):].split(STOP_STR)[0]
         chat_history = chat_history + [[message, model_output]]
         yield chat_history
         return
@@ -137,15 +144,15 @@ def chat():
             chat_history.pop(-1)
         return {chatbot: gr.update(value=chat_history)}
 
-    def run_retry(message: str, chat_history, instructions: str, temperature: float, top_p: float):
-        yield from run_chat(RETRY_COMMAND, chat_history, instructions, temperature, top_p)
+    def run_retry(message: str, chat_history, instructions: str, model: str, temperature: float, top_p: float):
+        yield from run_chat(RETRY_COMMAND, chat_history, instructions, model: str, temperature, top_p)
 
     def clear_chat():
         return []
 
     inputs.submit(
         run_chat,
-        [inputs, chatbot, instructions, temperature, top_p],
+        [inputs, chatbot, instructions, model, temperature, top_p],
         outputs=[chatbot],
         show_progress=False,
     )
@@ -153,7 +160,7 @@ def chat():
     delete_turn_button.click(delete_last_turn, inputs=[chatbot], outputs=[chatbot])
     retry_button.click(
         run_retry,
-        [inputs, chatbot, instructions, temperature, top_p],
+        [inputs, chatbot, instructions, model, temperature, top_p],
         outputs=[chatbot],
         show_progress=False,
     )