Update app.py

app.py

@@ -3,7 +3,8 @@
 import spaces
 import gradio as gr
 from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
 
 # Load the base model
 base_model = AutoModelForCausalLM.from_pretrained(
@@ -29,19 +30,23 @@ def generate_response(prompt):
     )
     inputs = tokenizer(reasoning_prompt, return_tensors="pt").to(model.device)
     
-    # Streamed response
-    stream = model.generate(
+    # Using TextIteratorStreamer for streaming responses
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+
+    generation_kwargs = dict(
         **inputs,
-        max_new_tokens=300,  # Increased token limit
+        max_new_tokens=300,
         do_sample=True,
         temperature=0.8,
         top_p=0.95,
-        stream=True
+        streamer=streamer
     )
-    
-    # Yield output tokens in real-time
-    for chunk in stream:
-        yield tokenizer.decode(chunk[0], skip_special_tokens=True)
+
+    thread = torch.Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    for new_text in streamer:
+        yield new_text
 
 demo = gr.Interface(
     fn=generate_response,
@@ -52,4 +57,4 @@ demo = gr.Interface(
     live=True
 )
 
-demo.launch()
+demo.launch(share=True)