Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import time
-import spaces  # ADDED THIS IMPORT
+import spaces
 
 # Model configuration
 MODEL_NAME = "krishna195/medgemma-anatomy-v1.2"
@@ -48,7 +48,7 @@ def load_model():
 print("Initializing MedGemma...")
 model, tokenizer = load_model()
 
-@spaces.GPU(duration=60)  # MOVED DECORATOR HERE - applied to inference function
+@spaces.GPU(duration=60)
 def generate_response(question, max_tokens=512, temperature=0.7, top_p=0.9):
     """
     Generate medical response for a given question
@@ -59,51 +59,59 @@ def generate_response(question, max_tokens=512, temperature=0.7, top_p=0.9):
         temperature: Sampling temperature (0.0-1.0)
         top_p: Nucleus sampling parameter
     """
-    if not question.strip():
-        return "Please enter a medical question."
-    
-    # Format prompt with Gemma chat template
-    prompt = f"""<start_of_turn>user
+    try:
+        if not question.strip():
+            return "⚠️ Please enter a medical question."
+        
+        # Show processing message
+        yield "🔄 **Processing your question...**\n\nGenerating response, please wait..."
+        
+        # Format prompt with Gemma chat template
+        prompt = f"""<start_of_turn>user
 {question}<end_of_turn>
 <start_of_turn>model
 """
-    
-    # Tokenize
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    # Generate
-    start_time = time.time()
-    
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            do_sample=True,
-            top_p=top_p,
-            repetition_penalty=1.1,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    
-    generation_time = time.time() - start_time
-    
-    # Decode response
-    full_output = tokenizer.decode(outputs[0], skip_special_tokens=False)
-    
-    # Extract model response
-    if "<start_of_turn>model" in full_output:
-        response = full_output.split("<start_of_turn>model")[-1]
-        response = response.split("<end_of_turn>")[0].strip()
-    else:
-        response = full_output.strip()
-    
-    # Add metadata
-    tokens_generated = outputs.shape[1] - inputs['input_ids'].shape[1]
-    tokens_per_sec = tokens_generated / generation_time if generation_time > 0 else 0
-    
-    metadata = f"\n\n---\n*Generated in {generation_time:.2f}s ({tokens_per_sec:.1f} tokens/sec)*"
-    
-    return response + metadata
+        
+        # Tokenize
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        # Generate
+        start_time = time.time()
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                do_sample=True,
+                top_p=top_p,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        generation_time = time.time() - start_time
+        
+        # Decode response
+        full_output = tokenizer.decode(outputs[0], skip_special_tokens=False)
+        
+        # Extract model response
+        if "<start_of_turn>model" in full_output:
+            response = full_output.split("<start_of_turn>model")[-1]
+            response = response.split("<end_of_turn>")[0].strip()
+        else:
+            response = full_output.strip()
+        
+        # Add metadata
+        tokens_generated = outputs.shape[1] - inputs['input_ids'].shape[1]
+        tokens_per_sec = tokens_generated / generation_time if generation_time > 0 else 0
+        
+        metadata = f"\n\n---\n✅ *Generated in {generation_time:.2f}s ({tokens_per_sec:.1f} tokens/sec) | Device: {DEVICE.upper()}*"
+        
+        yield response + metadata
+        
+    except Exception as e:
+        error_msg = f"❌ **Error occurred:**\n\n```\n{str(e)}\n```\n\nPlease try again or contact support if the issue persists."
+        yield error_msg
 
 # Example questions
 examples = [
@@ -120,6 +128,7 @@ css = """
 #warning {background-color: #FFCCCB; padding: 10px; border-radius: 5px; margin-bottom: 10px;}
 .generate-btn {background: linear-gradient(90deg, #667eea 0%, #764ba2 100%); color: white;}
 footer {visibility: hidden;}
+#output-box {min-height: 200px; border: 1px solid #e0e0e0; border-radius: 8px; padding: 15px;}
 """
 
 # Build Gradio interface
@@ -183,29 +192,36 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                     info="Nucleus sampling parameter"
                 )
             
-            generate_btn = gr.Button("Generate Response", variant="primary", elem_classes="generate-btn")
+            generate_btn = gr.Button("🚀 Generate Response", variant="primary", elem_classes="generate-btn")
+            clear_btn = gr.ClearButton([question_input], value="🗑️ Clear")
             
         with gr.Column(scale=3):
-            output = gr.Markdown(label="Response")
+            output = gr.Markdown(
+                label="Response", 
+                value="*Your medical answer will appear here...*",
+                elem_id="output-box"
+            )
     
     with gr.Row():
         gr.Examples(
             examples=examples,
             inputs=question_input,
-            label="Example Questions"
+            label="📋 Example Questions - Click to try"
         )
     
     # Event handlers
     generate_btn.click(
         fn=generate_response,
         inputs=[question_input, max_tokens, temperature, top_p],
-        outputs=output
+        outputs=output,
+        show_progress=True
     )
     
     question_input.submit(
         fn=generate_response,
         inputs=[question_input, max_tokens, temperature, top_p],
-        outputs=output
+        outputs=output,
+        show_progress=True
     )
     
     gr.Markdown(