Update app.py

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from datasets import load_dataset
 import torch
+import re
 
 # Cache to avoid reloading the model
 model_cache = {}
@@ -19,9 +20,20 @@ def load_model(model_id):
     return generator
 
 def format_prompt(item):
-    prompt = f"{item['question']}\nA. {item['choices'][0]}\nB. {item['choices'][1]}\nC. {item['choices'][2]}\nD. {item['choices'][3]}\nAnswer:"
+    system_instruction = "
+Only answer with a single letter: A, B, C, or D."
+    prompt = f"{item['question']}
+A. {item['choices'][0]}
+B. {item['choices'][1]}
+C. {item['choices'][2]}
+D. {item['choices'][3]}
+Answer:{system_instruction}"
     return prompt, item['answer']
 
+def extract_choice_letter(output):
+    match = re.search(r"\b([ABCD])\b", output.strip())
+    return match.group(1) if match else None
+
 def evaluate(model_id, sample_count, config_name):
     gen = load_model(model_id)
     dataset = load_dataset("cais/mmlu", config_name, token=HF_TOKEN)["test"]
@@ -32,8 +44,8 @@ def evaluate(model_id, sample_count, config_name):
 
     for item in dataset:
         prompt, answer = format_prompt(item)
-        output = gen(prompt, max_new_tokens=10, do_sample=False)[0]["generated_text"]
-        output_letter = next((char for char in reversed(output) if char in "ABCD"), None)
+        output = gen(prompt, max_new_tokens=20, do_sample=False)[0]["generated_text"]
+        output_letter = extract_choice_letter(output)
         is_correct = output_letter == answer
         correct += is_correct
         results.append((prompt, output.strip(), answer, output_letter, is_correct))
@@ -93,4 +105,4 @@ with gr.Blocks(css="body {font-family: Inter, sans-serif; padding: 1em; max-widt
     run_button.click(run, inputs=[model_id, sample_count, config_name], outputs=[acc_output, detail_output])
     download_button.click(save_text, inputs=detail_output, outputs=gr.File())
 
-demo.launch()
+demo.launch()