Changing the loop methodology

main.py

@@ -4,7 +4,7 @@ import re
 
 import gradio as gr
 import spaces
-from transformers import AsyncTextIteratorStreamer, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 JS = """
 () => {
@@ -46,6 +46,12 @@ print(model.config)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
+async def stream(streamer):
+    for txt in streamer:
+        await asyncio.sleep(0.01)
+        yield txt
+
+
 def reformat_math(text):
     """Fix MathJax delimiters to use the Gradio syntax.
 
@@ -58,7 +64,7 @@ def reformat_math(text):
 
 
 @spaces.GPU
-def _generate(history):
+def generate(history):
     text = tokenizer.apply_chat_template(
         history,
         tokenize=False,
@@ -72,7 +78,7 @@ def _generate(history):
         asyncio.set_event_loop(loop)
 
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-    streamer = AsyncTextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
 
     task = loop.run_in_executor(
         None,
@@ -97,15 +103,15 @@ async def chat(prompt, history):
     history = [] if history is None else history
     message_list = history + [message]
 
-    task, streamer = _generate(message_list)
+    task, streamer = generate(message_list)
 
     buffer = ""
     reasoning = ""
     thinking = False
 
     try:
-        async for new_text in streamer:
-            if task.done() or task.cancelled():
+        async for new_text in stream(streamer):
+            if task.cancelled():
                 print("Cancelled")
                 break  # Stop le streaming si la tâche est annulée
 
@@ -127,7 +133,7 @@ async def chat(prompt, history):
 
     except asyncio.CancelledError:
         # this doesn't work, I don't find a way to stop generation thread
-        print("Cancelled")
+        print("Cancelled by exception")
         streamer.on_finalized_text("cancelled", True)
         print("Signal sent")
         raise