Update app.py

app.py

@@ -5,6 +5,8 @@ import threading
 import queue
 import time
 import spaces
+import sys
+from io import StringIO
 
 # Model configuration
 model_name = "HelpingAI/Dhanishtha-2.0-preview"
@@ -30,30 +32,30 @@ def load_model():
     
     print("Model loaded successfully!")
 
-class GradioTextStreamer(TextStreamer):
-    """Custom TextStreamer for Gradio integration"""
-    def __init__(self, tokenizer, skip_prompt=True):
-        # TextStreamer only accepts tokenizer and skip_prompt parameters
-        super().__init__(tokenizer, skip_prompt)
+class StreamCapture:
+    """Capture streaming output from TextStreamer"""
+    def __init__(self):
         self.text_queue = queue.Queue()
-        self.generated_text = ""
-        self.skip_special_tokens = True  # Handle this manually if needed
-        
-    def on_finalized_text(self, text: str, stream_end: bool = False):
-        """Called when text is finalized"""
-        self.generated_text += text
-        self.text_queue.put(text)
-        if stream_end:
-            self.text_queue.put(None)
-            
-    def get_generated_text(self):
-        """Get all generated text so far"""
-        return self.generated_text
+        self.captured_text = ""
         
+    def write(self, text):
+        """Capture written text"""
+        if text and text.strip():
+            self.captured_text += text
+            self.text_queue.put(text)
+        return len(text)
+    
+    def flush(self):
+        """Flush method for compatibility"""
+        pass
+    
+    def get_text(self):
+        """Get all captured text"""
+        return self.captured_text
+    
     def reset(self):
-        """Reset the streamer"""
-        self.generated_text = ""
-        # Clear the queue
+        """Reset the capture"""
+        self.captured_text = ""
         while not self.text_queue.empty():
             try:
                 self.text_queue.get_nowait()
@@ -89,11 +91,16 @@ def generate_response(message, history, max_tokens, temperature, top_p):
     # Tokenize input
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
     
-    # Create and setup streamer
-    streamer = GradioTextStreamer(tokenizer, skip_prompt=True)
-    streamer.reset()
+    # Create stream capture
+    stream_capture = StreamCapture()
     
-    # Start generation in a separate thread
+    # Create TextStreamer with our capture
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    # Temporarily redirect the streamer's output
+    original_stdout = sys.stdout
+    
+    # Generation parameters
     generation_kwargs = {
         **model_inputs,
         "max_new_tokens": max_tokens,
@@ -102,17 +109,21 @@ def generate_response(message, history, max_tokens, temperature, top_p):
         "do_sample": True,
         "pad_token_id": tokenizer.eos_token_id,
         "streamer": streamer,
-        "return_dict_in_generate": True
     }
     
-    # Run generation in thread
+    # Start generation in a separate thread
     def generate():
         try:
+            # Redirect stdout to capture streamer output
+            sys.stdout = stream_capture
             with torch.no_grad():
                 model.generate(**generation_kwargs)
         except Exception as e:
-            streamer.text_queue.put(f"Error: {str(e)}")
-            streamer.text_queue.put(None)
+            stream_capture.text_queue.put(f"Error: {str(e)}")
+        finally:
+            # Restore stdout
+            sys.stdout = original_stdout
+            stream_capture.text_queue.put(None)  # Signal end
     
     thread = threading.Thread(target=generate)
     thread.start()
@@ -121,7 +132,7 @@ def generate_response(message, history, max_tokens, temperature, top_p):
     generated_text = ""
     while True:
         try:
-            new_text = streamer.text_queue.get(timeout=30)
+            new_text = stream_capture.text_queue.get(timeout=30)
             if new_text is None:
                 break
             generated_text += new_text