assets/evaluation.py

import os
import json
import re
import argparse
import torch
import gc
from tqdm import tqdm

from transformers import AutoModelForCausalLM, AutoModel, AutoTokenizer

os.environ['CUDA_VISIBLE_DEVICES'] = '1'


def parse_args():
    parser = argparse.ArgumentParser(description='Validation')
    parser.add_argument('--model_path', dest="model_path")
    parser.add_argument('--dataset_path', dest="dataset_path")
    parser.add_argument('--save_path', dest='save_path')
    parser.add_argument('--max_length', dest="max_length", default=2000)
    parser.add_argument('--max_new_tokens', dest="max_new_tokens", default=48)
    parser.add_argument('--input_template', dest="input_template",
                        default="材料：{material}\n问题：{question}\n{options}\n答案：{response}")
    parser.add_argument('--query_template', dest="query_template", default="问题：{question}\n{options}\n答案：{response}")
    parser.add_argument('--system_prompt', dest="system_prompt",
                        default="你是一名考生，请回答问题的正确选项，比如C。如果有多个正确选项，请按顺序回答所有正确选项，比如ABD。")
    parser.add_argument('--level_delimiter', dest="level_delimiter", default="|")
    args = parser.parse_args()
    return args


class Validator:
    def __init__(self, args):
        self.load_model(args.model_path)
        self.load_dataset(args.dataset_path)
        self.save_dir = os.path.join(args.save_path, os.path.split(model_path)[-1])
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)

        self.max_length = args.max_length
        self.max_new_tokens = args.max_new_tokens
        self.input_template = args.input_template
        self.query_template = args.query_template
        self.system_prompt = args.system_prompt
        self.level_delimiter = args.level_delimiter

    def load_model(self, model_path):
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

        self.model = AutoModelForCausalLM.from_pretrained(model_path, device_map={"": 0}, trust_remote_code=True)
        self.model.eval()

    def load_dataset(self, dataset_path):
        self.dataset = json.load(open(dataset_path, encoding="utf-8"))

    def format_prompt(self, material, question, options, response):
        if material:
            return self.input_template.format(material=material, question=question, options=options,
                                              response=response).strip()
        return self.query_template.format(question=question, options=options, response=response).strip()

    def build_prompt(self, item):
        query_prompt = self.format_prompt(item['material'], item['question'], item['options'], "")
        history_prompts = []
        for sub in item['history']:
            history_prompts.append(self.format_prompt(sub['material'], sub['question'], sub['options'], sub['choice']))

            final_prompt = self.system_prompt + "\n" + "\n".join(history_prompts) + "\n" + query_prompt

            if len(self.tokenizer.tokenize(final_prompt)) > self.max_length:
                history_prompts.pop()
                break
        return self.system_prompt + "\n" + "\n".join(history_prompts) + "\n" + query_prompt

    def get_predict(self, prompt):
        gen_kwargs = {"do_sample": False, "max_new_tokens": self.max_new_tokens}
        inputs = self.tokenizer([prompt], return_tensors="pt")
        inputs = inputs.to(self.model.device)
        with torch.no_grad():
            outputs = self.model.generate(**inputs, return_dict_in_generate=True, **gen_kwargs)
        predict = self.tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)[len(prompt):]
        return predict

    def extract_answer(self, results):
        predict_result = {"acc": 0, "wrong_value": 0, "human_acc": 0, "hit": 0, "total": 0, "wrong_hit": 0,
                          "wrong_total": 0, 'detail': []}
        for result in results:
            answer = result['item']['choice']
            most_wrong = result['item']['most_wrong']
            human_acc = result['item']['human_acc'] * 0.01

            predict = ""
            for e in re.split(r'[、,.，和\s]\s*', result['predict']):
                if not e.isascii():
                    break
                predict += e
            result['predict'] = predict
            predict_result['detail'].append(
                {"answer": answer, "most_wrong": most_wrong, "predict": predict, "human_acc": human_acc})
            predict_result['hit'] += 1 if answer == predict else 0
            predict_result['wrong_hit'] += 1 if most_wrong == predict else 0
            predict_result['wrong_value'] += 1 - human_acc if most_wrong == predict else 0
            predict_result['human_acc'] += human_acc
            predict_result['total'] += 1
        predict_result['acc'] = predict_result['hit'] / predict_result['total']
        predict_result['wrong_total'] = predict_result['total'] - predict_result['hit']
        predict_result['wrong_value'] = predict_result['wrong_value'] / predict_result['wrong_total']
        predict_result['human_acc'] = predict_result['human_acc'] / len(results)

        json.dump(predict_result, open(os.path.join(self.save_dir, "acc_result.json"), "w", encoding="utf-8"),
                  ensure_ascii=False)

    def category_summary(self, results):
        category_result = {"总计": {"hit": 0, "all": 0, "difficulty": {}, "human_acc": 0}}
        for result in results:
            hit = 1 if result['item']['choice'] == result['predict'] else 0
            categories_list = result['item']['categories']
            difficulty = result['item']['difficulty']
            human_acc = result['item']['human_acc']
            for categories in categories_list:
                if difficulty not in category_result["总计"]["difficulty"]:
                    category_result["总计"]["difficulty"][difficulty] = {"hit": 0, "all": 0}
                category_result["总计"]["difficulty"][difficulty]['hit'] += hit
                category_result["总计"]["difficulty"][difficulty]['all'] += 1
                category_result["总计"]['hit'] += hit
                category_result["总计"]['all'] += 1
                category_result["总计"]['human_acc'] += human_acc
                category_subset = []
                for category in categories:
                    category_subset.append(category)
                    category_name = self.level_delimiter.join(category_subset)
                    if not category_name:
                        category_name = "未分类"
                    if category_name not in category_result:
                        category_result[category_name] = {"hit": 0, "all": 0, "difficulty": {}, "human_acc": 0}
                    if difficulty not in category_result[category_name]["difficulty"]:
                        category_result[category_name]["difficulty"][difficulty] = {"hit": 0, "all": 0}
                    category_result[category_name]["difficulty"][difficulty]['hit'] += hit
                    category_result[category_name]["difficulty"][difficulty]['all'] += 1
                    category_result[category_name]['hit'] += hit
                    category_result[category_name]['all'] += 1
                    category_result[category_name]['human_acc'] += human_acc

        for k, v in category_result.items():
            v['acc'] = v['hit'] / v['all']
            v['human_acc'] = v['human_acc'] / v['all']
            for d, sub_v in v['difficulty'].items():
                sub_v['acc'] = sub_v['hit'] / sub_v['all']
        json.dump(category_result, open(os.path.join(self.save_dir, "category_result.json"), "w", encoding="utf-8"),
                  ensure_ascii=False)

    def difficulty_summary(self, results):
        difficulty_result = {"总计": {"hit": 0, "all": 0}}
        for result in results:
            hit = 1 if result['item']['choice'] == result['predict'] else 0
            difficulty = result['item']['difficulty']
            if difficulty not in difficulty_result:
                difficulty_result[difficulty] = {"hit": 0, "all": 0}
            difficulty_result[difficulty]['hit'] += hit
            difficulty_result[difficulty]['all'] += 1
            difficulty_result["总计"]['hit'] += hit
            difficulty_result["总计"]['all'] += 1

        for k in difficulty_result:
            difficulty_result[k]['acc'] = difficulty_result[k]['hit'] / difficulty_result[k]['all']

        json.dump(difficulty_result, open(os.path.join(self.save_dir, "difficulty_result.json"), "w", encoding="utf-8"),
                  ensure_ascii=False)

    def __call__(self):
        results = []
        for item in tqdm(self.dataset):
            prompt = self.build_prompt(item)
            predict = self.get_predict(prompt)
            results.append({"item": item, "predict": predict})
            gc.collect()

        json.dump(results, open(os.path.join(self.save_dir, "raw.json"), "w", encoding="utf-8"), ensure_ascii=False)
        self.extract_answer(results)
        self.category_summary(results)
        self.difficulty_summary(results)


if __name__ == "__main__":
    args = parse_args()
    Validator(args)()