Adding stop generating button to Prompt mode.

chat.py

@@ -96,11 +96,9 @@ def _generate(state, prompt, model, endseq, max_length,
             ):
 
             if not state['generate']:
-                print("Stopping generation.")
                 client.close_session()
                 yield state, [], None, ''
                 return
-                #return state, state['history'], None, prompt2 + output_raw
 
             output_raw += out
             output += out
@@ -168,8 +166,7 @@ with gr.Blocks() as iface_chat:
 
     with gr.Row():
         button_generate = gr.Button("Generate")
-        button_reset = gr.Button("Reset/Clear session")
-        # button_stop = gr.Button("Stop") # TODO, not supported by websocket API yet.
+        button_reset = gr.Button("Reset session")
 
     with gr.Accordion("Raw prompt log", open=False):
         output = gr.Textbox(lines=3, show_label=False).style(container=False)
@@ -187,7 +184,7 @@ with gr.Blocks() as iface_chat:
 
     examples = gr.Examples(inputs=[context, prompt, model, do_sample, top_k, top_p, temperature],
         examples=[
-        ["A Human talks to a powerful AI that follows the human's instructions. "
+        ["A Human talks to a powerful AI that follows the Human's instructions. "
          "AI is talkative, friendly, positive and provides detailed answers to any question.</s>\n"
          "Human: Hi!</s>\n"
          "AI: Hi! How can I help you?",

prompt.py

@@ -8,7 +8,17 @@ import chat_client
 CHAT_URL='ws://chat.petals.ml/api/v2/generate'
 #CHAT_URL='ws://localhost:8000/api/v2/generate'
 
-def generate(prompt, model, endseq, max_length,
+def generate(state, *args):
+    # Save that we're in generating loop
+    state['generate'] = True
+
+    try:
+        for x in _generate(state, *args):
+            yield x
+    finally:
+        state['generate'] = False
+
+def _generate(state, prompt, model, endseq, max_length,
         do_sample, top_k, top_p, temperature,
         add_stoptoken, copy_output):
 
@@ -17,7 +27,7 @@ def generate(prompt, model, endseq, max_length,
         client.open_session(f"bigscience/{model}-petals", max_length)
     except Exception:
         print(traceback.format_exc())
-        yield [prompt, "Error: " + traceback.format_exc()]
+        yield state, prompt, "Error: " + traceback.format_exc()
         return
 
     if add_stoptoken:
@@ -51,7 +61,7 @@ def generate(prompt, model, endseq, max_length,
 
     # This render prompt dialog immediately and
     # don't wait to generator to return first result
-    yield [prompt2, output]
+    yield [state, prompt2, output]
 
     try:
         for out in client.generate(prompt,
@@ -63,16 +73,25 @@ def generate(prompt, model, endseq, max_length,
                     extra_stop_sequences=seq
             ):
 
+            if not state['generate']:
+                client.close_session()
+                return
+
             output += out
             if copy_output:
                 prompt2 += out
 
-            yield [prompt2, output]
+            yield state, prompt2, output
     except Exception:
         print(traceback.format_exc())
-        yield [prompt, output + "\nError: " + traceback.format_exc()]
+        yield state, prompt, output + "\nError: " + traceback.format_exc()
         return
 
+def stop(state):
+    """Stops generating."""
+    state.update({"generate": False})
+    return state
+
 with gr.Blocks() as iface_prompt:
     gr.Markdown("""**Useful for testing raw prompts with zero, one or few-shot prompting.**""")
 
@@ -84,7 +103,7 @@ with gr.Blocks() as iface_prompt:
             value=["\\n", "</s>"], label='Extra end sequences')
 
         # Maximum length of inference session
-        max_length = gr.Radio([64, 128, 256, 512, 1024, 2048], value=128, interactive=True, label="Max length")
+        max_length = gr.Radio([64, 128, 256, 512, 1024, 2048], value=512, interactive=True, label="Max length")
 
     with gr.Row():
         with gr.Column():
@@ -103,26 +122,29 @@ with gr.Blocks() as iface_prompt:
         temperature = gr.Number(value=0.75, precision=2, interactive=True, label="Temperature")
 
     prompt = gr.Textbox(lines=3, label='Prompt', placeholder="Prompt Here...")
+    state = gr.State({'generate': False})
 
     with gr.Row():
         button_generate = gr.Button("Generate")
-        # button_stop = gr.Button("Stop") # TODO, not supported by websocket API yet.
+        button_stop = gr.Button("Stop")
 
         # Automatically copy the output at the end of prompt
         copy_output = gr.Checkbox(label="Output -> Prompt")
 
     output = gr.Textbox(lines=3, label='Output')
 
-    inputs = [prompt, model, endseq, max_length, do_sample,
+    inputs = [state, prompt, model, endseq, max_length, do_sample,
             top_k, top_p, temperature, add_stoptoken, copy_output]
-    outputs = [prompt, output]
+    outputs = [state, prompt, output]
     button_generate.click(generate, inputs=inputs, outputs=outputs)
+    button_stop.click(stop, inputs=[state], outputs=[state])
 
     examples = gr.Examples(inputs=[prompt, model, do_sample, top_k, top_p, temperature, add_stoptoken],
         examples=[
         ["The SQL command to extract all the users whose name starts with A is: ", "bloom-7b1", False, 0, 0, 1, False],
         ["The Spanish translation of thank you for your help is: ", "bloom-7b1", False, 0, 0, 1, False],
-        ["A human talks to a powerful AI that follows the human's instructions.</s>\n"
+        ["A human talks to a powerful AI that follows the Human's instructions.\n"
+         "AI is talkative, friendly, positive and provides detailed answers to any question.</s>\n"
          "Human: Hi!</s>\n"
          "AI: Hi! How can I help you?</s>\n"
          "Human: What's the capital of Portugal?</s>\n"