Update app.py

app.py

@@ -12,8 +12,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_P = "https://api-inference.huggingface.co/models/bigcode/starcoder/"
-API_URL ="https://api-inference.huggingface.co/models/bigcode/starcoderbase/"
+API_URL_G = "https://api-inference.huggingface.co/models/ArmelR/starcoder-gradio-v0/"
 
 with open("./HHH_prompt_short.txt", "r") as f:
     HHH_PROMPT = f.read() + "\n\n"
@@ -21,6 +20,8 @@ with open("./HHH_prompt_short.txt", "r") as f:
 with open("./TA_prompt_v0.txt", "r") as f:
     TA_PROMPT = f.read()
 
+NO_PROMPT = ""
+
 FIM_PREFIX = "<fim_prefix>"
 FIM_MIDDLE = "<fim_middle>"
 FIM_SUFFIX = "<fim_suffix>"
@@ -48,16 +49,10 @@ theme = gr.themes.Monochrome(
     ],
 )
 
-client = Client(
-    API_URL, headers={"Authorization": f"Bearer {HF_TOKEN}"},
-)
-
-client_p = Client(
-    API_URL_P, headers={"Authorization": f"Bearer {HF_TOKEN}"},
+client_g = Client(
+    API_URL_G, headers={"Authorization": f"Bearer {HF_TOKEN}"},
 )
 
-
-
 def generate(
     prompt,
     temperature=0.9,
@@ -82,28 +77,27 @@ def generate(
         truncate=7500,
         do_sample=True,
         seed=42,
-        stop_sequences=["\nHuman", "\n-----"],
+        stop_sequences=["\nHuman", "\n-----", "Question:", "Answer:"],
     )
     
     if chat_mode == "HHH prompt":
         base_prompt = HHH_PROMPT
-    else:
+    elif chat_mode == "TA prompt":
         base_prompt = TA_PROMPT
+    else :
+        base_prompt = NO_PROMPT
 
-    chat_prompt = prompt + "\n\nAssistant:"
+    chat_prompt = prompt + "\n\nAnswer:"
     prompt = base_prompt + chat_prompt
 
-    if version == "StarCoderBase":
-        stream = client.generate_stream(prompt, **generate_kwargs)
-    else:
-        stream = client_p.generate_stream(prompt, **generate_kwargs)
+    stream = client_g.generate_stream(prompt, **generate_kwargs)
         
     output = ""
     previous_token = ""
     
     for response in stream:
         if (
-            (response.token.text in ["Human", "-----"]
+            (response.token.text in ["Question:", "-----"]
             and previous_token in ["\n", "-----"])
             or response.token.text == "<|endoftext|>"
         ):
@@ -131,7 +125,7 @@ def bot(
 ):
     # concat history of prompts with answers expect for last empty answer only add prompt
     prompt = "\n".join(
-        [f"Human: {prompt}\n\nAssistant: {answer}" for prompt, answer in history[:-1]] + [f"\nHuman: {history[-1][0]}"]
+        [f"Question: {prompt}\n\nAnswer: {answer}" for prompt, answer in history[:-1]] + [f"\nQuestion: {history[-1][0]}"]
     )
     
     bot_message = generate(
@@ -174,8 +168,8 @@ _Note:_ this is an internal chat playground - **please do not share**. The deplo
             column_1, column_2 = gr.Column(scale=3), gr.Column(scale=1)
             with column_2:
                 chat_mode = gr.Dropdown(
-                    ["TA prompt", "HHH prompt"],
-                    value="TA prompt",
+                    ["NO prompt","TA prompt", "HHH prompt"],
+                    value="NO prompt",
                     label="Chat mode",
                     info="Use Anthropic's HHH prompt or our custom tech prompt to turn the model into an assistant.",
                 )
@@ -215,12 +209,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(
+                #    ["StarCoderBase", "StarCoder"],
+                #    value="StarCoderBase",
+                #    label="Version",
+                #    info="",
+                #)
             with column_1:
                 # output = gr.Code(elem_id="q-output")
                 # add visibl=False and update if chat_mode True