Function read_dataset

evaluation/utils.py:210–254

(
    data_file: str = None,
    dataset_type: str = "humanevalx",
)

def read_dataset(
    data_file: str = None,
    dataset_type: str = "humanevalx",
) -> Dict:
    if "humanevalx" in dataset_type.lower():
        dataset = {task["task_id"]: task for task in stream_jsonl(data_file)}
    elif "mbpp" in dataset_type.lower():
        problems = {task["task_id"]: task for task in stream_jsonl(data_file)}
        task_ids = sorted(problems.keys())[10:510]
        dataset = {}
        for task_id in task_ids:
            sample = problems[task_id]
            description = sample["text"]
            test_example = sample["test_list"][0]
            prompt = f'"""\n{description}\n{test_example}\n"""\n'
            sample["prompt"] = prompt
            dataset[task_id] = sample
    elif "ds1000" in dataset_type.lower():
        # install ds1000 from https://github.com/HKUNLP/DS-1000
        from ds1000 import DS1000Dataset
        ds1000 = DS1000Dataset(source_dir=data_file, libs="all", mode="Completion")
        for lib in ds1000.libs:
            for problem_id in range(len(ds1000[lib])):
                prefix = ""
                suffix = ""
                insert_flag = False
                first_line_flag = True
                # extract prefix and suffix of the prompt
                for line in ds1000[lib][problem_id]["prompt"].split("\n"):
                    if "[insert]" in line:
                        insert_flag = True
                        continue
                    if first_line_flag:
                        first_line_flag = False
                    else:
                        line = "\n" + line
                    if not insert_flag:
                        prefix += line
                    else:
                        suffix += line
            
    else:
        raise f"Dataset: {dataset_type} not supported."

    return dataset


def read_translation_dataset(