Update app.py

app.py

@@ -147,7 +147,7 @@ def generate(
             output = system_prompt + prompt
             generation_start_idx = len(output)
             highlighted_output = [
-                (prompt, "query"),
+                (prompt, "QUERY"),
             ]
             yield highlighted_output, output[generation_start_idx:]
             for response in stream:
@@ -181,13 +181,13 @@ def generate(
                 # I am sorry about the code
                 print("-------", generation_start_idx, request_idx, call_idx, response_idx)
                 highlighted_output = [
-                    (prompt, "query"),
-                    (output[generation_start_idx:], "model") if request_idx == -1 else ("", ""),
-                    (output[generation_start_idx:generation_start_idx+request_idx], "model"),
-                    (output[generation_start_idx+request_idx:], "model") if call_idx == -1 else  ("", ""),
-                    (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "tool request"),
-                    (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "tool call"),
-                    (output[generation_start_idx+response_idx:], "model") if submit_idx != -1 else ("", ""),
+                    (prompt, "QUERY"),
+                    (output[generation_start_idx:], "MODEL") if request_idx == -1 else ("", ""),
+                    (output[generation_start_idx:generation_start_idx+request_idx], "MODEL"),
+                    (output[generation_start_idx+request_idx:], "MODEL") if call_idx == -1 else  ("", ""),
+                    (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "TOOL_REQUEST"),
+                    (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "TOOL_CALL"),
+                    (output[generation_start_idx+response_idx:], "MODEL") if submit_idx != -1 else ("", ""),
                     # (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:], "call"),
@@ -207,11 +207,11 @@ def generate(
                     submit_idx = output[generation_start_idx:].find("<submit>")
                 # print("-------", generation_start_idx, request_idx, call_idx, response_idx)
                 highlighted_output = [
-                    (prompt, "query"),
-                    (output[generation_start_idx:generation_start_idx+request_idx], "model"),
-                    (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:], "model") if submit_idx != -1 else ("", ""),
+                    (prompt, "QUERY"),
+                    (output[generation_start_idx:generation_start_idx+request_idx], "MODEL"),
+                    (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "TOOL_REQUEST"),
+                    (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "TOOL_CALL"),
+                    (output[generation_start_idx+response_idx:], "MODEL") if submit_idx != -1 else ("", ""),
                 ]
                 # print(highlighted_output, output[generation_start_idx:])
                 yield highlighted_output, output[generation_start_idx:]
@@ -291,9 +291,9 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
                 
                 output = gr.HighlightedText(
                     label="Output",
-                    color_map={"query": "red", "tool call": "green", "tool response": "blue", "model": "pink"},
+                    color_map={"QUERY": "red", "TOOL_CALL": "green", "TOOL_RESPONSE": "blue", "MODEL": "pink"},
                 )
-
+                gr.Markdown("_Note:_ The trivia model is trained to give an answer first and then refine it with a Wiki call.")
                 gr_examples = gr.Examples(
                     examples=[example for client in clients.values() for example in client[3]],
                     inputs=[instruction],