Add stop button

app.py

@@ -22,6 +22,7 @@ pipe = pipeline("image-to-text", model=model_id, model_kwargs={"quantization_con
 def extract_response_pairs(text):
     turns = re.split(r'(USER:|ASSISTANT:)', text)[1:]
     turns = [turn.strip() for turn in turns if turn.strip()]
+    print(f"conv turns are {turns[1::2]}")
     conv_list = []
     for i in range(0, len(turns[1::2]), 2):
         if i + 1 < len(turns[1::2]):
@@ -63,7 +64,7 @@ def bot(history_chat, text_input, image,
             min_length,
             top_p):
   chat_history = " ".join(history_chat) # history as a str to be passed to model
-  chat_history = chat_history + f"USER: <image>\n{text_input}\nASSISTANT:" # add text input for prompting
+  chat_history = chat_history + f"USER: <image>\n{text_input}\nASSISTANT:" # add text input for prompting
 
 
   inference_result = infer(image, chat_history,
@@ -76,16 +77,16 @@ def bot(history_chat, text_input, image,
   # return inference and parse for new history
   chat_val = extract_response_pairs(inference_result)
 
-  # create history list for yielding the last inference response
+  # create history list for yielding the last inference response
   chat_state_list = copy.deepcopy(chat_val)
-  chat_state_list[-1][1] = "" # empty last response
+  chat_state_list[-1][1] = "" # empty last response
 
   # add characters iteratively
 
   for character in chat_val[-1][1]:
     chat_state_list[-1][1] += character
     time.sleep(0.05)
-    # yield history but with last response being streamed
+    # yield history but with last response being streamed
     print(chat_state_list)
     yield chat_state_list
 
@@ -178,9 +179,10 @@ with gr.Blocks(css="style.css") as demo:
     ]
     with gr.Row():
       clear_chat_button = gr.Button("Clear")
+      cancel_btn = gr.Button("Stop Generation")
       chat_button = gr.Button("Submit", variant="primary")
       
-      chat_button.click(add_text, [chatbot, text_input], [chatbot, text_input]).then(bot, [chatbot, text_input,
+    chat_event1 = chat_button.click(add_text, [chatbot, text_input], [chatbot, text_input]).then(bot, [chatbot, text_input,
                                                                                            image, temperature,
         length_penalty,
         repetition_penalty,
@@ -188,7 +190,7 @@ with gr.Blocks(css="style.css") as demo:
         min_length,
         top_p], chatbot)
     
-    text_input.submit(
+    chat_event2 = text_input.submit(
         add_text,
         [chatbot, text_input],
         [chatbot, text_input]
@@ -220,6 +222,10 @@ with gr.Blocks(css="style.css") as demo:
             history_chat
         ],
         queue=False)
+    cancel_btn.click(
+        None, [], [], 
+        cancels=[chat_event1, chat_event2]
+    )  
     examples = [["./examples/baklava.png", "How to make this pastry?"],["./examples/bee.png","Describe this image."]]
     gr.Examples(examples=examples, inputs=[image, text_input, chat_inputs])
 
@@ -227,4 +233,4 @@ with gr.Blocks(css="style.css") as demo:
     
 
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch(debug=True)
+    demo.queue(max_size=10).launch(debug=True)