Update app.py

app.py

@@ -1,13 +1,13 @@
 import os
 import time
 import gc
-import threading
+from queue import Queue
+from threading import Thread, Event
 from itertools import islice
 from datetime import datetime
 import re  # for parsing <think> blocks
 import gradio as gr
 import torch
-from transformers import TextIteratorStreamer
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from duckduckgo_search import DDGS
 # import spaces  # Import spaces early to enable ZeroGPU support
@@ -23,7 +23,7 @@ else:
 # ------------------------------
 # Global Cancellation Event
 # ------------------------------
-cancel_event = threading.Event()
+cancel_event = Event()
 
 # ------------------------------
 # Torch-Compatible Model Definitions with Adjusted Descriptions
@@ -38,6 +38,42 @@ MODELS = {
 # Global cache for pipelines to avoid re-loading.
 PIPELINES = {}
 
+class TextIterStreamer:
+    def __init__(self, tokenizer, skip_prompt=True, skip_special_tokens=True):
+        self.tokenizer = tokenizer
+        self.skip_prompt = skip_prompt
+        self.skip_special_tokens = skip_special_tokens
+        self.tokens = []
+        self.text_queue = Queue()
+        # self.text_queue = []
+        self.next_tokens_are_prompt = True
+
+    def put(self, value):
+        if self.skip_prompt and self.next_tokens_are_prompt:
+            self.next_tokens_are_prompt = False
+        else:
+            if len(value.shape) > 1:
+                value = value[0]
+            self.tokens.extend(value.tolist())
+            word = self.tokenizer.decode(self.tokens, skip_special_tokens=self.skip_special_tokens)
+            # self.text_queue.append(word)
+            self.text_queue.put(word)
+
+    def end(self):
+        # self.text_queue.append(None)
+        self.text_queue.put(None)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value = self.text_queue.get()
+        if value is None:
+            raise StopIteration()
+        else:
+            return value
+
+
 def load_pipeline(model_name):
     """
     Load and cache a transformers pipeline for text generation.
@@ -101,7 +137,7 @@ def chat_response(user_msg, chat_history, system_prompt,
     search_results = []
     if enable_search:
         debug = 'Search task started.'
-        thread_search = threading.Thread(
+        thread_search = Thread(
             target=lambda: search_results.extend(
                 retrieve_context(user_msg, int(max_results), int(max_chars))
             )
@@ -142,20 +178,20 @@ def chat_response(user_msg, chat_history, system_prompt,
                                         skip_prompt=True,
                                         skip_special_tokens=True)
         generation_config = dict(
-                temperature=temperature,
-                top_k=top_k,
-                top_p=top_p,
-                max_new_tokens=max_tokens,
-                do_sample=True,
-                repetition_penalty=repeat_penalty,
-                streamer=streamer,
-            )
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            repetition_penalty=repeat_penalty,
+            streamer=streamer,
+        )
         inputs = pipe["tokenizer"](prompt, return_tensors="pt")
         if device == "auto":
             input_ids = inputs["input_ids"].cuda()
         else:
             input_ids = inputs["input_ids"]
-        gen_thread = threading.Thread(target=lambda: pipe["model"].generate(input_ids=input_ids, **generation_config))
+        gen_thread = Thread(target=lambda: pipe["model"].generate(input_ids=input_ids, **generation_config))
         gen_thread.start()
 
         # Buffers for thought vs answer