Implement richtextiteratorstreamer

app.py

@@ -9,6 +9,34 @@ import threading
 from transformers import TextIteratorStreamer
 import threading
 
+from transformers import TextIteratorStreamer
+import queue
+
+class RichTextStreamer(TextIteratorStreamer):
+    def __init__(self, tokenizer, **kwargs):
+        super().__init__(tokenizer, **kwargs)
+        self.token_queue = queue.Queue()
+    
+    def put(self, value):
+        # Instead of just decoding here, we emit full info per token
+        token_id = value.item() if hasattr(value, "item") else value
+        token_str = self.tokenizer.decode([token_id], **self.decode_kwargs)
+        is_special = token_id in self.tokenizer.all_special_ids
+        self.token_queue.put({
+            "token_id": token_id,
+            "token": token_str,
+            "is_special": is_special
+        })
+
+    def __iter__(self):
+        while True:
+            try:
+                token_info = self.token_queue.get(timeout=self.timeout)
+                yield token_info
+            except queue.Empty:
+                if self.end_of_generation.is_set():
+                    break
+
 
 @spaces.GPU
 def chat_with_model(messages):
@@ -31,7 +59,9 @@ def chat_with_model(messages):
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
 
-    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
+    # streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
+    streamer = RichTextStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
+
 
     generation_kwargs = dict(
         **inputs,
@@ -49,16 +79,16 @@ def chat_with_model(messages):
     messages = messages.copy()
     messages.append({"role": "assistant", "content": ""})
 
-    for new_text in streamer:
-        output_text += new_text
-        if "\nUser:" in output_text:
-            output_text = output_text.split("\nUser:")[0].rstrip()
-            messages[-1]["content"] = output_text
-            yield messages
-            break
+    for token_info in streamer:
+        token_str = token_info["token"]
+        is_special = token_info["is_special"]
+        output_text += token_str
         messages[-1]["content"] = output_text
         yield messages
 
+        if is_special and token_info["token_id"] == current_tokenizer.eos_token_id:
+            break
+
     current_model.to("cpu")
     torch.cuda.empty_cache()