Update app.py

app.py

@@ -2,9 +2,9 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
 import os
-import time
+from threading import Thread
 
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -69,17 +69,24 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
     
     input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
     
+    streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": True, "skip_prompt": True, 'clean_up_tokenization_spaces':False,})
 
-    gen_tokens= model.generate(
+    generate_kwargs = dict(
         input_ids, 
+        streamer=streamer,
         max_new_tokens=max_new_tokens, 
         do_sample=True, 
         temperature=temperature,
     )
+    
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
 
-    gen_text = tokenizer.batch_decode(gen_tokens[0], skip_special_tokens=True)
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield buffer
 
-    return gen_text
 
 
 chatbot = gr.Chatbot(height=450)