Update app.py

app.py

@@ -9,7 +9,7 @@ import os
 import time
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
 from threading import Thread
 
@@ -83,28 +83,25 @@ def stream_chat(
     conversation.append({"role": "user", "content": message})
 
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(device)
-    
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-    
-    generate_kwargs = dict(
+        
+    gen_out = model.multi_byte_generate(
         input_ids=input_ids, 
         max_new_tokens = max_new_tokens,
         do_sample = False if temperature == 0 else True,
         top_p = top_p,
         temperature = temperature,
-        streamer=streamer,
     )
 
-    with torch.no_grad():
-        thread = Thread(target=model.multi_byte_generate, kwargs=generate_kwargs)
-        thread.start()
-        
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
+    response = tokenizer.decode(
+        gen_out[0][input_ids.shape[1]:], 
+        skip_special_tokens=False,
+        clean_up_tokenization_spaces=False
+    )
 
-            
+    for i in range(len(response)):
+        time.sleep(0.05)
+        yield response[: i + 1]
+                
 chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)
 
 with gr.Blocks(css=CSS, theme="soft") as demo: