def prepare_groups()

in src/scripts/extract_data.py [0:0]


def prepare_groups(entity_name, dataset_name, indices_path):
    keys = create_indices(dataset_name, indices_path)
    with open("./data/prompts.json", "r") as f:
        groups = json.load(f)["group"][entity_name]
    with open(f"./data/{dataset_name}/{entity_name}/grouped_completions.json", "r") as f:
        tmp = json.load(f)
        data = {key: tmp[key] for key in keys}
    for key in keys:
        data[key] = process_group(data[key], groups)
    return data