Enable ChatInterface stop button

app.py

@@ -5,7 +5,7 @@ from threading import Thread
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, TextIteratorStreamer
 
 DESCRIPTION = """\
 # Gemma 2 9B IT
@@ -32,6 +32,14 @@ model.config.sliding_window = 4096
 model.eval()
 
 
+class StopOnSignal(StoppingCriteria):
+    def __init__(self) -> None:
+        self.stopped = False
+
+    def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor, **kwargs: object) -> bool:  # noqa: ARG002
+        return self.stopped
+
+
 @spaces.GPU(duration=90)
 def _generate_on_gpu(
     input_ids: torch.Tensor,
@@ -44,9 +52,11 @@ def _generate_on_gpu(
     input_ids = input_ids.to(model.device)
 
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+    stop_criteria = StopOnSignal()
     generate_kwargs = {
         "input_ids": input_ids,
         "streamer": streamer,
+        "stopping_criteria": [stop_criteria],
         "max_new_tokens": max_new_tokens,
         "do_sample": True,
         "top_p": top_p,
@@ -69,9 +79,16 @@ def _generate_on_gpu(
     thread.start()
 
     chunks: list[str] = []
-    for text in streamer:
-        chunks.append(text)
-        yield "".join(chunks)
+    try:
+        for text in streamer:
+            chunks.append(text)
+            yield "".join(chunks)
+    except GeneratorExit:
+        stop_criteria.stopped = True
+        for _ in streamer:
+            pass
+        thread.join()
+        raise
 
     thread.join()
     if exception_holder:
@@ -165,7 +182,6 @@ demo = gr.ChatInterface(
             value=1.2,
         ),
     ],
-    stop_btn=False,
     examples=[
         ["Hello there! How are you doing?"],
         ["Can you explain briefly to me what is the Python programming language?"],