Update app.py

app.py

@@ -159,7 +159,7 @@ def generate(
         repetition_penalty=repetition_penalty,
         do_sample=True,
         seed=42,
-        stop_sequences=["<call>"]
+        stop_sequences=["<call>", "<submit>"]
     )
     generation_still_running = True
     request_idx = -1
@@ -179,13 +179,10 @@ def generate(
             highlighted_output = [
                 (prompt, ""),
             ]
-            yield highlighted_output
+            yield highlighted_output, output[generation_start_idx:]
             for response in stream:
                 i += 1
-                if response.token.text == "<|endoftext|>":
-                    return output
-                else:
-                    output += response.token.text
+                output += response.token.text
                 tool, query = parse_tool_call(output[generation_start_idx:])
                 
                 if tool is not None and query is not None:
@@ -205,17 +202,24 @@ def generate(
                     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)
+                if submit_idx == -1:
+                    submit_idx = output[generation_start_idx:].find("<submit>")
+                # I am sorry about the code
+                print("-------", generation_start_idx, request_idx, call_idx, response_idx)
                 highlighted_output = [
                     (prompt, ""),
+                    (output[generation_start_idx:], "") if request_idx == -1 else ("", ""),
                     (output[generation_start_idx:generation_start_idx+request_idx], ""),
+                    (output[generation_start_idx+request_idx:], "") if call_idx == -1 else  ("", ""),
                     (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "request"),
-                    (output[generation_start_idx+call_idx:-1], "call"),
+                    (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "call"),
+                    (output[generation_start_idx+response_idx:], "submit") 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"),
                 ]
-                # print(i, highlighted_output, output)
-                yield highlighted_output
+                print(i, highlighted_output, output[generation_start_idx:])
+                yield highlighted_output, output[generation_start_idx:]
 
             # breakpoint()
             call_output = copy.deepcopy(output)
@@ -224,26 +228,23 @@ def generate(
             generate_kwargs["stop_sequences"] = ["<submit>"]
             stream = client.generate_stream(output, **generate_kwargs)
             for response in stream:
-                if response.token.text == "<|endoftext|>":
-                    return output
-                else:
-                    output += response.token.text
+                output += response.token.text
                 if submit_idx == -1:
                     submit_idx = output[generation_start_idx:].find("<submit>")
-                # print("-------", generation_start_idx, request_idx, call_idx, response_idx, submit_idx)
+                # 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:generation_start_idx+response_idx], "call"),
-                    (output[generation_start_idx+response_idx:-1], "submit"),
+                    (output[generation_start_idx+response_idx:], "submit") if submit_idx != -1 else ("", ""),
                 ]
-                # print(highlighted_output, output)
-                yield highlighted_output
-            print("-------", generation_start_idx, request_idx, call_idx, response_idx, submit_idx)
-            print(highlighted_output, output)
+                # print(highlighted_output, output[generation_start_idx:])
+                yield highlighted_output, output[generation_start_idx:]
+            print("-------", generation_start_idx, request_idx, call_idx, response_idx)
+            print(highlighted_output, output[generation_start_idx:])
 
-            return highlighted_output
+            return highlighted_output, output[generation_start_idx:]
         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)")
@@ -281,14 +282,13 @@ css += share_btn_css + monospace_css + ".gradio-container {color: black}"
 
 description = """
 <div style="text-align: center;">
-    <h1> ⭐ StarCoderBase TriviaQA <span style='color: #e6b800;'>Models</span> Playground</h1>
+    <h1> ⭐ TRL + TextEnvironment <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 StarCoderBase TriviaQA models:</p>
+    <p>This is a demo to generate text and code with the following StarCoderBase 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;'>StarCoderBase TriviaQA</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/starcoderplus" style='color: #e6b800;'>StarCoderBase TriviaQA</a>: A finetuned version of StarCoderBase on on the TriviaQA dataset using reinforcement learning via TRL's TextEnvironment (https://github.com/huggingface/trl/pull/424)</li>
+        <li><a href="https://huggingface.co/bigcode/starcoderbase" style='color: #e6b800;'>StarCoderBase GSM8K</a>: A finetuned version of StarCoderBase on on the GSM8K dataset using reinforcement learning via TRL's TextEnvironment (https://github.com/huggingface/trl/pull/424).</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>
@@ -326,11 +326,12 @@ 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.HighlightedText(
                     label="Output",
                     color_map={"query": "red", "call": "green", "response": "blue", "submit": "yellow", "model": "pink"},
                 )
+                output2 = gr.Code(elem_id="q-output", lines=30, label="Raw output")
                 with gr.Row():
                     with gr.Column():
                         with gr.Accordion("Advanced settings", open=False):
@@ -387,14 +388,14 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
                     inputs=[instruction],
                     cache_examples=False,
                     fn=process_example,
-                    outputs=[output],
+                    outputs=[output, output2],
                 )
                 # gr.Markdown(FORMATS)
 
     submit.click(
         generate,
         inputs=[instruction, system_prompt, version, temperature, max_new_tokens, top_p, repetition_penalty],
-        outputs=[output],
+        outputs=[output, output2],
     )
     share_button.click(None, [], [], _js=share_js)
 demo.queue(concurrency_count=16).launch(debug=True)