Update app.py

app.py

@@ -94,6 +94,21 @@ 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>
+"""
+    ],
+    "StarCoderBase TriviaQA2": [
+        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>
 """
     ],
 }
@@ -146,6 +161,11 @@ def generate(
         stop_sequences=["<call>"]
     )
     generation_still_running = True
+    request_idx = -1
+    call_idx = -1
+    response_idx = -1
+    submit_idx = -1
+    
     while generation_still_running:
         try:
             stream = client.generate_stream(system_prompt + prompt, **generate_kwargs)
@@ -153,16 +173,19 @@ def generate(
 
             # call env phase
             output = system_prompt + prompt
-            previous_token = ""
+            generation_start_idx = len(output)
+            highlighted_output = [
+                (prompt, ""),
+            ]
+            yield highlighted_output
             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)
+                print(response.token.text)
+                tool, query = parse_tool_call(output[generation_start_idx:])
+                # print("tool", tool, query)
                 if tool is not None and query is not None:
                     if tool not in tools:
                         response = f"Unknown tool {tool}."
@@ -171,22 +194,48 @@ def generate(
                         output += response + "<response>"
                     except Exception as error:
                         response = f"Tool error: {str(error)}"
-                yield output[len(system_prompt + prompt):]
+
+                if request_idx == -1:
+                    request_idx = output[generation_start_idx:].find("<request>")
+                if call_idx == -1:
+                    call_idx = output[generation_start_idx:].find("<call>")
+                if response_idx == -1:
+                    response_idx = output[generation_start_idx:].find("<response>")
+
+                # if `<request>` is in the output, highlight it, if `<call>` is in the output, highlight it
+                # print(generation_start_idx, request_idx, call_idx, response_idx)
+                highlighted_output = [
+                    (prompt, ""),
+                    (output[generation_start_idx:generation_start_idx+request_idx], ""),
+                    (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "request"),
+                    (output[generation_start_idx+call_idx:-1], "call"),
+                ]
+                # print(highlighted_output, output)
+                yield highlighted_output
 
             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):]
-
-            return output
+                if submit_idx == -1:
+                    submit_idx = output[generation_start_idx:].find("<submit>")
+                # print(generation_start_idx, request_idx, call_idx, response_idx, submit_idx)
+                highlighted_output = [
+                    (prompt, ""),
+                    (output[generation_start_idx:generation_start_idx+request_idx], ""),
+                    (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "request"),
+                    (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "call"),
+                    (output[generation_start_idx+response_idx:-1], "submit"),
+                ]
+                # print(highlighted_output, output)
+                yield highlighted_output
+
+            return highlighted_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)")
@@ -264,7 +313,11 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
                     elem_id="q-input",
                 )
                 submit = gr.Button("Generate", variant="primary")
-                output = gr.Code(elem_id="q-output", lines=30, label="Output")
+                # output = gr.Code(elem_id="q-output", lines=30, label="Output")
+                output = gr.HighlightedText(
+                    label="Output",
+                    color_map={"query": "red", "call": "green", "response": "blue", "submit": "yellow", "model": "pink"},
+                )
                 with gr.Row():
                     with gr.Column():
                         with gr.Accordion("Advanced settings", open=False):