Update app.py

app.py

@@ -23,14 +23,13 @@ def load_model(model_id):
     return generator
 
 def format_prompt(item):
-    # 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."
+    # Simplified prompt: rely on max_new_tokens=1 and model's understanding for single-letter answer
     prompt = f"""{item['question']}
 A. {item['choices'][0]}
 B. {item['choices'][1]}
 C. {item['choices'][2]}
 D. {item['choices'][3]}
-Answer: {system_instruction}""" # Place instruction after 'Answer:' with a space
+Answer:""" # Removed direct instruction from here
     return prompt, item['answer']
 
 def extract_choice_letter(output):
@@ -38,6 +37,10 @@ def extract_choice_letter(output):
     match = re.search(r"\b([ABCD])\b", output.strip())
     return match.group(1) if match else None
 
+def get_choice_letter(index):
+    """Converts a numerical choice index (0-3) to a capital letter (A-D)."""
+    return chr(ord('A') + index)
+
 def evaluate(model_id, sample_count, config_name, progress=gr.Progress()):
     if config_name == "ALL":
         # Dynamically get all MMLU subjects
@@ -50,15 +53,19 @@ def evaluate(model_id, sample_count, config_name, progress=gr.Progress()):
         for i, subject in enumerate(progress.tqdm(subjects, desc="Evaluating subjects")):
             dataset = load_dataset("cais/mmlu", subject, token=HF_TOKEN)["test"]
             dataset = dataset.shuffle(seed=42).select(range(min(sample_count, len(dataset))))
-            correct = 0
+            correct_subject = 0
             for j, item in enumerate(progress.tqdm(dataset, desc=f"Processing {subject} samples")):
-                prompt, answer = format_prompt(item)
+                prompt, answer_idx = format_prompt(item) # answer_idx is 0, 1, 2, or 3
+                expected_letter = get_choice_letter(answer_idx) # Convert to 'A', 'B', 'C', 'D'
+
                 # 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))
-            total_correct += correct
+                output_letter = extract_choice_letter(output) # Extract the letter from model's output
+                
+                is_correct = output_letter == expected_letter
+                correct_subject += is_correct
+                all_results.append((prompt, output.strip(), expected_letter, output_letter, is_correct)) # Store expected_letter
+            total_correct += correct_subject
             total_samples += len(dataset)
         avg_accuracy = total_correct / total_samples * 100
         return avg_accuracy, all_results
@@ -71,13 +78,16 @@ def evaluate(model_id, sample_count, config_name, progress=gr.Progress()):
     results = []
 
     for i, item in enumerate(progress.tqdm(dataset, desc=f"Processing {config_name} samples")):
-        prompt, answer = format_prompt(item)
+        prompt, answer_idx = format_prompt(item) # answer_idx is 0, 1, 2, or 3
+        expected_letter = get_choice_letter(answer_idx) # Convert to 'A', 'B', 'C', 'D'
+
         # 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
+        output_letter = extract_choice_letter(output) # Extract the letter from model's output
+        
+        is_correct = output_letter == expected_letter
         correct += is_correct
-        results.append((prompt, output.strip(), answer, output_letter, is_correct))
+        results.append((prompt, output.strip(), expected_letter, output_letter, is_correct)) # Store expected_letter
 
     accuracy = correct / len(dataset) * 100
     return accuracy, results