Update app.py

app.py

@@ -1,14 +1,15 @@
-import json
 import os
-import shutil
-import requests
+import re
+import copy
 
 import gradio as gr
-from huggingface_hub import Repository
 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)
 
@@ -80,15 +81,39 @@ client = Client(
     API_URL,
     headers={"Authorization": f"Bearer {HF_TOKEN}"},
 )
-client_base = Client(
-    API_URL_BASE, headers={"Authorization": f"Bearer {HF_TOKEN}"},
-)
-client_plus = Client(
-    API_URL_PLUS, 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}
+
+def parse_tool_call(text, request_token="<request>", call_token="<call>"):
+    """
+    Parse request string. Expected format: <request><tool_name>query<call>
+    """
+    result = re.search(f"(?<={request_token}).*?(?={call_token})", text, re.DOTALL)
+
+    # if we can't find a <request>/<call> span we return none
+    if result is None:
+        return None, None
+    else:
+        extracted_text = result.group()
+
+    result = re.search(r"<(.*?)>", extracted_text)
+
+    # if we can't find a tool name we return none
+    if result is None:
+        return None, None
+    else:
+        tool = result.group(1)
+
+    # split off the tool name
+    query = ">".join(extracted_text.split(">")[1:])
+
+    return tool, query
+
+
 
 def generate(
-    prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, version="StarCoder",
+    prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, version="StarCoderBase TriviaQA",
 ):
 
     temperature = float(temperature)
@@ -97,6 +122,7 @@ def generate(
     top_p = float(top_p)
     fim_mode = False
 
+    # TextEnv tool
     generate_kwargs = dict(
         temperature=temperature,
         max_new_tokens=max_new_tokens,
@@ -104,39 +130,49 @@ def generate(
         repetition_penalty=repetition_penalty,
         do_sample=True,
         seed=42,
+        stop_sequences=["<call>"]
     )
 
-    if FIM_INDICATOR in prompt:
-        fim_mode = True
-        try:
-            prefix, suffix = prompt.split(FIM_INDICATOR)
-        except:
-            raise ValueError(f"Only one {FIM_INDICATOR} allowed in prompt!")
-        prompt = f"{FIM_PREFIX}{prefix}{FIM_SUFFIX}{suffix}{FIM_MIDDLE}"
-
-    if version == "StarCoder":
+    if version == "StarCoderBase TriviaQA":
         stream = client.generate_stream(prompt, **generate_kwargs)
-    elif version == "StarCoderPlus":
-        stream = client_plus.generate_stream(prompt, **generate_kwargs)
-    else:
-        stream = client_base.generate_stream(prompt, **generate_kwargs)
 
-    if fim_mode:
-        output = prefix
-    else:
-        output = prompt
 
+    # call env phase
+    output = prompt
     previous_token = ""
     for response in stream:
         if response.token.text == "<|endoftext|>":
-            if fim_mode:
-                output += suffix
-            else:
-                return output
+            return output
         else:
             output += response.token.text
         previous_token = response.token.text
-        yield output
+        # 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
 
 
@@ -168,14 +204,14 @@ css += share_btn_css + monospace_css + ".gradio-container {color: black}"
 
 description = """
 <div style="text-align: center;">
-    <h1> ⭐ StarCoder <span style='color: #e6b800;'>Models</span> Playground</h1>
+    <h1> ⭐ StarCoderBase TriviaQA <span style='color: #e6b800;'>Models</span> Playground</h1>
 </div>
 <div style="text-align: left;">
-    <p>This is a demo to generate text and code with the following StarCoder models:</p>
+    <p>This is a demo to generate text and code with the following StarCoderBase TriviaQA models:</p>
     <ul>
         <li><a href="https://huggingface.co/bigcode/starcoderplus" style='color: #e6b800;'>StarCoderPlus</a>: A finetuned version of StarCoderBase on English web data, making it strong in both English text and code generation.</li>
         <li><a href="https://huggingface.co/bigcode/starcoderbase" style='color: #e6b800;'>StarCoderBase</a>: A code generation model trained on 80+ programming languages, providing broad language coverage for code generation tasks.</li>
-        <li><a href="https://huggingface.co/bigcode/starcoder" style='color: #e6b800;'>StarCoder</a>: A finetuned version of StarCoderBase specifically focused on Python, while also maintaining strong performance on other programming languages.</li>
+        <li><a href="https://huggingface.co/bigcode/starcoder" style='color: #e6b800;'>StarCoderBase TriviaQA</a>: A finetuned version of StarCoderBase specifically focused on Python, while also maintaining strong performance on other programming languages.</li>
     </ul>
     <p><b>Please note:</b> These models are not designed for instruction purposes. If you're looking for instruction or want to chat with a fine-tuned model, you can visit the <a href="https://huggingface.co/spaces/HuggingFaceH4/starchat-playground">StarChat Playground</a>.</p>
 </div>
@@ -188,8 +224,8 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
         gr.Markdown(description)
         with gr.Row():
             version = gr.Dropdown(
-                        ["StarCoderPlus", "StarCoderBase", "StarCoder"],
-                        value="StarCoder",
+                        ["StarCoderBase TriviaQA"],
+                        value="StarCoderBase TriviaQA",
                         label="Model",
                         info="Choose a model from the list",
                         )
@@ -269,4 +305,20 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
         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)
+
+
+"""
+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>
+
+Q: In what state is Philadelphia located?"""