add CMMLU evaluation results

eval/evaluate_cmmlu.py

@@ -0,0 +1,271 @@
+import os
+import pandas as pd
+import numpy as np
+import argparse
+import datasets
+import torch
+from collections import defaultdict
+
+from typing import List
+from tqdm import tqdm
+from transformers.trainer_utils import set_seed
+
+
+'''
+wget https://huggingface.co/datasets/haonan-li/cmmlu/resolve/main/cmmlu_v1_0_1.zip
+mkdir data/cmmlu
+mv cmmlu_v1_0_1.zip data/cmmlu
+cd data/cmmlu; unzip cmmlu_v1_0_1.zip
+cd ../../
+python evaluate_cmmlu.py -d data/cmmlu/
+'''
+
+def load_models_tokenizer(args):
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers.generation import GenerationConfig
+
+    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint_path, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(args.checkpoint_path, device_map="auto", trust_remote_code=True).eval()
+    model.generation_config = GenerationConfig.from_pretrained(args.checkpoint_path, trust_remote_code=True)
+    return model, tokenizer
+
+
+def format_example(line, include_answer=True):
+    example = '问题：' + line['Question']
+    for choice in choices:
+        example += f'\n{choice}. {line[f"{choice}"]}'
+
+    if include_answer:
+        example += '\n答案：' + line["Answer"] + '\n\n'
+    else:
+        example += '\n答案：'
+    return example
+
+
+def generate_few_shot_prompt(k, subject, dev_df):
+    prompt = ''
+    if k == -1:
+        k = dev_df.shape[0]
+    for i in range(k):
+        prompt += format_example(
+            dev_df.iloc[i, :],
+            include_answer=True,
+        )
+    return prompt
+
+
+def get_logits(tokenizer, model, inputs: List[str]):
+    input_ids = tokenizer(inputs, padding=False)['input_ids']
+    input_ids = torch.tensor(input_ids, device=model.device)
+    tokens = {'input_ids': input_ids}
+
+    outputs = model(input_ids)['logits']
+    logits = outputs[:, -1, :]
+    log_probs = torch.nn.functional.softmax(logits, dim=-1)
+    return log_probs, {'tokens': tokens}
+
+
+@torch.no_grad()
+def eval_subject(
+        model,
+        tokenizer,
+        subject_name,
+        test_df,
+        k=5,
+        dev_df=None,
+        few_shot=False,
+        save_result_dir=None,
+        **kwargs
+):
+    result = []
+    score = []
+
+    few_shot_prompt = generate_few_shot_prompt(
+        k, subject_name, dev_df) if few_shot else []
+    all_probs = {'prob_A': [], 'prob_B': [], 'prob_C': [], 'prob_D': []}
+    if args.debug: print(f"few_shot_prompt: {few_shot_prompt}")
+
+    for _, row in tqdm(test_df.iterrows(), total=len(test_df)):
+        question = format_example(row, include_answer=False)
+        full_prompt = few_shot_prompt + question
+
+        output, input_info = get_logits(tokenizer, model, [full_prompt])
+        assert output.shape[0] == 1
+        logits = output.flatten()
+
+        softval = torch.nn.functional.softmax(
+                torch.tensor(
+                    [
+                        logits[tokenizer("A")['input_ids']],
+                        logits[tokenizer("B")['input_ids']],
+                        logits[tokenizer("C")['input_ids']],
+                        logits[tokenizer("D")['input_ids']],
+                    ]
+                ),
+                dim=0,
+            )
+        if softval.dtype in {torch.bfloat16, torch.float16}:
+            softval = softval.to(dtype=torch.float32)
+        probs = softval.detach().cpu().numpy()
+
+        for i, choice in enumerate(choices):
+            all_probs[f'prob_{choice}'].append(probs[i])
+        pred = {0: "A", 1: "B", 2: "C", 3: "D"}[np.argmax(probs)]
+
+        if 'Answer' in row:
+            correct = 1 if pred == row['Answer'] else 0
+            score.append(correct)
+            if args.debug: print(f'{question} pred: {pred} ref: {row["Answer"]}')
+        result.append(pred)
+
+    if score:
+        correct_ratio = 100 * sum(score) / len(score)
+        if args.debug: print(subject_name, correct_ratio)
+    else:
+        correct_ratio = 0
+    if save_result_dir:
+        test_df['model_output'] = result
+        for i, choice in enumerate(choices):
+            test_df[f'prob_{choice}'] = (all_probs[f'prob_{choice}'])
+        if score:
+            test_df["correctness"] = score
+        os.makedirs(save_result_dir, exist_ok=True)
+        test_df.to_csv(os.path.join(
+            save_result_dir, f'{subject_name}_result.csv'), encoding="utf-8", index=False)
+
+    return correct_ratio
+
+
+def cal_cmmlu(res):
+    print('\n\n\n')
+    res = {k.split('-')[-1]:float(v) for k,v in res.items()}
+    for k, v in TASK_NAME_MAPPING.items():
+        avg_acc = np.mean(list(map(lambda x: res[x], v)))
+        print(f"{k} acc: {avg_acc:.2f}")
+    avg_all_acc = np.mean(list(res.values()))
+    print(f"AVERAGE acc: {avg_all_acc:.2f}")
+
+
+subcategories = {
+    "agronomy": ['other'],
+    "anatomy": ['biology'],
+    "ancient_chinese": ['linguistics','china specific'],
+    "arts": ['arts'],
+    "astronomy": ['physics'],
+    "business_ethics": ['business'],
+    "chinese_civil_service_exam": ['politics','china specific'],
+    "chinese_driving_rule": ['other','china specific'],
+    "chinese_food_culture": ['culture','china specific'],
+    "chinese_foreign_policy": ['politics','china specific'],
+    "chinese_history":['history','china specific'],
+    "chinese_literature": ['literature','china specific'],
+    "chinese_teacher_qualification": ['education','china specific'],
+    "college_actuarial_science":['math'],
+    "college_education":['education'],
+    "college_engineering_hydrology": ['engineering'],
+    "college_law": ['law'],
+    "college_mathematics": ['math'],
+    "college_medical_statistics":['statistics'],
+    "clinical_knowledge": ['other'],
+    "college_medicine": ['other'],
+    "computer_science": ['computer science'],
+    "computer_security": ['other'],
+    "conceptual_physics": ['physics'],
+    "construction_project_management": ['other','china specific'],
+    "economics": ['economics'],
+    "education": ['education'],
+    "elementary_chinese":['linguistics','china specific'],
+    "elementary_commonsense":['other','china specific'],
+    "elementary_information_and_technology": ['other'],
+    "electrical_engineering": ['engineering'],
+    "elementary_mathematics": ['math'],
+    "ethnology": ['culture','china specific'],
+    "food_science": ['other'],
+    "genetics": ['biology'],
+    "global_facts": ['global'],
+    "high_school_biology": ['biology'],
+    "high_school_chemistry": ['chemistry'],
+    "high_school_geography": ['geography'],
+    "high_school_mathematics": ['math'],
+    "high_school_physics": ['physics'],
+    "high_school_politics": ['politics','china specific'],
+    "human_sexuality": ['other'],
+    "international_law": ['law'],
+    "journalism": ['sociology'],
+    "jurisprudence": ['law'],
+    "legal_and_moral_basis": ['other'],
+    "logical": ['philosophy'],
+    "machine_learning": ['computer science'],
+    "management": ['business'],
+    "marketing": ['business'],
+    "marxist_theory": ['philosophy'],
+    "modern_chinese": ['linguistics','china specific'],
+    "nutrition": ['other'],
+    "philosophy": ['philosophy'],
+    "professional_accounting": ['business'],
+    "professional_law": ['law'],
+    "professional_medicine": ['other'],
+    "professional_psychology": ['psychology'],
+    "public_relations": ['politics'],
+    "security_study": ['politics'],
+    "sociology": ['culture'],
+    "sports_science": ['other'],
+    "traditional_chinese_medicine": ['other','china specific'],
+    "virology": ['biology'],
+    "world_history":['history'],
+    "world_religions": ['global'],
+}
+
+categories = {
+    "STEM": ["physics", "chemistry", "biology", "computer science", "math", "engineering", "statistics"],
+    "Humanities": ["history", "philosophy", "law", "arts", "literature", "global"],
+    "Social Science": ['linguistics',"business", "politics", "culture", "economics", "geography", "psychology", "education", "sociology"],
+    "Other":["other"],
+    "China specific": ["china specific"],
+}
+
+TASK_NAME_MAPPING = defaultdict(list)
+for k,v in categories.items():
+    for subject, subcat in subcategories.items():
+        for c in subcat:
+            if c in v:
+                TASK_NAME_MAPPING[k].append(subject)
+
+
+choices = ["A", "B", "C", "D"]
+
+
+def main(args):
+    model, tokenizer = load_models_tokenizer(args)
+
+    test_result = {}
+    for subject_name in tqdm(subcategories.keys()):
+        dev_file_path = os.path.join(args.eval_data_path, 'dev', f'{subject_name}.csv')
+        test_file_path = os.path.join(args.eval_data_path, 'test', f'{subject_name}.csv')
+        dev_df = pd.read_csv(dev_file_path)
+        test_df = pd.read_csv(test_file_path)
+
+        score = eval_subject(model, tokenizer, subject_name, dev_df=dev_df, test_df=test_df, k=5, few_shot=True,
+                             save_result_dir=f"outs/cmmlu_eval_result")
+        test_result[subject_name] = score
+    cal_cmmlu(test_result)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Test HF checkpoint.')
+    parser.add_argument('-c', '--checkpoint-path', type=str, help='Checkpoint path', default="Qwen/Qwen-7B")
+    parser.add_argument('-s', '--seed', type=int, default=1234, help='Random seed')
+
+    """Provide extra arguments required for tasks."""
+    group = parser.add_argument_group(title='Evaluation options')
+    group.add_argument('-d', '--eval_data_path', type=str, required=True,
+                       help='Path to eval data')
+    group.add_argument("--max-seq-len", type=int, default=2048,
+                       help='Size of the output generated text.')
+    group.add_argument("--debug", action='store_true', default=False,
+                       help='Print infos.')
+
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    main(args)