Update app.py

app.py

@@ -3,9 +3,10 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from datasets import load_dataset
 import random
+import re
 
 # Load model and tokenizer
-model_name = "rgb2gbr/BioXP-0.5B-MedMCQA"
+model_name = "rgb2gbr/GRPO_BioMedmcqa_Qwen2.5-0.5B"
 model = AutoModelForCausalLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -26,10 +27,27 @@ def get_random_question():
         question_data['opa'],
         question_data['opb'],
         question_data['opc'],
-        question_data['opd']
+        question_data['opd'],
+        question_data.get('cop', None),  # Correct option (0-3)
+        question_data.get('exp', None)   # Explanation
     )
 
-def predict(question: str, option_a: str, option_b: str, option_c: str, option_d: str):
+def extract_answer(prediction: str) -> tuple:
+    """Extract answer and reasoning from model output"""
+    # Try to find the answer part
+    answer_match = re.search(r"Answer:\s*([A-D])", prediction, re.IGNORECASE)
+    answer = answer_match.group(1).upper() if answer_match else "Not found"
+    
+    # Try to find reasoning part
+    reasoning = ""
+    if "Reasoning:" in prediction:
+        reasoning = prediction.split("Reasoning:")[-1].strip()
+    elif "Explanation:" in prediction:
+        reasoning = prediction.split("Explanation:")[-1].strip()
+    
+    return answer, reasoning
+
+def predict(question: str, option_a: str, option_b: str, option_c: str, option_d: str, correct_option: int = None, explanation: str = None):
     # Format the prompt
     prompt = f"Question: {question}\n\nOptions:\nA. {option_a}\nB. {option_b}\nC. {option_c}\nD. {option_d}\n\nAnswer:"
     
@@ -40,15 +58,30 @@ def predict(question: str, option_a: str, option_b: str, option_c: str, option_d
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=10,
-            temperature=0.7,
-            do_sample=False,
+            max_new_tokens=256,
+            temperature=0.6,
+            top_p=0.9,
+            do_sample=True,
             pad_token_id=tokenizer.eos_token_id
         )
     
     # Get prediction
     prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return prediction
+    model_answer, model_reasoning = extract_answer(prediction)
+    
+    # Format output with evaluation if available
+    output = prediction
+    if correct_option is not None:
+        correct_letter = chr(65 + correct_option)  # Convert 0-3 to A-D
+        is_correct = model_answer == correct_letter
+        output += f"\n\n---\nEvaluation:\n"
+        output += f"Correct Answer: {correct_letter}\n"
+        output += f"Model's Answer: {model_answer}\n"
+        output += f"Result: {'✅ Correct' if is_correct else '❌ Incorrect'}\n"
+        if explanation:
+            output += f"\nExpert Explanation:\n{explanation}"
+    
+    return output
 
 # Create Gradio interface with Blocks for more control
 with gr.Blocks(title="Medical MCQ Predictor") as demo:
@@ -64,25 +97,29 @@ with gr.Blocks(title="Medical MCQ Predictor") as demo:
             option_c = gr.Textbox(label="Option C", interactive=True)
             option_d = gr.Textbox(label="Option D", interactive=True)
             
+            # Hidden fields for correct answer and explanation
+            correct_option = gr.Number(visible=False)
+            expert_explanation = gr.Textbox(visible=False)
+            
             # Buttons
             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=5)
+            output = gr.Textbox(label="Model's Answer", lines=10)
     
     # Set up button actions
     predict_btn.click(
         fn=predict,
-        inputs=[question, option_a, option_b, option_c, option_d],
+        inputs=[question, option_a, option_b, option_c, option_d, correct_option, expert_explanation],
         outputs=output
     )
     
     random_btn.click(
         fn=get_random_question,
         inputs=[],
-        outputs=[question, option_a, option_b, option_c, option_d]
+        outputs=[question, option_a, option_b, option_c, option_d, correct_option, expert_explanation]
     )
 
 # Launch the app