Update app.py

app.py

@@ -13,6 +13,7 @@ from share_btn import community_icon_html, loading_icon_html, share_js, share_bt
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
 API_URL_G = "https://api-inference.huggingface.co/models/ArmelR/starcoder-gradio-v0/"
+API_URL_S = "https://api-inference.huggingface.co/models/HuggingFaceH4/starcoderbase-finetuned-oasst1"
 
 with open("./HHH_prompt_short.txt", "r") as f:
     HHH_PROMPT = f.read() + "\n\n"
@@ -53,6 +54,10 @@ client_g = Client(
     API_URL_G, headers={"Authorization": f"Bearer {HF_TOKEN}"},
 )
 
+client_starchat = Client(
+    API_URL_S, headers={"Authorization": f"Bearer {HF_TOKEN}"},
+)
+
 def generate(
     prompt,
     temperature=0.9,
@@ -60,7 +65,7 @@ def generate(
     top_p=0.95,
     repetition_penalty=1.0,
     chat_mode="TA prompt",
-    version=None,
+    version="StarChat-alpha",
 ):
 
     temperature = float(temperature)
@@ -90,14 +95,19 @@ def generate(
     chat_prompt = prompt + "\n\nAnswer:"
     prompt = base_prompt + chat_prompt
 
-    stream = client_g.generate_stream(prompt, **generate_kwargs)
+    if version == "StarCoder-gradio" :
+        stream = client_g.generate_stream(prompt, **generate_kwargs)
+    elif version == "StarChat-alpha" :
+        stream = client_s.generate_stream(prompt, **generate_kwargs)
+    else :
+        pass
         
     output = ""
     previous_token = ""
     
     for response in stream:
         if (
-            (response.token.text in ["Question:", "-----"]
+            (response.token.text in ["Human", "-----", "Question:"]
             and previous_token in ["\n", "-----"])
             or response.token.text == "<|endoftext|>"
         ):
@@ -121,12 +131,17 @@ def bot(
     top_p=0.95,
     repetition_penalty=1.0,
     chat_mode=None,
-    version=None, 
+    version="starchat-alpha", 
 ):
     # concat history of prompts with answers expect for last empty answer only add prompt
-    prompt = "\n".join(
-        [f"Question: {prompt}\n\nAnswer: {answer}" for prompt, answer in history[:-1]] + [f"\nQuestion: {history[-1][0]}"]
-    )
+    if version == "StarCoder-gradio"
+        prompt = "\n".join(
+            [f"Question: {prompt}\n\nAnswer: {answer}" for prompt, answer in history[:-1]] + [f"\nQuestion: {history[-1][0]}"]
+        )
+    else :
+        prompt = "\n".join(
+            [f"Human: {prompt}\n\nAssistant: {answer}" for prompt, answer in history[:-1]] + [f"\nHuman: {history[-1][0]}"]
+        )
     
     bot_message = generate(
         prompt,
@@ -210,12 +225,12 @@ _Note:_ this is an internal chat playground - **please do not share**. The deplo
                     interactive=True,
                     info="Penalize repeated tokens",
                 )
-                #version = gr.Dropdown(
-                #    ["StarCoderBase", "StarCoder"],
-                #    value="StarCoderBase",
-                #    label="Version",
-                #    info="",
-                #)
+                version = gr.Dropdown(
+                    ["StarCoder-gradio", "StarChat-alpha"],
+                    value="StarCoderBase",
+                    label="Version",
+                    info="",
+                )
             with column_1:
                 # output = gr.Code(elem_id="q-output")
                 # add visibl=False and update if chat_mode True
@@ -251,7 +266,7 @@ _Note:_ this is an internal chat playground - **please do not share**. The deplo
             user, [instruction, chatbot], [instruction, chatbot], queue=False
         ).then(
             bot,
-            [chatbot, temperature, max_new_tokens, top_p, repetition_penalty, chat_mode],
+            [chatbot, temperature, max_new_tokens, top_p, repetition_penalty, chat_mode, version],
             chatbot,
         )
 
@@ -259,7 +274,7 @@ _Note:_ this is an internal chat playground - **please do not share**. The deplo
             user, [instruction, chatbot], [instruction, chatbot], queue=False
         ).then(
             bot,
-            [chatbot, temperature, max_new_tokens, top_p, repetition_penalty, chat_mode],
+            [chatbot, temperature, max_new_tokens, top_p, repetition_penalty, chat_mode, version],
             chatbot,
         )
         clear.click(lambda: None, None, chatbot, queue=False)