Update app.py

app.py

@@ -47,7 +47,9 @@ def extract_answer(prediction: str) -> tuple:
     
     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, 
+           correct_option: int = None, explanation: str = None,
+           temperature: float = 0.6, top_p: float = 0.9, max_tokens: int = 256):
     # Format the prompt
     prompt = f"Question: {question}\n\nOptions:\nA. {option_a}\nB. {option_b}\nC. {option_c}\nD. {option_d}\n\nAnswer:"
     
@@ -58,9 +60,9 @@ 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=20,
-            temperature=0.6,
-            top_p=0.9,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id
         )
@@ -97,6 +99,33 @@ 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)
             
+            # 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"
+                )
+                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"
+                )
+                max_tokens = gr.Slider(
+                    minimum=32,
+                    maximum=512,
+                    value=256,
+                    step=32,
+                    label="Max Tokens",
+                    info="Maximum length of the generated response"
+                )
+            
             # Hidden fields for correct answer and explanation
             correct_option = gr.Number(visible=False)
             expert_explanation = gr.Textbox(visible=False)
@@ -112,7 +141,11 @@ with gr.Blocks(title="Medical MCQ Predictor") as demo:
     # 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,
+            correct_option, expert_explanation,
+            temperature, top_p, max_tokens
+        ],
         outputs=output
     )