Update app.py

app.py

@@ -4,6 +4,9 @@ os.system("pip install git+https://github.com/shumingma/transformers.git")
 
 import threading
 import torch
+import torch._dynamo
+torch._dynamo.config.suppress_errors = True
+
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -32,6 +35,21 @@ def respond(
     temperature: float,
     top_p: float,
 ):
+    """
+    Generate a chat response using streaming with TextIteratorStreamer.
+
+    Args:
+        message: User's current message.
+        history: List of (user, assistant) tuples from previous turns.
+        system_message: Initial system prompt guiding the assistant.
+        max_tokens: Maximum number of tokens to generate.
+        temperature: Sampling temperature.
+        top_p: Nucleus sampling probability.
+
+    Yields:
+        The growing response text as new tokens are generated.
+    """
+    # Assemble messages
     messages = [{"role": "system", "content": system_message}]
     for user_msg, bot_msg in history:
         if user_msg:
@@ -40,18 +58,33 @@ def respond(
             messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    # Prepare prompt and tokenize
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    outputs = model.generate(
+    # Set up streamer for real-time output
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+    generate_kwargs = dict(
         **inputs,
+        streamer=streamer,
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-        do_sample=True
+        do_sample=True,
     )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    yield response
+    # Start generation in a separate thread
+    thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
+
+    # Stream tokens back to user
+    response = ""
+    for new_text in streamer:
+        response += new_text
+        yield response
 
 # Initialize Gradio chat interface
 
@@ -106,4 +139,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()