Update app.py

app.py

@@ -32,21 +32,6 @@ 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:
@@ -55,33 +40,18 @@ def respond(
             messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    # Prepare prompt and tokenize
-    prompt = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    # Set up streamer for real-time output
-    streamer = TextIteratorStreamer(
-        tokenizer, skip_prompt=True, skip_special_tokens=True
-    )
-    generate_kwargs = dict(
+    outputs = model.generate(
         **inputs,
-        streamer=streamer,
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-        do_sample=True,
+        do_sample=True
     )
-    # 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
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    yield response
 
 # Initialize Gradio chat interface