Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import pipeline
+from transformers import pipeline, TextStreamer
 import gradio as gr
 import threading
 import time
@@ -55,6 +55,10 @@ def chat(message, history, temperature, max_new_tokens):
     start_time = time.time()
     generated_tokens = 0
     
+    # Create a TextStreamer for token streaming
+    tokenizer = model_pipeline.tokenizer
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
     def generate():
         nonlocal generated_tokens
         for response in model_pipeline(
@@ -63,9 +67,8 @@ def chat(message, history, temperature, max_new_tokens):
             temperature=temperature,
             do_sample=True,
             truncation=True,
-            pad_token_id=50256,
-            return_full_text=False,
-            streamer=True
+            pad_token_id=tokenizer.eos_token_id,
+            streamer=streamer  # Use the TextStreamer here
         ):
             bot_text = response[0]["generated_text"]
             bot_text = bot_text.split("Assistant:")[-1].strip()