Fix cancel generation to gracefully stop ongoing response generation

- Add StoppingCriteria import and CancelStoppingCriteria class
- Integrate stopping criteria into generation pipeline to halt token generation when cancel event is set
- Ensures generation stops at the model level, preventing unnecessary GPU usage

app.py

@@ -7,7 +7,7 @@ from datetime import datetime
 import re  # for parsing <think> blocks
 import gradio as gr
 import torch
-from transformers import pipeline, TextIteratorStreamer
+from transformers import pipeline, TextIteratorStreamer, StoppingCriteria
 from transformers import AutoTokenizer
 from ddgs import DDGS
 import spaces  # Import spaces early to enable ZeroGPU support
@@ -23,6 +23,13 @@ access_token=os.environ['HF_TOKEN']
 # ------------------------------
 cancel_event = threading.Event()
 
+# ------------------------------
+# Stopping Criteria for Cancellation
+# ------------------------------
+class CancelStoppingCriteria(StoppingCriteria):
+    def __call__(self, input_ids, scores, **kwargs):
+        return cancel_event.is_set()
+
 # ------------------------------
 # Torch-Compatible Model Definitions with Adjusted Descriptions
 # ------------------------------
@@ -499,6 +506,7 @@ def chat_response(user_msg, chat_history, system_prompt,
                 'top_p': top_p,
                 'repetition_penalty': repeat_penalty,
                 'streamer': streamer,
+                'stopping_criteria': [CancelStoppingCriteria()],
                 'return_full_text': False,
             }
         )