Update app.py

app.py

@@ -18,38 +18,20 @@ def load_model(model_id):
     model_cache[model_id] = generator
     return generator
 
-def format_prompt(item, source):
-    if source == "cais/mmlu":
-        prompt = f"{item['question']}\nA. {item['choices'][0]}\nB. {item['choices'][1]}\nC. {item['choices'][2]}\nD. {item['choices'][3]}\nAnswer:"
-        answer = item['answer']
-    elif source == "TIGER-Lab/MMLU-Pro":
-        if all(opt in item for opt in ['A', 'B', 'C', 'D']):
-            prompt = f"{item['question']}\nA. {item['A']}\nB. {item['B']}\nC. {item['C']}\nD. {item['D']}\nAnswer:"
-        else:
-            choices = item.get("choices", ["", "", "", ""])
-            prompt = f"{item['question']}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nAnswer:"
-        answer = item['answer']
-    elif source == "cais/hle":
-        prompt = f"{item['question']}\n{item['A']}\n{item['B']}\n{item['C']}\n{item['D']}\nAnswer:"
-        answer = item['answer']
-    else:
-        prompt, answer = "", ""
-    return prompt, answer
-
-def evaluate(model_id, dataset_name, sample_count):
-    gen = load_model(model_id)
-    dataset = load_dataset(dataset_name, token=HF_TOKEN)
-    if 'test' in dataset:
-        dataset = dataset['test']
-    else:
-        dataset = dataset[list(dataset.keys())[0]]
+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:"
+    return prompt, item['answer']
 
+def evaluate(model_id, sample_count, config_name):
+    gen = load_model(model_id)
+    dataset = load_dataset("cais/mmlu", config_name, token=HF_TOKEN)["test"]
     dataset = dataset.shuffle(seed=42).select(range(min(sample_count, len(dataset))))
+
     correct = 0
     results = []
 
     for item in dataset:
-        prompt, answer = format_prompt(item, dataset_name)
+        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 output[::-1] if char in "ABCD"), None)
         is_correct = output_letter == answer
@@ -57,19 +39,17 @@ def evaluate(model_id, dataset_name, sample_count):
         results.append((prompt, output.strip(), answer, output_letter, is_correct))
 
     accuracy = correct / len(dataset) * 100
-    return accuracy, results
-
-def run(model_id, benchmark, sample_count):
-    if benchmark != "cais/mmlu":
-        return "Only MMLU (cais/mmlu) is available now. MMLU-Pro and Humanity's Last Exam are coming soon.", ""
+    return f"Accuracy: {accuracy:.2f}%", results
 
-    accuracy, details = evaluate(model_id, benchmark, sample_count)
+def run(model_id, sample_count, config_name):
+    if config_name == "coming soon":
+        return "Only MMLU is currently available. MMLU-Pro and HLE coming soon.", ""
+    score, details = evaluate(model_id, sample_count, config_name)
     formatted = "\n\n".join([
         f"### Question:\n{q}\n\n**Model Answer:** {o}\n**Expected:** {a}\n**Predicted:** {g}\n**Correct:** {c}"
         for q, o, a, g, c in details
     ])
-
-    return f"Accuracy: {accuracy:.2f}%", formatted
+    return score, formatted
 
 def save_text(text):
     return "evaluation_results.txt", text
@@ -81,15 +61,29 @@ with gr.Blocks(css="body {font-family: Inter, sans-serif; padding: 1em; max-widt
     Currently, only **MMLU** (`cais/mmlu`) is available for evaluation.
     **MMLU-Pro** and **Humanity's Last Exam** will be coming soon.
 
-    Enter your model ID, pick MMLU, and hit evaluate.
+    Enter your model ID, pick MMLU, choose a subject, and hit evaluate.
     """)
 
     with gr.Row():
         model_id = gr.Textbox(label="Your Hugging Face Model ID", placeholder="e.g., your-org/your-model")
-        benchmark = gr.Dropdown(
-            label="Choose Benchmark",
-            choices=["cais/mmlu"],
-            value="cais/mmlu"
+        config_name = gr.Dropdown(
+            label="Choose MMLU Subject",
+            choices=[
+                "abstract_algebra", "anatomy", "astronomy", "business_ethics", "college_biology",
+                "college_chemistry", "college_computer_science", "college_mathematics", "college_medicine",
+                "college_physics", "computer_security", "econometrics", "electrical_engineering",
+                "elementary_mathematics", "formal_logic", "global_facts", "high_school_biology",
+                "high_school_chemistry", "high_school_computer_science", "high_school_european_history",
+                "high_school_geography", "high_school_government_and_politics", "high_school_macroeconomics",
+                "high_school_microeconomics", "high_school_physics", "high_school_psychology",
+                "high_school_statistics", "high_school_us_history", "high_school_world_history", "human_aging",
+                "human_sexuality", "international_law", "jurisprudence", "logical_fallacies", "machine_learning",
+                "management", "marketing", "medical_genetics", "miscellaneous", "moral_disputes",
+                "moral_scenarios", "nutrition", "philosophy", "prehistory", "professional_accounting",
+                "professional_law", "professional_medicine", "professional_psychology", "public_relations",
+                "security_studies", "sociology", "us_foreign_policy", "virology", "world_religions"
+            ],
+            value="college_mathematics"
         )
         sample_count = gr.Slider(label="Number of Samples", minimum=1, maximum=100, value=10, step=1)
 
@@ -98,7 +92,11 @@ with gr.Blocks(css="body {font-family: Inter, sans-serif; padding: 1em; max-widt
     detail_output = gr.Textbox(label="Evaluation Details", lines=20, interactive=False)
     download_button = gr.Button("📥 Download Full Evaluation")
 
-    run_button.click(run, inputs=[model_id, benchmark, sample_count], outputs=[acc_output, detail_output])
+    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(share=True)
+    gr.Markdown("""
+    MMLU-Pro and HLE support will be added soon.
+    """)
+
+demo.launch()