Update app.py

app.py

@@ -1,6 +1,7 @@
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import gradio as gr
+from threading import Thread
 
 checkpoint = "WillHeld/soft-raccoon"
 device = "cuda"
@@ -13,20 +14,28 @@ def predict(message, history, temperature, top_p):
     input_text = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
     
-    streamer = gr.TelegramStreamer()  # Use Gradio's built-in streamer
+    # Create a streamer
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # Generate with streaming
-    model.generate(
-        inputs, 
-        max_new_tokens=1024, 
-        temperature=float(temperature), 
-        top_p=float(top_p), 
-        do_sample=True,
-        streamer=streamer
-    )
+    # Set up generation parameters
+    generation_kwargs = {
+        "input_ids": inputs,
+        "max_new_tokens": 1024,
+        "temperature": float(temperature),
+        "top_p": float(top_p),
+        "do_sample": True,
+        "streamer": streamer,
+    }
+    
+    # Run generation in a separate thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
     
-    # The streamer will handle returning the tokens
-    return streamer
+    # Yield from the streamer as tokens are generated
+    partial_text = ""
+    for new_text in streamer:
+        partial_text += new_text
+        yield partial_text
 
 with gr.Blocks() as demo:
     chatbot = gr.ChatInterface(