Fix richtextstreamer for multiple tokens input

app.py

@@ -16,17 +16,24 @@ 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
-        })
+        # Convert incoming tensor or list to flat list of token IDs
+        if isinstance(value, torch.Tensor):
+            token_ids = value.view(-1).tolist()
+        elif isinstance(value, list):
+            token_ids = value
+        else:
+            token_ids = [value]
+
+        for token_id in token_ids:
+            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: