tasks.py

from dataclasses import dataclass, field
from datasets import load_dataset, Dataset
from functools import cached_property
from tqdm.auto import tqdm
from typing import Any, Optional, Protocol, Iterable, Callable
import logging
import pandas as pd
from functools import partial
from datasets.utils.logging import disable_progress_bar

from .utils import *

from evaluate import load
from collections import defaultdict
import sys


# if sys.version_info >= (3, 9):
#     from functools import cache
# else:
#     from functools import lru_cache as cache


disable_progress_bar()


def mt_bench_prompt(example):
    judge_prompt = "You are ChatGPT, a large language model trained by OpenAI. Please act as an impartial judge and evaluate the quality of the response provided by an AI assistant to the user question displayed below. The  Your evaluation should consider factors such as the helpfulness, relevance, accuracy, depth, creativity, and level of detail of the response."
    judge_prompt = "You are ChatGPT, a large language model trained by OpenAI. Your task is to act as an impartial judge and evaluate the quality of the responses provided by an 'assistant' role in the displayed conversation. Your evaluation should focus on the helpfulness, relevance, accuracy, depth, creativity, language fluency, clarity, and level of detail in the assistant's responses. Please note that the evaluation should not consider the user's questions or the overall conversation, but solely the quality of the assistant's replies."
    multi_prompt = "You evaluation should focus on the assistant's answer to the second user question."
    ref_prompt = "In the conversation, you will encounter system messages labeled 'Reference Answer' followed by the assistant's response. Your task is to evaluate the quality of the assistant's response by comparing it with the reference answer."
    json_prompt = 'You must rate the response on a scale of 1 to 10 in JSON format, for example: {"rating": 5}.'
    prompt_list = [judge_prompt]
    conversations = example["conversation"]
    if example["turn"] == 2:
        prompt_list.append(multi_prompt)

    if example["reference"] is not None:
        conversations = []
        quesiotns = filter(lambda e: e["role"] == "user", example["conversation"])
        answers = filter(lambda e: e["role"] == "assistant", example["conversation"])
        for q, a, r in zip(quesiotns, answers, example["reference"]):
            conversations.append(q)
            conversations.append(
                {"role": "system", "content": "Reference Answer: " + r}
            )
            conversations.append(a)
        prompt_list.append(ref_prompt)
    prompt_list.append(json_prompt)

    messages = [{"role": "system", "content": " ".join(prompt_list)}] + conversations
    return messages


@dataclass
class Task:
    dataset_name: str | tuple[str, str] = ("gsm8k", "main")
    split: str = "test"
    # metrics: list[str] = field(default_factory=list)
    metric_name: str | tuple[str, str] = ("sustech/tlem", "gsm8k")
    input_column: str = "question"
    label_column: str = ""
    prompt: Optional[Callable | str] = None
    few_shot: int = 0
    few_shot_from: Optional[str] = None
    # results: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        names = (
            [self.dataset_name]
            if isinstance(self.dataset_name, str)
            else list(self.dataset_name)
        )
        names[0] = names[0].split("/")[-1]

        self.name = "-".join(names) + f"-{self.split}"
        if isinstance(self.prompt, str):
            self.prompt = lambda example: {
                self.input_column: self.prompt.format(
                    input_column=example[self.input_column]
                )
            }
        self.label_column = self.label_column or self.input_column

    @cached_property
    def samples(self):
        return self.dataset[self.input_column]

    @cached_property
    def dataset(self):
        ds = load_dataset(
            *self.dataset_name
            if isinstance(self.dataset_name, tuple)
            else self.dataset_name,
            # split=self.split,
        )
        test_ds = ds[self.split]
        if self.prompt is not None:
            test_ds = test_ds.map(self.prompt)

        if self.few_shot:
            if self.few_shot_from is None:
                for name in ["train", "validation", "val", "dev"]:
                    if name in ds:
                        self.few_shot_from = name
                        break

            assert self.few_shot_from != self.split
            shots = ds[self.few_shot_from].select(range(self.few_shot))
            # else:
            #     shots = ds.select(range(self.few_shot))
            if self.prompt is not None:
                shots = shots.map(self.prompt)

            shots = shots.map(
                lambda example: {
                    self.input_column: example[self.input_column]
                    + example[self.label_column],
                }
            )[self.input_column]
            few_shot_prompts = "\n\n".join(shots)

            test_ds = test_ds.map(
                lambda example: {
                    self.input_column: few_shot_prompts
                    + "\n\n"
                    + example[self.input_column],
                }
            )

        return test_ds

    @cached_property
    def metric(self):
        metric = (
            load(self.metric_name)
            if isinstance(self.metric_name, str)
            else load(*self.metric_name)
        )
        return metric

    # @cache
    def run(
        self,
        pipeline,
    ):
        if (outputs := pipeline(self.samples)) is None:
            logging.warning("pipeline returns None")
            return
        self.outputs = outputs
        try:
            try:
                result = self.metric._compute(
                    responses=outputs, references=self.dataset[self.label_column]
                )
            except Exception as e:
                result = self.metric.compute(
                    responses=outputs, references=self.dataset[self.label_column]
                )
        except Exception as e:
            result = outputs
        # if log:
        #     name = name or pipeline.__name__
        #     self.results[name] = result

        return result


def multichoice(responses: Any, references: list[str]):
    if isinstance(responses[0], str):
        responses = [extract_choice(response) for response in responses]
    else:
        responses = decode_choice(responses)

    return responses, references


def multichoice_zh(responses: Any, references: list[str]):
    if isinstance(responses[0], str):
        responses = [extract_choice_zh(response) for response in responses]
    else:
        responses = decode_choice(responses)

    return responses, references


class Metrics:
    cmmlu = multichoice_zh
    mmlu = multichoice
    
    def ceval(responses: list[str], answers: list[str | int]):
        responses = [extract_choice_zh(pred) for pred in responses]
        return responses, answers
    
    def winogrande(responses: list[str], answers: list[str | int]):
        responses = [first_option_postprocess(pred, options="AB") for pred in responses]
        return responses, answers
    
    def arc(responses: list[str], answers: list[str | int]):
        if len(responses) != len(answers):
            return {
                'error': 'predictions and references have different '
                'length'
            }
        responses = [first_option_postprocess(pred, options="ABCD") for pred in responses]

        return responses, answers
    
    def hellaswag(responses: list[str], answers: list[str | int]):
        if len(responses) != len(answers):
            return {
                'error': 'predictions and references have different '
                'length'
            }
        responses = [first_option_postprocess(pred, options="ABCD") for pred in responses]
        answers = ['ABCD'[int(ans)] for ans in answers]
        return responses, answers
    
    def drop(responses: list[str], answers: list[list]):
        if len(responses) != len(answers):
            return {
                'error': 'predictions and references have different '
                'length'
            }
        responses = [general_postprocess(pred) for pred in responses]
        processed_answers = [[general_postprocess(j) for j in i]
                             for i in answers]
        matched_answers = []
        for pred, ans, origin_ans in zip(responses, processed_answers,
                                         answers):
            
            if pred in ans or pred in origin_ans:
                matched_answers.append(pred)
            else:
                matched_answers.append(ans[0])
            
        return responses, matched_answers
    
    def bbh_mcq(responses: list[str], answers: list[str | int]):
        if len(responses) != len(answers):
            return {
                'error': 'predictions and references have different '
                'length'
            }
        responses = [bbh_mcq_postprocess(pred) for pred in responses]

        return responses, answers
    
    def bbh_freefrom(responses: list[str], answers: list[str | int]):
        if len(responses) != len(answers):
            return {
                'error': 'predictions and references have different '
                'length'
            }

        responses = [bbh_freeform_postprocess(pred) for pred in responses]

        return responses, answers

    def gsm8k(responses: list[str], answers: list[str | int]):
        # scores = []
        # for response, answer in zip(responses, answers):
        #     pred = extract_numeric(response)
        #     gold = extract_numeric(answer) if isinstance(answer, str) else str(answer)
        #     scores.append(1.0 * (pred == gold))
        responses = [extract_numeric(response) for response in responses]
        answers = [
            extract_numeric(answer) if isinstance(answer, str) else str(answer)
            for answer in answers
        ]

        return responses, answers

    def MATH(responses: list[str], answers: list[str]):
        scores = []

        for response, answer in zip(responses, answers):
            indices = [pos for pos, char in enumerate(response) if char == "$"]
            if len(indices) <= 2:
                scores.append(0)
                continue
            else:
                result = response[indices[-2] + 1 : indices[-1]]
                gold = get_answer(answer)
                scores.append(1.0 * is_equiv(result, gold))

        return scores

    def math23k(responses: list[str], answers: list[str]):
        scores = []
        for response, answer in zip(responses, answers):
            pred = extract_numeric(response, pattern=NUMERIC_IN_ZH)
            gold = extract_numeric(answer, pattern=NUMERIC_IN_ZH)
            scores.append(1.0 * (pred == gold))
        return scores


class CMMLU:
    input_column = "prompt"
    label_column = "Answer"

    def prompt_cmmlu(example, chat=False):
        prefix = "以下是一道多项选择题，请从A、B、C和D中选择最合适的答案作为这个问题的答案。\n\n" if chat else "问题："
        prompt = prefix + example["Question"]
        for choice in list("ABCD"):
            prompt += f"\n{choice}. {example[choice]}"

            prompt += "\n答案："
        return {"prompt": prompt}

    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"],
        "Test": ["computer science"],
    }

    @classmethod
    def suite(cls, chat=False):
        finer_categories = (
            pd.Series(cls.subcategories)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        cls.categories["all"] = list(finer_categories.keys())
        for k, v in cls.categories.items():
            for subject in v:
                suite[k].extend(
                    [
                        Task(
                            ("haonan-li/cmmlu", subcategories),
                            metric_name=("sustech/tlem", "cmmlu"),
                            input_column=cls.input_column,
                            label_column=cls.label_column,
                            prompt=partial(cls.prompt_cmmlu, chat=chat),
                            few_shot=0 if chat else 5,
                            few_shot_from="dev",
                        )
                        for subcategories in finer_categories[subject]
                    ]
                )
        return suite


class MMLU:
    input_column = "prompt"
    label_column = "target"

    @classmethod
    def prompt_mmlu(cls, example, chat=False):
        prefix = (
            "The following is a multiple-choice question. Please choose the most suitable one among A, B, C and D as the answer to this question.\n\n"
            if chat
            else "Question: "
        )
        prompt = prefix + example["input"]
        for choice in list("ABCD"):
            prompt += f"\n{choice}. {example[choice]}"

        prompt += "\nAnswer："
        return {"prompt": prompt}

    subcategories = {
        "abstract_algebra": ["math"],
        "anatomy": ["health"],
        "astronomy": ["physics"],
        "business_ethics": ["business"],
        "clinical_knowledge": ["health"],
        "college_biology": ["biology"],
        "college_chemistry": ["chemistry"],
        "college_computer_science": ["computer science"],
        "college_mathematics": ["math"],
        "college_medicine": ["health"],
        "college_physics": ["physics"],
        "computer_security": ["computer science"],
        "conceptual_physics": ["physics"],
        "econometrics": ["economics"],
        "electrical_engineering": ["engineering"],
        "elementary_mathematics": ["math"],
        "formal_logic": ["philosophy"],
        "global_facts": ["other"],
        "high_school_biology": ["biology"],
        "high_school_chemistry": ["chemistry"],
        "high_school_computer_science": ["computer science"],
        "high_school_european_history": ["history"],
        "high_school_geography": ["geography"],
        "high_school_government_and_politics": ["politics"],
        "high_school_macroeconomics": ["economics"],
        "high_school_mathematics": ["math"],
        "high_school_microeconomics": ["economics"],
        "high_school_physics": ["physics"],
        "high_school_psychology": ["psychology"],
        "high_school_statistics": ["math"],
        "high_school_us_history": ["history"],
        "high_school_world_history": ["history"],
        "human_aging": ["health"],
        "human_sexuality": ["culture"],
        "international_law": ["law"],
        "jurisprudence": ["law"],
        "logical_fallacies": ["philosophy"],
        "machine_learning": ["computer science"],
        "management": ["business"],
        "marketing": ["business"],
        "medical_genetics": ["health"],
        "miscellaneous": ["other"],
        "moral_disputes": ["philosophy"],
        "moral_scenarios": ["philosophy"],
        "nutrition": ["health"],
        "philosophy": ["philosophy"],
        "prehistory": ["history"],
        "professional_accounting": ["other"],
        "professional_law": ["law"],
        "professional_medicine": ["health"],
        "professional_psychology": ["psychology"],
        "public_relations": ["politics"],
        "security_studies": ["politics"],
        "sociology": ["culture"],
        "us_foreign_policy": ["politics"],
        "virology": ["health"],
        "world_religions": ["philosophy"],
    }

    categories = {
        "STEM": [
            "physics",
            "chemistry",
            "biology",
            "computer science",
            "math",
            "engineering",
        ],
        "humanities": ["history", "philosophy", "law"],
        "social sciences": [
            "politics",
            "culture",
            "economics",
            "geography",
            "psychology",
        ],
        "other": ["other", "business", "health"],
    }

    @classmethod
    def suite(cls, chat=False):
        finer_categories = (
            pd.Series(cls.subcategories)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        cls.categories["all"] = list(finer_categories.keys())
        for k, v in cls.categories.items():
            for subject in v:
                suite[k].extend(
                    [
                        Task(
                            ("lukaemon/mmlu", subcategories),
                            metric_name=("sustech/tlem", "mmlu"),
                            input_column=cls.input_column,
                            label_column=cls.label_column,
                            prompt=partial(cls.prompt_mmlu, chat=chat),
                            few_shot=0 if chat else 5,
                            few_shot_from="validation",
                        )
                        for subcategories in finer_categories[subject]
                    ]
                )
        return suite


class Winogrande:
    input_column = "input"
    label_column = "answer"
    
    categories = [
        "winogrande_debiased",
        "winogrande_l",
        "winogrande_m",
        "winogrande_s",
        "winogrande_xl",
        "winogrande_xs",
    ]
    
    @classmethod
    def prompt_winogrande(cls, example):
        option1 = example["sentence"].replace("_", example['option1'])
        option2 = example["sentence"].replace("_", example['option2'])
        answer = example[cls.label_column]
        prompt = f"Which of the following is a good sentence:\nA. {option1}\nB. {option2}\nAnswer:"
        
        return {
            cls.input_column: prompt,
            cls.label_column: ' AB'[int(answer)] if answer != '' else ''
        }
        
    @classmethod
    def suite(cls,):
        subcategories = {
            item: [item] for item in cls.categories
        }
        finer_categories = (
            pd.Series(subcategories)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        subcategories["all"] = list(finer_categories.keys())
        for cate, sub_cates in subcategories.items():
            for sub_cate in sub_cates:
                suite[cate].append(
                    Task(
                        ("winogrande", sub_cate),
                        metric_name=("sustech/tlem", "winogrande"),
                        input_column=cls.input_column,
                        label_column=cls.label_column,
                        prompt=partial(cls.prompt_winogrande),
                        few_shot=0,
                        split="validation"
                    )
                )
        
        return suite
        


class DROP:
    input_column = "input"
    label_column = "answers"
    
    icl_prompt = '''\
Text: In the county, the population was spread out with 23.50% under the age of 18, 8.70% from 18 to 24, 29.70% from 25 to 44, 24.70% from 45 to 64, and 13.30% who were 65 years of age or older.
Question: How many more percent are under the age of 18 compared to the 18 to 24 group?
Anawer: According to the text, 23.5% are under the age of 18, and 8.7% are from ages 18 to 24. 23.5%-8.7%=14.8%. So the answer is 14.8.

Text: Playing in their second straight Thanksgiving game, the Eagles struggled especially on defense, where they were unable to stop the much-hyped Lions offense. The worst of it all was how unproven rookie Eric Rowe was tasked with covering wide receiver Calvin Johnson, leading to Johnson catching 3 touchdowns. Stafford’s five passing touchdowns, including three of them to Johnson was too much for the Eagles to overcome and for the second consecutive time this season, the Eagles gave up 45 points in a game. With the loss, the Eagles drop to 4-7 on the season and 6-1 when playing on Thanksgiving.
Question: How many TD passes did Stafford throw other than to Johnson?
Anawer: According to the text, Stafford threw 5 TD passes, 3 of which were to Johnson. 5-3=2. So the answer is 2.

Text: [PROMPT]
Question: [QUESTION]
Anawer:'''
    
    categories = ["validation"]
    
    @classmethod
    def prompt_drop(cls, example):
        
        prompt = cls.icl_prompt.replace("[PROMPT]", example["passage"]).replace("[QUESTION]", example["question"])
        
        validated_answers = example["answers_spans"]["spans"]
        validated_types = example["answers_spans"]["types"]
        answers = []
        for answer_item, answer_type in zip(validated_answers, validated_types):
            # if answer_type == "number":
            #     answers.append(answer_item)
            # elif any(answer_item['date'][i] for i in ['day', 'month', 'year']):
            #     d = [answer_item['date'][i] for i in ['day', 'month', 'year']]
            #     answers.append(' '.join(d).strip())
            # else:
            #     for span in answer_item['spans']:
                    # answers.append(span)
            answers.append(answer_item)
        answers = list(set(answers))
        
        return {
            cls.input_column: prompt,
            cls.label_column: answers
        }
    
    @classmethod
    def suite(cls,):
        
        finer_categories = (
            pd.Series(cls.categories)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        categories = list(finer_categories.keys())
        for cate in categories:
            suite[cate].append(
                Task(
                    ("drop", cate),
                    metric_name=("sustech/tlem", "drop"),
                    input_column=cls.input_column,
                    label_column=cls.label_column,
                    prompt=partial(cls.prompt_drop),
                    few_shot=0,
                    split="validation"
                )
            )
        
        return suite


class HellaSwag:
    input_column = "input"
    label_column = "label"
    
    categories = ["validation"]
    
    @classmethod
    def prompt_hellaswag(cls, example):
        
        prompt = f"{example['ctx']}\nQuestion: Which ending makes the most sense?\n"
        prompt += f"A. {example['endings'][0]}\n"
        prompt += f"B. {example['endings'][1]}\n"
        prompt += f"C. {example['endings'][2]}\n"
        prompt += f"D. {example['endings'][3]}\n"
        prompt += "You may choose from 'A', 'B', 'C', 'D'.\nAnswer:"
        
        return {cls.input_column: prompt}
    
    @classmethod
    def suite(cls,):
        finer_categories = (
            pd.Series(cls.categories)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        categories = list(finer_categories.keys())
        for cate in categories:
            suite[cate].append(
                Task(
                    ("Rowan/hellaswag", cate),
                    metric_name=("sustech/tlem", "hellaswag"),
                    input_column=cls.input_column,
                    label_column=cls.label_column,
                    prompt=partial(cls.prompt_hellaswag),
                    few_shot=0,
                    split="validation"
                )
            )
        
        return suite

class ARC:
    input_column = "input"
    label_column = "answerKey"
    
    categories = [
        "ARC-Challenge",
        "ARC-Easy",
    ]
    
    @classmethod
    def prompt_arc(cls, example):
        choices = example["choices"]
        prompt = f"Question: {example['question']}"
        for label, choice in zip(choices["label"], choices["text"]):
            prompt += f"\n{label}. {choice}"
        prompt += "\nAnswer:"
        return {
            cls.input_column: prompt
        }
    
    @classmethod
    def suite(cls):
        finer_categories = (
            pd.Series(cls.categories)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        categories = list(finer_categories.keys())
        for cate in categories:
            suite[cate].append(
                Task(
                    ("ai2_arc", cate),
                    metric_name=("sustech/tlem", "arc"),
                    input_column=cls.input_column,
                    label_column=cls.label_column,
                    prompt=partial(cls.prompt_arc),
                    few_shot=0,
                )
            )
        
        return suite


class BBH:
    input_column = "input"
    label_column = "target"
    
    multiple_choice_prefix = "Follow the given examples and answer the question.\n[HINT]\n\nQ: [INPUT]\nA: Let's think step by step."
    free_form_prefix = "Follow the given examples and answer the question.\n[HINT]\n\nQ: [INPUT]\nA: Let's think step by step."
    
    bbh_multiple_choice_sets = [
        'temporal_sequences',
        'disambiguation_qa',
        'date_understanding',
        'tracking_shuffled_objects_three_objects',
        'penguins_in_a_table',
        'geometric_shapes',
        'snarks',
        'ruin_names',
        'tracking_shuffled_objects_seven_objects',
        'tracking_shuffled_objects_five_objects',
        'logical_deduction_three_objects',
        'hyperbaton',
        'logical_deduction_five_objects',
        'logical_deduction_seven_objects',
        'movie_recommendation',
        'salient_translation_error_detection',
        'reasoning_about_colored_objects',
    ]
    
    bbh_free_form_sets = [
        'multistep_arithmetic_two',
        'navigate',
        'dyck_languages',
        'word_sorting',
        'sports_understanding',
        'boolean_expressions',
        'object_counting',
        'formal_fallacies',
        'causal_judgement',
        'web_of_lies',
    ]
    
    @classmethod
    def prompt_bbh(cls, example, category:str):
        
        meta_prompt = cls.multiple_choice_prefix if category in cls.bbh_multiple_choice_sets else cls.free_form_prefix
        prompt = meta_prompt.replace("[HINT]", bbh_lib_prompt(category=category)).replace("[INPUT]", example[cls.input_column])
        
        return {"input": prompt}
    
    @classmethod
    def suite(cls,):
        finer_categories = (
            pd.Series(cls.bbh_free_form_sets + cls.bbh_multiple_choice_sets)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        categories = list(finer_categories.keys())
        for cate in categories:
            if cate in cls.bbh_multiple_choice_sets:
                suite[cate].append(
                    Task(
                        ("lukaemon/bbh", cate),
                        metric_name=("sustech/tlem", "bbh_mcq"),
                        input_column=cls.input_column,
                        label_column=cls.label_column,
                        prompt=partial(cls.prompt_bbh, category=cate),
                        few_shot=0,
                    )
                )
            else:
                suite[cate].append(
                    Task(
                        ("lukaemon/bbh", cate),
                        metric_name=("sustech/tlem", "bbh_freefrom"),
                        input_column=cls.input_column,
                        label_column=cls.label_column,
                        prompt=partial(cls.prompt_bbh, category=cate),
                        few_shot=0,
                    )
                )
        
        return suite
    

    
class CEVAL:
    input_column = "input"
    label_column = "answer"
    
    @classmethod
    def prompt_ceval(cls, example, cate:str, chat=False):
        _ch_name = cls.ceval_subject_mapping[cate][1]
        prefix = (
            f"以下是中国关于{_ch_name}考试的单项选择题，请选出其中的正确答案。\n"
            if chat
            else "问题："
        )
        
        prompt = prefix + f'{example["question"]}'
        for choice in list("ABCD"):
            prompt += f"\n{choice}. {example[choice]}"

        prompt += "\n答案："
        return {"input": prompt}
    
    ceval_subject_mapping = {
        "computer_network":
        ["Computer Network", "\u8ba1\u7b97\u673a\u7f51\u7edc", "STEM"],
        "operating_system":
        ["Operating System", "\u64cd\u4f5c\u7cfb\u7edf", "STEM"],
        "computer_architecture":
        ["Computer Architecture", "\u8ba1\u7b97\u673a\u7ec4\u6210", "STEM"],
        "college_programming":
        ["College Programming", "\u5927\u5b66\u7f16\u7a0b", "STEM"],
        "college_physics": ["College Physics", "\u5927\u5b66\u7269\u7406", "STEM"],
        "college_chemistry":
        ["College Chemistry", "\u5927\u5b66\u5316\u5b66", "STEM"],
        "advanced_mathematics":
        ["Advanced Mathematics", "\u9ad8\u7b49\u6570\u5b66", "STEM"],
        "probability_and_statistics":
        ["Probability and Statistics", "\u6982\u7387\u7edf\u8ba1", "STEM"],
        "discrete_mathematics":
        ["Discrete Mathematics", "\u79bb\u6563\u6570\u5b66", "STEM"],
        "electrical_engineer": [
            "Electrical Engineer", "\u6ce8\u518c\u7535\u6c14\u5de5\u7a0b\u5e08",
            "STEM"
        ],
        "metrology_engineer":
        ["Metrology Engineer", "\u6ce8\u518c\u8ba1\u91cf\u5e08", "STEM"],
        "high_school_mathematics":
        ["High School Mathematics", "\u9ad8\u4e2d\u6570\u5b66", "STEM"],
        "high_school_physics":
        ["High School Physics", "\u9ad8\u4e2d\u7269\u7406", "STEM"],
        "high_school_chemistry":
        ["High School Chemistry", "\u9ad8\u4e2d\u5316\u5b66", "STEM"],
        "high_school_biology": [
            "High School Biology", "\u9ad8\u4e2d\u751f\u7269", "STEM"
        ],
        "middle_school_mathematics": [
            "Middle School Mathematics", "\u521d\u4e2d\u6570\u5b66", "STEM"
        ],
        "middle_school_biology": [
            "Middle School Biology", "\u521d\u4e2d\u751f\u7269", "STEM"
        ],
        "middle_school_physics": [
            "Middle School Physics", "\u521d\u4e2d\u7269\u7406", "STEM"
        ],
        "middle_school_chemistry": [
            "Middle School Chemistry", "\u521d\u4e2d\u5316\u5b66", "STEM"
        ],
        "veterinary_medicine": [
            "Veterinary Medicine", "\u517d\u533b\u5b66", "STEM"
        ],
        "college_economics": [
            "College Economics", "\u5927\u5b66\u7ecf\u6d4e\u5b66", "Social Science"
        ],
        "business_administration": [
            "Business Administration", "\u5de5\u5546\u7ba1\u7406", "Social Science"
        ],
        "marxism": [
            "Marxism", "\u9a6c\u514b\u601d\u4e3b\u4e49\u57fa\u672c\u539f\u7406",
            "Social Science"
        ],
        "mao_zedong_thought": [
            "Mao Zedong Thought",
            "\u6bdb\u6cfd\u4e1c\u601d\u60f3\u548c\u4e2d\u56fd\u7279\u8272\u793e\u4f1a\u4e3b\u4e49\u7406\u8bba\u4f53\u7cfb\u6982\u8bba",
            "Social Science"
        ],
        "education_science": [
            "Education Science", "\u6559\u80b2\u5b66", "Social Science"
        ],
        "teacher_qualification": [
            "Teacher Qualification", "\u6559\u5e08\u8d44\u683c", "Social Science"
        ],
        "high_school_politics": [
            "High School Politics", "\u9ad8\u4e2d\u653f\u6cbb", "Social Science"
        ],
        "high_school_geography": [
            "High School Geography", "\u9ad8\u4e2d\u5730\u7406", "Social Science"
        ],
        "middle_school_politics": [
            "Middle School Politics", "\u521d\u4e2d\u653f\u6cbb", "Social Science"
        ],
        "middle_school_geography": [
            "Middle School Geography", "\u521d\u4e2d\u5730\u7406", "Social Science"
        ],
        "modern_chinese_history":
        ["Modern Chinese History", "\u8fd1\u4ee3\u53f2\u7eb2\u8981", "Humanities"],
        "ideological_and_moral_cultivation": [
            "Ideological and Moral Cultivation",
            "\u601d\u60f3\u9053\u5fb7\u4fee\u517b\u4e0e\u6cd5\u5f8b\u57fa\u7840",
            "Humanities"
        ],
        "logic": ["Logic", "\u903b\u8f91\u5b66", "Humanities"],
        "law": ["Law", "\u6cd5\u5b66", "Humanities"],
        "chinese_language_and_literature": [
            "Chinese Language and Literature",
            "\u4e2d\u56fd\u8bed\u8a00\u6587\u5b66", "Humanities"
        ],
        "art_studies": ["Art Studies", "\u827a\u672f\u5b66", "Humanities"],
        "professional_tour_guide": [
            "Professional Tour Guide", "\u5bfc\u6e38\u8d44\u683c", "Humanities"
        ],
        "legal_professional": [
            "Legal Professional", "\u6cd5\u5f8b\u804c\u4e1a\u8d44\u683c",
            "Humanities"
        ],
        "high_school_chinese": [
            "High School Chinese", "\u9ad8\u4e2d\u8bed\u6587", "Humanities"
        ],
        "high_school_history": [
            "High School History", "\u9ad8\u4e2d\u5386\u53f2", "Humanities"
        ],
        "middle_school_history": [
            "Middle School History", "\u521d\u4e2d\u5386\u53f2", "Humanities"
        ],
        "civil_servant": ["Civil Servant", "\u516c\u52a1\u5458", "Other"],
        "sports_science": ["Sports Science", "\u4f53\u80b2\u5b66", "Other"],
        "plant_protection": [
            "Plant Protection", "\u690d\u7269\u4fdd\u62a4", "Other"
        ],
        "basic_medicine": ["Basic Medicine", "\u57fa\u7840\u533b\u5b66", "Other"],
        "clinical_medicine": [
            "Clinical Medicine", "\u4e34\u5e8a\u533b\u5b66", "Other"
        ],
        "urban_and_rural_planner": [
            "Urban and Rural Planner",
            "\u6ce8\u518c\u57ce\u4e61\u89c4\u5212\u5e08", "Other"
        ],
        "accountant": ["Accountant", "\u6ce8\u518c\u4f1a\u8ba1\u5e08", "Other"],
        "fire_engineer": [
            "Fire Engineer", "\u6ce8\u518c\u6d88\u9632\u5de5\u7a0b\u5e08", "Other"
        ],
        "environmental_impact_assessment_engineer": [
            "Environmental Impact Assessment Engineer",
            "\u73af\u5883\u5f71\u54cd\u8bc4\u4ef7\u5de5\u7a0b\u5e08", "Other"
        ],
        "tax_accountant": ["Tax Accountant", "\u7a0e\u52a1\u5e08", "Other"],
        "physician": ["Physician", "\u533b\u5e08\u8d44\u683c", "Other"]
    }
    
    @classmethod
    def suite(cls, chat: bool):
        suite = defaultdict(list)
        cls.categories = defaultdict(list)
        for task, info in cls.ceval_subject_mapping.items():
            cls.categories[info[0]].append(task)
            cls.categories[info[2]].append(task)
        cls.categories["all"] = list(cls.ceval_subject_mapping.keys())
        for k, v in cls.categories.items():
            for subject in v:
                suite[k].append(
                    Task(
                        dataset_name=("ceval/ceval-exam", subject),
                        metric_name=("sustech/tlem", "ceval"),
                        input_column=cls.input_column,
                        label_column=cls.label_column,
                        prompt=partial(cls.prompt_ceval, cate=subject, chat=chat),
                        few_shot=0 if chat else 5,
                        few_shot_from="dev",
                        split="val"
                    )
                )
                
        return suite