Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from datasets import load_dataset
+from datasets import load_dataset, get_dataset_config_names # Import get_dataset_config_names
 import torch
 import re
 import json
@@ -23,37 +23,26 @@ def load_model(model_id):
     return generator
 
 def format_prompt(item):
-    system_instruction = " Only answer with a single letter: A, B, C, or D."
+    # Emphasize the single letter answer instruction to encourage concise output
+    system_instruction = "Respond ONLY with a single capital letter: A, B, C, or D. No other text."
     prompt = f"""{item['question']}
 A. {item['choices'][0]}
 B. {item['choices'][1]}
 C. {item['choices'][2]}
 D. {item['choices'][3]}
-Answer:{system_instruction}"""
+Answer: {system_instruction}""" # Place instruction after 'Answer:' with a space
     return prompt, item['answer']
 
 def extract_choice_letter(output):
+    # This function should now be more reliable as max_new_tokens is set to 1
     match = re.search(r"\b([ABCD])\b", output.strip())
     return match.group(1) if match else None
 
-# Modified evaluate function to return accuracy as a float directly
 def evaluate(model_id, sample_count, config_name, progress=gr.Progress()):
     if config_name == "ALL":
-        subjects = [
-            "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"
-        ]
+        # Dynamically get all MMLU subjects
+        subjects = get_dataset_config_names("cais/mmlu", token=HF_TOKEN) 
+        
         gen = load_model(model_id)
         total_correct = 0
         total_samples = 0
@@ -64,16 +53,16 @@ def evaluate(model_id, sample_count, config_name, progress=gr.Progress()):
             correct = 0
             for j, item in enumerate(progress.tqdm(dataset, desc=f"Processing {subject} samples")):
                 prompt, answer = format_prompt(item)
-                output = gen(prompt, max_new_tokens=20, do_sample=False)[0]["generated_text"]
+                # Crucial change: Limit generation to 1 new token
+                output = gen(prompt, max_new_tokens=1, do_sample=False)[0]["generated_text"]
                 output_letter = extract_choice_letter(output)
                 correct += output_letter == answer
                 all_results.append((prompt, output.strip(), answer, output_letter, output_letter == answer))
-            # No need to write subject-level record here, only aggregate
             total_correct += correct
             total_samples += len(dataset)
         avg_accuracy = total_correct / total_samples * 100
-        # Return the float accuracy value
         return avg_accuracy, all_results
+    
     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))))
@@ -83,19 +72,17 @@ def evaluate(model_id, sample_count, config_name, progress=gr.Progress()):
 
     for i, item in enumerate(progress.tqdm(dataset, desc=f"Processing {config_name} samples")):
         prompt, answer = format_prompt(item)
-        output = gen(prompt, max_new_tokens=20, do_sample=False)[0]["generated_text"]
+        # Crucial change: Limit generation to 1 new token
+        output = gen(prompt, max_new_tokens=1, 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))
 
     accuracy = correct / len(dataset) * 100
-    # Return the float accuracy value
     return accuracy, results
 
-# Pass progress to evaluate function
 def run(model_id, sample_count, config_name, progress=gr.Progress()):
-    # Receive accuracy_value directly as a float
     accuracy_value, details = evaluate(model_id, sample_count, config_name, progress)
 
     formatted = "\n\n".join([
@@ -103,17 +90,15 @@ def run(model_id, sample_count, config_name, progress=gr.Progress()):
         for q, o, a, g, c in details
     ])
     
-    # Format the score string based on config_name
     if config_name == "ALL":
         score_string = f"Average Accuracy: {accuracy_value:.2f}% across all subjects"
     else:
-        # Assuming len(details) corresponds to the number of samples processed for a single subject
         score_string = f"Accuracy: {accuracy_value:.2f}%, out of {len(details)} samples"
 
     record = {"model_id": model_id, "subject": config_name, "accuracy": accuracy_value}
     with open("eval.jsonl", "a") as f:
         f.write(json.dumps(record) + "\n")
-    return score_string, formatted # Return the formatted string and details
+    return score_string, formatted
 
 def save_text(text):
     return "evaluation_results.txt", text
@@ -128,14 +113,17 @@ with gr.Blocks(css="body {font-family: Inter, sans-serif; padding: 1em; max-widt
     Enter your model ID, pick MMLU, choose a subject, and hit evaluate.
     """)
 
+    # Get all MMLU subject config names dynamically
+    mmlu_subjects = ["ALL"] + get_dataset_config_names("cais/mmlu", token=HF_TOKEN)
+
     with gr.Row():
         model_id = gr.Textbox(label="Your Hugging Face Model ID", placeholder="e.g., your-org/your-model")
         config_name = gr.Dropdown(
-    label="Choose MMLU Subject",
-    choices=["ALL"],
-    value="ALL",
-    interactive=False
-)
+            label="Choose MMLU Subject",
+            choices=mmlu_subjects, # Populate with all subjects
+            value="ALL",
+            interactive=True # Make interactive now that there are more choices
+        )
         sample_count = gr.Slider(label="Number of Samples", minimum=1, maximum=100, value=10, step=1)
 
     run_button = gr.Button("🚀 Run Evaluation")
@@ -143,7 +131,6 @@ 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")
 
-    # Pass progress to the run function
     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())
 
@@ -172,4 +159,4 @@ with gr.Blocks(css="body {font-family: Inter, sans-serif; padding: 1em; max-widt
 
     demo.load(load_leaderboard, inputs=[], outputs=[leaderboard_plot, leaderboard_table])
 
-demo.launch()
+demo.launch()