Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from datasets import load_dataset
 import random
+import re
 
 # Load model and tokenizer
 model_name = "rgb2gbr/GRPO_BioMedmcqa_Qwen2.5-0.5B"
@@ -26,12 +27,38 @@ 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:"
+    prompt = f"""Question: {question}
+
+Options:
+A. {option_a}
+B. {option_b}
+C. {option_c}
+D. {option_d}
+
+Please provide your answer and reasoning.
+Answer:"""
     
     # Tokenize and generate
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
@@ -49,7 +76,32 @@ def predict(question: str, option_a: str, option_b: str, option_c: str, option_d
     
     # Get prediction
     prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return prediction
+    model_answer, model_reasoning = extract_answer(prediction)
+    
+    # Format the output
+    output = f"""## Model's Response
+
+### Answer
+{model_answer}
+
+### Reasoning
+{model_reasoning if model_reasoning else "No reasoning provided"}
+
+### Evaluation
+"""
+    
+    # Add evaluation if correct answer is available
+    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"- 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"\n### Expert Explanation\n{explanation}"
+    
+    return output
 
 # Create Gradio interface with Blocks for more control
 with gr.Blocks(title="Medical MCQ Predictor") as demo:
@@ -65,25 +117,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 with markdown support
+            output = gr.Markdown(label="Model's Answer")
     
     # 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