Update app.py

app.py

@@ -3,10 +3,9 @@ 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"
+model_name = "rgb2gbr/BioXP-0.5B-MedMCQA"
 model = AutoModelForCausalLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -27,38 +26,12 @@ def get_random_question():
         question_data['opa'],
         question_data['opb'],
         question_data['opc'],
-        question_data['opd'],
-        question_data.get('cop', None),  # Correct option (0-3)
-        question_data.get('exp', None)   # Explanation
+        question_data['opd']
     )
 
-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):
+def predict(question: str, option_a: str, option_b: str, option_c: str, option_d: str):
     # Format the prompt
-    prompt = f"""Question: {question}
-
-Options:
-A. {option_a}
-B. {option_b}
-C. {option_c}
-D. {option_d}
-
-Please provide your answer and reasoning.
-Answer:"""
+    prompt = f"Question: {question}\n\nOptions:\nA. {option_a}\nB. {option_b}\nC. {option_c}\nD. {option_d}\n\nAnswer:"
     
     # Tokenize and generate
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
@@ -67,47 +40,15 @@ Answer:"""
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=256,
-            temperature=0.6,
-            top_p=0.9,
-            do_sample=True,
+            max_new_tokens=10,
+            temperature=0.7,
+            do_sample=False,
             pad_token_id=tokenizer.eos_token_id
         )
     
     # Get prediction
     prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    model_answer, model_reasoning = extract_answer(prediction)
-    
-    # Format the output
-    output = f"""## Raw Model Output
-```
-{prediction}
-```
-
-## Evaluation
-
-### Answer
-{model_answer}
-
-### Reasoning
-{model_reasoning if model_reasoning else "No reasoning provided"}
-"""
-    
-    # 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"""
-### Results
-- Correct Answer: {correct_letter}
-- Model's Answer: {model_answer}
-- Result: {'✅ Correct' if is_correct else '❌ Incorrect'}
-
-### Expert Explanation
-{explanation if explanation else "No expert explanation available"}
-"""
-    
-    return output
+    return prediction
 
 # Create Gradio interface with Blocks for more control
 with gr.Blocks(title="Medical MCQ Predictor") as demo:
@@ -123,29 +64,25 @@ 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 with markdown support
-            output = gr.Markdown(label="Model's Answer")
+            # Output
+            output = gr.Textbox(label="Model's Answer", lines=5)
     
     # Set up button actions
     predict_btn.click(
         fn=predict,
-        inputs=[question, option_a, option_b, option_c, option_d, correct_option, expert_explanation],
+        inputs=[question, option_a, option_b, option_c, option_d],
         outputs=output
     )
     
     random_btn.click(
         fn=get_random_question,
         inputs=[],
-        outputs=[question, option_a, option_b, option_c, option_d, correct_option, expert_explanation]
+        outputs=[question, option_a, option_b, option_c, option_d]
     )
 
 # Launch the app