Repository: hpcaitech/ColossalAI
File: applications/ColossalEval/colossal_eval/dataset/safetybench_en.py (Python)

Code before:
INST = lang2inst[LANGUAGE]
INPUT_FORMAT = lang2input_format[LANGUAGE]
FILES = lang2files[LANGUAGE]

CHOICE_TEMP = ["A. {}", "B. {}", "C. {}", "D. {}"]
IDX2CHOICE = {0: 'A', 1: 'B', 2: 'C', 3: 'D'}

default_inference_kwargs = {
    "calculate_loss": False,
    "all_classes": ["A", "B", "C", "D"],
    "language": LANGUAGE,
    "pretrain": False,
    "max_new_tokens": 32}

def get_query_str(question, options, choices_templates=CHOICE_TEMP, pad=True):
    # {'questions': 'what is xxx?\n', options: ['aaa', 'bbb', 'ccc', 'ddd'], ...}
    # --> 'what is xxx?\nA. aaa\n'
    query = ""
    query += question if question.endswith("\n") else question + "\n"
    num_choices = len(choices_templates)
    choices = []
    for idx, option in enumerate(options):
        choices.append(choices_templates[idx].format(option + '\n'))  # e.g. "A. xxxx\n", "B. xxxx\n", ...
    remain_choice = num_choices - len(choices)
    if pad and remain_choice > 0:  # use NULL choice to pad choices to max choices number 
        fake_choice = 'NULL'
        for i in range(remain_choice, num_choices):
            choices.append(choices_templates[i].format(fake_choice + '\n'))
    query += ''.join(choices)
    query = INPUT_FORMAT.format(query)
    return query

def process_test(sample_list):
    test_dict = {}
    for sample in sample_list:
        category = sample['category']
        if category not in test_dict:
            test_dict[category] = {'data': [], 'inference_kwargs': default_inference_kwargs}
        question = sample['question']
        options = sample['options']
        query_str = get_query_str(question, options, pad=True)
        data_sample = {
            'dataset': EVAL_NAME,
            'split': 'test',
            'category': category,
            'instruction': INST,
            'input': query_str,
            'output': '',
            'target': '',
            'id': sample['id']
        }

Review comment (suggestion):
Adding `NULL` choices may affect model's generation. Should just add the exact number of choices and set `"all_classes"` as `None` and thus rely on the generated content to decide which option the model chooses instead of the softmax over choices. (Or rearrange questions with the same length of options into the same category and ignore the original category such as Offensiveness)

Code after:
INPUT_FORMAT = lang2input_format[LANGUAGE]
FILES = lang2files[LANGUAGE]
PAD_CHOICES = True

CHOICE_TEMP = ["A. {}", "B. {}", "C. {}", "D. {}"]
IDX2CHOICE = {0: "A", 1: "B", 2: "C", 3: "D"}

default_inference_kwargs = {
    "calculate_loss": False,
    "all_classes": ["A", "B", "C", "D"],
    "language": LANGUAGE,
    "pretrain": False,
    "max_new_tokens": 32,
}


def get_query_str(question, options, choices_templates=CHOICE_TEMP, pad=True):
    # {'questions': 'what is xxx?\n', options: ['aaa', 'bbb', 'ccc', 'ddd'], ...}
    # --> 'what is xxx?\nA. aaa\nB. bbb\nC. ccc\nD. ddd\n'
    query = question if question.endswith("\n") else question + "\n"
    num_choices = len(choices_templates)

    choices = []
    for idx, option in enumerate(options):
        choices.append(choices_templates[idx].format(option + "\n"))  # e.g. "A. xxxx\n", "B. xxxx\n", ...
    remain_choice = num_choices - len(choices)
    if pad and remain_choice > 0:  # use NULL choice to pad choices to max choices number
        fake_choice = "NULL"
        for i in range(num_choices - remain_choice, num_choices):
            choices.append(choices_templates[i].format(fake_choice + "\n"))
    query += "".join(choices)
    query = INPUT_FORMAT.format(query)
    return query


def process_test(sample_list, pad_choices=False):
    test_dict = {}
    for sample in sample_list:
        num_options = len(sample["options"])
        category = sample["category"]
        inference_kwargs = deepcopy(default_inference_kwargs)
        if not pad_choices:
            category += "_{}".format(num_options)
            inference_kwargs["all_classes"] = inference_kwargs["all_classes"][:num_options]
        if category not in test_dict:
            test_dict[category] = {"data": [], "inference_kwargs": inference_kwargs}
        question = sample["question"]
        options = sample["options"]
        query_str = get_query_str(question, options, pad=pad_choices)
        data_sample = {
            "dataset": EVAL_NAME,