Update app.py

app.py

@@ -1,22 +1,17 @@
 import os
 import re
 import copy
+import time
 
 import gradio as gr
 from text_generation import Client
 from transformers import load_tool
-
-
 from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css
 
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 print(HF_TOKEN)
 
-API_URL = "https://api-inference.huggingface.co/models/vwxyzjn/starcoderbase-triviaqa"
-API_URL_BASE ="https://api-inference.huggingface.co/models/bigcode/starcoderbase"
-API_URL_PLUS = "https://api-inference.huggingface.co/models/bigcode/starcoderplus"
-
 FIM_PREFIX = "<fim_prefix>"
 FIM_MIDDLE = "<fim_middle>"
 FIM_SUFFIX = "<fim_suffix>"
@@ -77,13 +72,31 @@ theme = gr.themes.Monochrome(
     ],
 )
 
-client = Client(
-    API_URL,
-    headers={"Authorization": f"Bearer {HF_TOKEN}"},
-)
 tool = load_tool("vwxyzjn/pyserini-wikipedia-kilt-doc")
-tool_fn = lambda x: tool(x).split("\n")[1][:600] # limit the amount if tokens
-tools = {"Wiki": tool_fn}
+tool_fn = lambda x: tool(x).split("\n")[1][:600] # limit the amount if token, system_prompts
+
+clients = {
+    "StarCoderBase TriviaQA": [
+        Client(
+            "https://api-inference.huggingface.co/models/vwxyzjn/starcoderbase-triviaqa",
+            headers={"Authorization": f"Bearer {HF_TOKEN}"},
+        ),
+        {"Wiki": tool_fn},
+        """\
+Answer the following question:
+
+Q: In which branch of the arts is Patricia Neary famous?
+A: Ballets
+A2: <request><Wiki>Patricia Neary<call>Patricia Neary (born October 27, 1942) is an American ballerina, choreographer and ballet director, who has been particularly active in Switzerland. She has also been a highly successful ambassador for the Balanchine Trust, bringing George Balanchine's ballets to 60 cities around the globe.<response>
+Result=Ballets<submit>
+
+Q: Who won Super Bowl XX?
+A: Chicago Bears
+A2: <request><Wiki>Super Bowl XX<call>Super Bowl XX was an American football game between the National Football Conference (NFC) champion Chicago Bears and the American Football Conference (AFC) champion New England Patriots to decide the National Football League (NFL) champion for the 1985 season. The Bears defeated the Patriots by the score of 46–10, capturing their first NFL championship (and Chicago's first overall sports victory) since 1963, three years prior to the birth of the Super Bowl. Super Bowl XX was played on January 26, 1986 at the Louisiana Superdome in New Orleans.<response>
+Result=Chicago Bears<submit>
+"""
+    ],
+}
 
 def parse_tool_call(text, request_token="<request>", call_token="<call>"):
     """
@@ -113,9 +126,9 @@ def parse_tool_call(text, request_token="<request>", call_token="<call>"):
 
 
 def generate(
-    prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, version="StarCoderBase TriviaQA",
+    prompt, system_prompt, version, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,
 ):
-
+    client, tools, _ = clients[version]
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
@@ -132,48 +145,55 @@ def generate(
         seed=42,
         stop_sequences=["<call>"]
     )
+    generation_still_running = True
+    while generation_still_running:
+        try:
+            stream = client.generate_stream(system_prompt + prompt, **generate_kwargs)
+
+
+            # call env phase
+            output = system_prompt + prompt
+            previous_token = ""
+            for response in stream:
+                if response.token.text == "<|endoftext|>":
+                    return output
+                else:
+                    output += response.token.text
+                previous_token = response.token.text
+                # text env logic:
+                tool, query = parse_tool_call(output[len(system_prompt + prompt):])
+                print("tool", tool, query)
+                if tool is not None and query is not None:
+                    if tool not in tools:
+                        response = f"Unknown tool {tool}."
+                    try:
+                        response = tools[tool](query)
+                        output += response + "<response>"
+                    except Exception as error:
+                        response = f"Tool error: {str(error)}"
+                yield output[len(system_prompt + prompt):]
+
+            call_output = copy.deepcopy(output)
+            # response phase
+            generate_kwargs["stop_sequences"] = ["<submit>"]
+            stream = client.generate_stream(output, **generate_kwargs)
+            previous_token = ""
+            for response in stream:
+                if response.token.text == "<|endoftext|>":
+                    return output
+                else:
+                    output += response.token.text
+                previous_token = response.token.text
+                yield output[len(system_prompt + prompt):]
 
-    if version == "StarCoderBase TriviaQA":
-        stream = client.generate_stream(prompt, **generate_kwargs)
-
-
-    # call env phase
-    output = prompt
-    previous_token = ""
-    for response in stream:
-        if response.token.text == "<|endoftext|>":
-            return output
-        else:
-            output += response.token.text
-        previous_token = response.token.text
-        # text env logic:
-        tool, query = parse_tool_call(output[len(prompt):])
-        if tool is not None and query is not None:
-            if tool not in tools:
-                response = f"Unknown tool {tool}."
-            try:
-                response = tools[tool](query)
-                output += response + "<response>"
-            except Exception as error:
-                response = f"Tool error: {str(error)}"
-        yield output[len(prompt):]
-
-    call_output = copy.deepcopy(output)
-    # response phase
-    generate_kwargs["stop_sequences"] = ["<submit>"]
-    stream = client.generate_stream(output, **generate_kwargs)
-    previous_token = ""
-    for response in stream:
-        if response.token.text == "<|endoftext|>":
             return output
-        else:
-            output += response.token.text
-        previous_token = response.token.text
-        yield output[len(prompt):]
-
-
-
-    return output
+        except Exception as e:
+            if "loading" in str(e):
+                gr.Warning("waiting for model to load... (this could take up to 20 minutes, after which things are much faster)")
+                time.sleep(7)
+                continue
+            else:
+                raise gr.Error(str(e))           
 
 
 examples = [
@@ -224,15 +244,21 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
         gr.Markdown(description)
         with gr.Row():
             version = gr.Dropdown(
-                        ["StarCoderBase TriviaQA"],
-                        value="StarCoderBase TriviaQA",
+                        list(clients.keys()),
+                        value=list(clients.keys())[0],
                         label="Model",
                         info="Choose a model from the list",
                         )
+            system_prompt = gr.Textbox(
+                value=clients[list(clients.keys())[0]][2],
+                label="System prompt",
+            )
+            
         with gr.Row():
             with gr.Column():
                 instruction = gr.Textbox(
-                    placeholder="Enter your code here",
+                    value="Q: In which country is Oberhofen situated?",
+                    # placeholder="Enter your question here. E.g., Q: In which country is Oberhofen situated?",
                     lines=5,
                     label="Input",
                     elem_id="q-input",
@@ -297,15 +323,15 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
                     fn=process_example,
                     outputs=[output],
                 )
-                gr.Markdown(FORMATS)
+                # gr.Markdown(FORMATS)
 
     submit.click(
         generate,
-        inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty, version],
+        inputs=[instruction, system_prompt, version, temperature, max_new_tokens, top_p, repetition_penalty],
         outputs=[output],
     )
     share_button.click(None, [], [], _js=share_js)
-demo.queue(concurrency_count=16).launch(debug=True)
+demo.queue(concurrency_count=16).launch(share=True)
 
 
 """