Update app.py

app.py

@@ -6,11 +6,10 @@ import random
 import re
 
 # Load model and tokenizer
-# model_name = "rgb2gbr/GRPO_BioMedmcqa_Qwen2.5-0.5B"
 model_name = "rgb2gbr/BioXP-0.5B-MedMCQA"
 
 SYSTEM_PROMPT = """
-You're a medical expert. Answer the question with careful analysis and explain why the selected option is correct in 150 words without reapeating.
+You are a medical expert. Answer the medical question with careful analysis and explain why the selected option is correct in 200 words without repeating.
 Respond in the following format:
 <answer>
 [correct answer]
@@ -45,15 +44,32 @@ def get_random_question():
         question_data.get('exp', None)   # Explanation
     )
 
-def predict(question: str, option_a: str, option_b: str, option_c: str, option_d: str, 
+def predict(question: str, option_a: str = "", option_b: str = "", option_c: str = "", option_d: str = "", 
            correct_option: int = None, explanation: str = None,
            temperature: float = 0.6, top_p: float = 0.9, max_tokens: int = 256):
-    # Format the question with options
-    formatted_question = f"Question: {question}\n\nOptions:\nA. {option_a}\nB. {option_b}\nC. {option_c}\nD. {option_d}"
+    
+    # Determine if this is an MCQ by checking if any option is provided
+    # Only treat as MCQ if at least one option is non-empty
+    is_mcq = any(opt.strip() for opt in [option_a, option_b, option_c, option_d])
+    
+    if is_mcq:
+        # Format MCQ question with only non-empty options
+        options = []
+        if option_a.strip(): options.append(f"A. {option_a}")
+        if option_b.strip(): options.append(f"B. {option_b}")
+        if option_c.strip(): options.append(f"C. {option_c}")
+        if option_d.strip(): options.append(f"D. {option_d}")
+        
+        formatted_question = f"Question: {question}\n\nOptions:\n" + "\n".join(options)
+        system_prompt = SYSTEM_PROMPT
+    else:
+        # Format regular question
+        formatted_question = f"Question: {question}"
+        system_prompt = SYSTEM_PROMPT
     
     # Create chat-style prompt
     prompt = [
-        {'role': 'system', 'content': SYSTEM_PROMPT},
+        {'role': 'system', 'content': system_prompt},
         {'role': 'user', 'content': formatted_question}
     ]
     
@@ -69,20 +85,23 @@ def predict(question: str, option_a: str, option_b: str, option_c: str, option_d
             max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
-            # repetition_penalty=1.1,
         )
     
-    # Get only the generated response (excluding the prompt)
+    # Get only the generated response
     generated_ids = generated_ids[0, model_inputs.input_ids.shape[1]:]
     model_response = tokenizer.decode(generated_ids, skip_special_tokens=True)
     
-    # Format output with evaluation if available
-    output = model_response
+    # Clean up the response by removing tags and formatting
+    cleaned_response = model_response
+    cleaned_response = re.sub(r'<answer>\s*([A-D])\s*</answer>', r'Answer: \1', cleaned_response, flags=re.IGNORECASE)
+    cleaned_response = re.sub(r'<reasoning>\s*(.*?)\s*</reasoning>', r'Reasoning:\n\1', cleaned_response, flags=re.IGNORECASE | re.DOTALL)
     
-    if correct_option is not None:
-        correct_letter = chr(65 + correct_option)  # Convert 0-3 to A-D
-        # Extract answer from model response for evaluation
-        answer_match = re.search(r"<answer>\s*([A-D])\s*</answer>", model_response, re.IGNORECASE)
+    # Format output with evaluation if available (only for MCQs)
+    output = cleaned_response
+    
+    if is_mcq and correct_option is not None:
+        correct_letter = chr(65 + correct_option)
+        answer_match = re.search(r"Answer:\s*([A-D])", cleaned_response, re.IGNORECASE)
         model_answer = answer_match.group(1).upper() if answer_match else "Not found"
         
         is_correct = model_answer == correct_letter
@@ -95,61 +114,106 @@ def predict(question: str, option_a: str, option_b: str, option_c: str, option_d
     
     return output
 
-# Create Gradio interface with Blocks for more control
-with gr.Blocks(title="Medical-QA (MedMCQA) Predictor") as demo:
-    gr.Markdown("# Medical-QA (MedMCQA) Predictor")
-    gr.Markdown("Get a random medical question or enter your own question and options.")
+# Create Gradio interface with mobile-optimized design
+with gr.Blocks(
+    title="BioXP Medical MCQ Assistant",
+    theme=gr.themes.Soft(
+        primary_hue="blue",
+        secondary_hue="blue",
+        neutral_hue="slate",
+        radius_size="md",
+        font=["Inter", "ui-sans-serif", "system-ui", "sans-serif"],
+    )
+) as demo:
+    gr.Markdown("""
+    # BioXP Medical MCQ Assistant
+    A specialized AI assistant for medical multiple-choice questions.
+    """)
     
     with gr.Row():
-        with gr.Column():
-            # Input fields
-            question = gr.Textbox(label="Question", lines=3, interactive=True)
+        with gr.Column(scale=1):
+            # Input fields with mobile-friendly spacing
+            question = gr.Textbox(
+                label="Medical Question",
+                placeholder="Enter your medical question here...",
+                lines=3,
+                interactive=True,
+                elem_classes=["mobile-input"]
+            )
             
-            # Options in an expandable accordion
-            with gr.Accordion("Options", open=False):
-                option_a = gr.Textbox(label="Option A", interactive=True)
-                option_b = gr.Textbox(label="Option B", interactive=True)
-                option_c = gr.Textbox(label="Option C", interactive=True)
-                option_d = gr.Textbox(label="Option D", interactive=True)
-            
-            # Generation parameters
-            with gr.Accordion("Generation Parameters", open=False):
-                temperature = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.6,
-                    step=0.1,
-                    label="Temperature",
-                    info="Higher values make output more random, lower values more focused"
+            # Options in a mobile-friendly accordion
+            with gr.Accordion("Options", open=True):
+                option_a = gr.Textbox(
+                    label="Option A",
+                    placeholder="Enter option A...",
+                    interactive=True,
+                    elem_classes=["mobile-input"]
+                )
+                option_b = gr.Textbox(
+                    label="Option B",
+                    placeholder="Enter option B...",
+                    interactive=True,
+                    elem_classes=["mobile-input"]
                 )
-                top_p = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.9,
-                    step=0.1,
-                    label="Top P",
-                    info="Higher values allow more diverse tokens, lower values more focused"
+                option_c = gr.Textbox(
+                    label="Option C",
+                    placeholder="Enter option C...",
+                    interactive=True,
+                    elem_classes=["mobile-input"]
                 )
+                option_d = gr.Textbox(
+                    label="Option D",
+                    placeholder="Enter option D...",
+                    interactive=True,
+                    elem_classes=["mobile-input"]
+                )
+            
+            # Generation parameters in a collapsible section
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        temperature = gr.Slider(
+                            minimum=0.1,
+                            maximum=1.0,
+                            value=0.6,
+                            step=0.1,
+                            label="Temperature",
+                            info="Higher = more creative, Lower = more focused"
+                        )
+                    with gr.Column(scale=1):
+                        top_p = gr.Slider(
+                            minimum=0.1,
+                            maximum=1.0,
+                            value=0.9,
+                            step=0.1,
+                            label="Top P",
+                            info="Controls response diversity"
+                        )
                 max_tokens = gr.Slider(
                     minimum=50,
                     maximum=512,
                     value=256,
                     step=32,
-                    label="Max Tokens",
-                    info="Maximum length of the generated response (recommended: 256)"
+                    label="Max Response Length",
+                    info="Maximum length of the response"
                 )
             
-            # Hidden fields for correct answer and explanation
+            # Hidden fields
             correct_option = gr.Number(visible=False)
             expert_explanation = gr.Textbox(visible=False)
             
-            # Buttons
+            # Buttons with mobile-friendly spacing
             with gr.Row():
-                predict_btn = gr.Button("Predict", variant="primary")
-                random_btn = gr.Button("Get Random Question", variant="secondary")
-            
-            # Output
-            output = gr.Textbox(label="Model's Answer", lines=10)
+                predict_btn = gr.Button("Get Answer", variant="primary", size="lg", elem_classes=["mobile-button"])
+                random_btn = gr.Button("Random Question", variant="secondary", size="lg", elem_classes=["mobile-button"])
+        
+        with gr.Column(scale=1):
+            # Output with mobile-friendly styling
+            output = gr.Textbox(
+                label="Model's Response",
+                lines=12,
+                elem_classes=["response-box", "mobile-output"]
+            )
     
     # Set up button actions
     predict_btn.click(
@@ -168,6 +232,76 @@ with gr.Blocks(title="Medical-QA (MedMCQA) Predictor") as demo:
         outputs=[question, option_a, option_b, option_c, option_d, correct_option, expert_explanation]
     )
 
+    # Add mobile-optimized CSS
+    gr.HTML("""
+    <style>
+        /* Mobile-friendly base styles */
+        .container {
+            max-width: 100%;
+            padding: 0.5rem;
+        }
+        
+        /* Input styling */
+        .mobile-input textarea {
+            font-size: 1rem;
+            padding: 0.75rem;
+            border-radius: 0.5rem;
+            min-height: 2.5rem;
+        }
+        
+        /* Button styling */
+        .mobile-button {
+            width: 100%;
+            margin: 0.5rem 0;
+            padding: 0.75rem;
+            font-size: 1rem;
+            font-weight: 500;
+        }
+        
+        /* Response box styling */
+        .response-box {
+            font-family: 'Inter', sans-serif;
+            line-height: 1.6;
+        }
+        .response-box textarea {
+            font-size: 1rem;
+            padding: 1rem;
+            border-radius: 0.5rem;
+        }
+        
+        /* Mobile-specific adjustments */
+        @media (max-width: 768px) {
+            .gr-form {
+                padding: 0.75rem;
+            }
+            .gr-box {
+                margin: 0.5rem 0;
+            }
+            .gr-button {
+                min-height: 2.5rem;
+            }
+            .gr-accordion {
+                margin: 0.5rem 0;
+            }
+            .gr-input {
+                margin-bottom: 0.5rem;
+            }
+        }
+        
+        /* Dark mode support */
+        @media (prefers-color-scheme: dark) {
+            .gr-box {
+                background-color: #1a1a1a;
+            }
+            .mobile-input textarea,
+            .response-box textarea {
+                background-color: #2a2a2a;
+                color: #ffffff;
+            }
+        }
+    </style>
+    """)
+
 # Launch the app
 if __name__ == "__main__":
-    demo.launch() 
+    demo.launch(share=False)