Update app.py

app.py

@@ -7,8 +7,7 @@ import sys
 # Install specific transformers version
 subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers==4.48.3"])
 
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load model and tokenizer
 model_name = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
@@ -28,7 +27,7 @@ def load_model():
 
 @spaces.GPU(duration=120)
 def generate_response(message, history, enable_reasoning, temperature, top_p, max_tokens):
-    """Generate response from the model with streaming"""
+    """Generate response from the model"""
     
     # Prepare messages with reasoning control
     messages = []
@@ -59,47 +58,33 @@ def generate_response(message, history, enable_reasoning, temperature, top_p, ma
         return_tensors="pt"
     ).to(model.device)
     
-    # Create streamer for real-time token generation
-    streamer = TextIteratorStreamer(
-        tokenizer, 
-        skip_prompt=True, 
-        skip_special_tokens=True,
-        timeout=10.0
-    )
-    
     # Set generation parameters based on reasoning mode
     if enable_reasoning:
         # Recommended settings for reasoning
         generation_kwargs = {
-            "input_ids": tokenized_chat,
             "temperature": temperature if temperature > 0 else 0.6,
             "top_p": top_p if top_p < 1 else 0.95,
             "do_sample": True,
             "max_new_tokens": max_tokens,
-            "eos_token_id": tokenizer.eos_token_id,
-            "streamer": streamer
+            "eos_token_id": tokenizer.eos_token_id
         }
     else:
         # Greedy search for non-reasoning
         generation_kwargs = {
-            "input_ids": tokenized_chat,
             "do_sample": False,
             "max_new_tokens": max_tokens,
-            "eos_token_id": tokenizer.eos_token_id,
-            "streamer": streamer
+            "eos_token_id": tokenizer.eos_token_id
         }
     
-    # Generate response in a separate thread
-    generation_thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    generation_thread.start()
+    # Generate response
+    with torch.no_grad():
+        outputs = model.generate(tokenized_chat, **generation_kwargs)
     
-    # Stream the response
-    response = ""
-    for new_text in streamer:
-        response += new_text
-        yield response
+    # Decode and extract the assistant's response
+    generated_tokens = outputs[0][tokenized_chat.shape[-1]:]  # Get only new tokens
+    response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
     
-    generation_thread.join()
+    return response
 
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -126,7 +111,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Row():
         submit = gr.Button("Send", variant="primary")
         clear = gr.Button("Clear")
-        stop = gr.Button("Stop")
     
     with gr.Accordion("Advanced Settings", open=False):
         enable_reasoning = gr.Checkbox(
@@ -163,36 +147,29 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         )
     
     def user_submit(message, history):
-        """Add user message to chat history"""
         return "", history + [[message, None]]
     
     def bot_response(history, enable_reasoning, temperature, top_p, max_tokens):
-        """Generate bot response with streaming"""
-        if not history or history[-1][0] is None:
-            yield history
-            return
+        if not history:
+            return history
         
         message = history[-1][0]
-        history[-1][1] = ""
-        
         try:
-            # Stream the response
-            for partial_response in generate_response(
+            response = generate_response(
                 message, 
                 history[:-1], 
                 enable_reasoning, 
                 temperature, 
                 top_p, 
                 max_tokens
-            ):
-                history[-1][1] = partial_response
-                yield history
+            )
+            history[-1][1] = response
         except Exception as e:
             history[-1][1] = f"Error generating response: {str(e)}"
-            yield history
+        
+        return history
     
-    # Handle message submission
-    submit_event = msg.submit(
+    msg.submit(
         user_submit, 
         [msg, chatbot], 
         [msg, chatbot], 
@@ -203,8 +180,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         chatbot
     )
     
-    # Handle button click
-    click_event = submit.click(
+    submit.click(
         user_submit, 
         [msg, chatbot], 
         [msg, chatbot], 
@@ -215,15 +191,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         chatbot
     )
     
-    # Stop generation
-    stop.click(
-        None, 
-        None, 
-        None, 
-        cancels=[submit_event, click_event]
-    )
-    
-    # Clear chat
     clear.click(lambda: None, None, chatbot, queue=False)
     
     # Example prompts
@@ -239,4 +206,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch()