nvidia/NVLM-D-72B

1import json2import os3import sys4import time5import yaml6import spacy7import ast8from PIL import Image9from glob import glob10from tqdm import tqdm11from collections import defaultdict12import pandas as pd13from io import BytesIO14import base6415from anls import anls_score16import torch17from torch.utils.data import Dataset, DataLoader, DistributedSampler18import torchvision.transforms as T19from eval import conversation as conversation_lib20from eval.mmmu_utils import CAT_SHORT2LONG, DOMAIN_CAT2SUB_CAT, parse_multi_choice_response, parse_open_response, \21    process_single_sample, construct_prompt, mmmu_main_eval, process_single_sample_pro, construct_prompt_pro22from eval.mmmu_utils import evaluate as evaluate_mmmu23from torchvision.transforms.functional import InterpolationMode24from datasets import load_dataset, concatenate_datasets25 26IMAGENET_MEAN = (0.485, 0.456, 0.406)27IMAGENET_STD = (0.229, 0.224, 0.225)28 29 30def build_transform(input_size):31    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD32    transform = T.Compose([33        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),34        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),35        T.ToTensor(),36        T.Normalize(mean=MEAN, std=STD)37    ])38    return transform39 40 41def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):42    best_ratio_diff = float('inf')43    best_ratio = (1, 1)44    area = width * height45    for ratio in target_ratios:46        target_aspect_ratio = ratio[0] / ratio[1]47        ratio_diff = abs(aspect_ratio - target_aspect_ratio)48        if ratio_diff < best_ratio_diff:49            best_ratio_diff = ratio_diff50            best_ratio = ratio51        elif ratio_diff == best_ratio_diff:52            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:53                best_ratio = ratio54    return best_ratio55 56 57def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnail=False):58    orig_width, orig_height = image.size59    aspect_ratio = orig_width / orig_height60 61    # calculate the existing image aspect ratio62    target_ratios = set(63        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if64        i * j <= max_num and i * j >= min_num)65    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])66 67    # find the closest aspect ratio to the target68    target_aspect_ratio = find_closest_aspect_ratio(69        aspect_ratio, target_ratios, orig_width, orig_height, image_size)70 71    # calculate the target width and height72    target_width = image_size * target_aspect_ratio[0]73    target_height = image_size * target_aspect_ratio[1]74    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]75 76    # resize the image77    resized_img = image.resize((target_width, target_height))78    processed_images = []79    for i in range(blocks):80        box = (81            (i % (target_width // image_size)) * image_size,82            (i // (target_width // image_size)) * image_size,83            ((i % (target_width // image_size)) + 1) * image_size,84            ((i // (target_width // image_size)) + 1) * image_size85        )86        # split the image87        split_img = resized_img.crop(box)88        processed_images.append(split_img)89    assert len(processed_images) == blocks90    if use_thumbnail and len(processed_images) != 1:91        thumbnail_img = image.resize((image_size, image_size))92        processed_images.append(thumbnail_img)93    return processed_images94 95 96def load_image(image, input_size=448, max_num=6, decoded=False):97    if not decoded:98        image = Image.open(image).convert('RGB')99    transform = build_transform(input_size=input_size)100    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)101    pixel_values = [transform(image) for image in images]102    pixel_values = torch.stack(pixel_values)103    return pixel_values104 105 106def levenshtein_distance(s1, s2):107    if len(s1) > len(s2):108        s1, s2 = s2, s1109 110    distances = range(len(s1) + 1)111    for i2, c2 in enumerate(s2):112        distances_ = [i2 + 1]113        for i1, c1 in enumerate(s1):114            if c1 == c2:115                distances_.append(distances[i1])116            else:117                distances_.append(1 + min((distances[i1], distances[i1 + 1], distances_[-1])))118        distances = distances_119    return distances[-1]120 121 122def get_anls_score(pred, gold_labels, threshold, llava_eval=False):123    values = []124    for answer in gold_labels:125        # preprocess both the answers - gt and prediction126        gt_answer = ' '.join(answer.strip().lower().split())127        det_answer = ' '.join(pred.strip().lower().split())128 129        dist = levenshtein_distance(gt_answer, det_answer)130        length = max(len(answer.upper()), len(pred.upper()))131        values.append(0.0 if length == 0 else float(dist) / float(length))132 133    question_result = 1 - min(values)134 135    if llava_eval:136        question_result = 1.0 if question_result >= threshold else 0.0137    else:138        if (question_result < threshold):139            question_result = 0140 141    return question_result142 143 144def isNumber(n: str):145    try:146        float(n)147        return True148    except ValueError:149        return False150 151 152class COCOEvalDataset(Dataset):153    def __init__(self, args, img_dir, subset=None):154        self.args = args155        self.img_files = sorted(glob(os.path.join(img_dir, "*")))156 157        if subset:158            self.img_files = self.img_files[:subset]159 160        self.image_ids = [int(img_file.split("_")[-1].split(".")[0]) for img_file in self.img_files]161 162    def __len__(self):163        return len(self.img_files)164 165    def __getitem__(self, idx):166        img_path = self.img_files[idx]167        img = load_image(img_path, max_num=6).to(torch.bfloat16)168 169        return self.image_ids[idx], img170 171 172class Flickr30KEvalDataset(Dataset):173    def __init__(self, args, img_dir, subset=None):174        self.args = args175        self.img_dir = img_dir176        self.test_samples = json.load(open(os.path.join(img_dir, "flickr30k_test.json"), encoding='utf-8'))177 178        if subset:179            self.test_samples = self.test_samples[:subset]180 181    def __len__(self):182        return len(self.test_samples)183 184    def __getitem__(self, idx):185        img_path = os.path.join(self.img_dir, self.test_samples[idx]["image"])186        img = load_image(img_path, max_num=6).to(torch.bfloat16)187 188        image_id = int(self.test_samples[idx]["image"].split("/")[-1].replace(".jpg", ""))189 190        return image_id, img191 192 193class VQAv2EvalDataset(Dataset):194    def __init__(self, args, img_dir, gt_path, subset=None):195        self.args = args196        self.img_dir = img_dir197        self.gt = json.load(open(gt_path, encoding='utf-8'))198 199        if subset:200            self.gt = self.gt[:subset]201 202    def __len__(self):203        return len(self.gt)204 205    def __getitem__(self, idx):206        img_path = os.path.join(self.img_dir, self.gt[idx]["image"])207        img = load_image(img_path, max_num=6).to(torch.bfloat16)208 209        question_id = self.gt[idx]["question_id"]210        question = self.gt[idx]["question"]211        answer = self.gt[idx]["answer"]212 213        return img, question_id, question, answer214 215 216class TextVQAEvalDataset(Dataset):217    def __init__(self, args, img_dir, gt_path, subset=None):218        self.args = args219        self.img_dir = img_dir220        self.gt = json.load(open(gt_path, encoding='utf-8'))['data']221 222        if subset:223            self.gt = self.gt[:subset]224 225    def __len__(self):226        return len(self.gt)227 228    def __getitem__(self, idx):229        img_path = os.path.join(self.img_dir, self.gt[idx]["image_id"] + '.jpg')230        if not os.path.exists(img_path):231            img_path = img_path.replace('.jpg', '.png')232        img = load_image(img_path, max_num=6).to(torch.bfloat16)233 234        question_id = self.gt[idx]["question_id"]235        question = self.gt[idx]["question"]236        answer = self.gt[idx]["answers"]237 238        return img, question_id, question, answer239 240 241class GQAEvalDataset(Dataset):242    def __init__(self, args, img_dir, gt_path, subset=None):243        self.args = args244        self.img_dir = img_dir245        self.gt = json.load(open(gt_path, encoding='utf-8'))246        self.gt = [{247            "question_id": int(k),248            "image": v['imageId'] + ".jpg",249            "question": v['question'],250            "answer": v['answer']251        } for k, v in self.gt.items()]252 253        if subset:254            self.gt = self.gt[:subset]255 256    def __len__(self):257        return len(self.gt)258 259    def __getitem__(self, idx):260        img_path = os.path.join(self.img_dir, self.gt[idx]["image"])261        img = load_image(img_path, max_num=6).to(torch.bfloat16)262 263        question_id = self.gt[idx]["question_id"]264        question = self.gt[idx]["question"]265        answer = self.gt[idx]["answer"]266 267        return img, question_id, question, [answer]268 269 270class ChartQAEvalDataset(Dataset):271    def __init__(self, args, img_dir, gt_path, subset=None):272        self.args = args273        self.img_dir = img_dir274        self.gt = json.load(open(gt_path, encoding='utf-8'))275        for i in range(len(self.gt)):276            self.gt[i]['question_id'] = i277 278        if subset:279            self.gt = self.gt[:subset]280 281    def __len__(self):282        return len(self.gt)283 284    def __getitem__(self, idx):285        img_path = os.path.join(self.img_dir, self.gt[idx]["imgname"])286        img = load_image(img_path, max_num=6).to(torch.bfloat16)287 288        question_id = self.gt[idx]["question_id"]289        question = self.gt[idx]["query"]290        answer = self.gt[idx]["label"]291 292        return img, question_id, question, [answer]293 294 295class OKVQAEvalDataset(Dataset):296    def __init__(self, args, img_dir, gt_path, question_path, subset=None):297        self.args = args298        self.img_dir = img_dir299        self.gt = json.load(open(gt_path, encoding='utf-8'))['annotations']300        self.questions = json.load(open(question_path, 'r'))['questions']301 302        if subset:303            self.gt = self.gt[:subset]304 305        qid2q = {q['question_id']: q['question'] for q in self.questions}306 307        for ann in self.gt:308            ann['answers'] = [ans['answer'] for ans in ann['answers']]309            ann['question'] = qid2q[ann['question_id']]310 311    def __len__(self):312        return len(self.gt)313 314    def __getitem__(self, idx):315        img_id = str(self.gt[idx]["image_id"])316        img_id = '0' * (12 - len(img_id)) + img_id317        img_file_name = f"COCO_val2014_{img_id}.jpg"318        img_path = os.path.join(self.img_dir, img_file_name)319        img = load_image(img_path, max_num=6).to(torch.bfloat16)320 321        question_id = self.gt[idx]["question_id"]322        question = self.gt[idx]["question"]323        answer = self.gt[idx]["answers"]324 325        return img, question_id, question, answer326 327 328class DocVQAEvalDataset(Dataset):329    def __init__(self, args, img_dir, gt_path, split='val', subset=None):330        self.args = args331        self.img_dir = img_dir332        self.gt = json.load(open(gt_path, encoding='utf-8'))['data']333 334        if subset:335            self.gt = self.gt[:subset]336 337        self.split = split338 339    def __len__(self):340        return len(self.gt)341 342    def __getitem__(self, idx):343        img_path = os.path.join(self.img_dir, self.gt[idx]['image'].split('/')[-1])344        img = load_image(img_path, max_num=6).to(torch.bfloat16)345 346        question_id = self.gt[idx]["questionId"]347        question = self.gt[idx]["question"]348 349        if self.split == 'val':350            answer = self.gt[idx]["answers"]351        else:352            answer = ['']353 354        return img, question_id, question, answer355 356 357class OCRBenchEvalDataset(Dataset):358    def __init__(self, args, img_dir, gt_path, subset=None):359        self.args = args360        self.img_dir = img_dir361        self.gt = json.load(open(gt_path, encoding='utf-8'))362 363        if subset:364            self.gt = self.gt[:subset]365 366    def __len__(self):367        return len(self.gt)368 369    def __getitem__(self, idx):370        img_path = os.path.join(self.img_dir, self.gt[idx]['image_path'])371        img = load_image(img_path, max_num=6).to(torch.bfloat16)372 373        dataset_name = self.gt[idx]["dataset_name"]374        question_id = f"{idx}"375        question = self.gt[idx]["question"]376        answer = self.gt[idx]["answers"]377        data_type = self.gt[idx]["type"]378 379        return img, question_id, question, answer, dataset_name, data_type380 381 382class AI2DiagramEvalDataset(Dataset):383    def __init__(self, args, img_dir, gt_path, subset=None):384        self.args = args385        self.img_dir = img_dir386 387        with open(gt_path, 'r') as json_file:388            json_list = list(json_file)389        self.gt = [json.loads(json_str) for json_str in json_list]390 391        if subset:392            self.gt = self.gt[:subset]393 394    def __len__(self):395        return len(self.gt)396 397    def __getitem__(self, idx):398        img_path = os.path.join(self.img_dir, self.gt[idx]['image'])399        img = load_image(img_path, max_num=6).to(torch.bfloat16)400 401        question_id = self.gt[idx]["question_id"]402        question = self.gt[idx]["question"]403        answer = self.gt[idx]["answer"]404 405        return img, question_id, question, answer406 407 408class AI2DiagramNoMaskEvalDataset(Dataset):409    def __init__(self, args, img_dir, gt_path, subset=None):410        self.args = args411        self.img_dir = img_dir412 413        with open(gt_path, 'r') as json_file:414            json_list = list(json_file)415        self.gt = [json.loads(json_str) for json_str in json_list]416 417        if subset:418            self.gt = self.gt[:subset]419 420    def __len__(self):421        return len(self.gt)422 423    def __getitem__(self, idx):424        img_file_name = self.gt[idx]['image'].replace("AI2D_TEST", "AI2D_TEST_NO_MASK_IMAGES")425        img_path = os.path.join(self.img_dir, img_file_name)426        img = load_image(img_path, max_num=6).to(torch.bfloat16)427 428        question_id = self.gt[idx]["question_id"]429        question = self.gt[idx]["question"]430        answer = self.gt[idx]["answer"]431 432        return img, question_id, question, answer433 434 435class RealworldQAEvalDataset(Dataset):436    def __init__(self, args, img_dir, gt_path, subset=None):437        self.args = args438        self.img_dir = img_dir439        self.gt = json.load(open(gt_path, encoding='utf-8'))440 441        if subset:442            self.gt = self.gt[:subset]443 444    def __len__(self):445        return len(self.gt)446 447    def __getitem__(self, idx):448        img_path = os.path.join(self.img_dir, self.gt[idx]['image'])449        img = load_image(img_path, max_num=6).to(torch.bfloat16)450 451        question_id = int(self.gt[idx]['image'].replace(".webp", ""))452        question = self.gt[idx]["question"]453 454        if self.gt[idx]['question_type'] == "multi-choice":455            choices = self.gt[idx]["choices"]456            start_chr = 'A'457            choices_str = ''458            index2ans = {}459            all_choices = []460            for choice in choices:461                all_choices.append(start_chr)462                index2ans[start_chr] = choice463                choices_str += f"{start_chr}. {choice}\n"464                start_chr = chr(ord(start_chr) + 1)465 466            question = question + '\n' + choices_str467            question = question + "Answer with the option's letter from the given choices directly."468            answer = chr(ord('A') + self.gt[idx]['correct_choice_index'])469        else:470            question = question + "\nAnswer the question using a single word or phrase."471            answer = self.gt[idx]['answer']472 473        return img, question_id, question, [answer]474 475 476class MathVistaEvalDataset(Dataset):477    def __init__(self, args, task_cfg, gt_path=None):478        self.args = args479        self.task_cfg = task_cfg480        self.dataset = load_dataset("AI4Math/MathVista")['testmini']481 482    def __len__(self):483        return len(self.dataset)484 485    def __getitem__(self, idx):486        img = self.dataset[idx]['decoded_image']487        img = load_image(img.convert("RGB"), max_num=6, decoded=True).to(torch.bfloat16)488 489        question_id = self.dataset[idx]["pid"]490        question = self.dataset[idx]["question"]491        question_type = self.dataset[idx]["question_type"]  # free_form or multi_choice492        query = self.dataset[idx]["query"]493        choices = self.dataset[idx]["choices"]494        answer = self.dataset[idx]["answer"]495 496        if question_type == 'multi_choice':497            start_chr = 'A'498            choices_str = ''499            index2ans = {}500            all_choices = []501            for choice in choices:502                all_choices.append(start_chr)503                index2ans[start_chr] = choice504                choices_str += f"{start_chr}. {choice}\n"505                start_chr = chr(ord(start_chr) + 1)506 507            question = question + '\n' + choices_str508            question = question + "Answer with the option's letter from the given choices directly."509            answer = chr(ord('A') + choices.index(answer))510        else:511            question = query.replace("Hint: ", "")512            index2ans = {}513            all_choices = []514 515        return img, question_id, question_type, question, answer, str(index2ans), str(all_choices)516 517 518def construct_prompt_for_fewshot(sample):519    config = {520        "task_instructions": "",521        "multi_choice_example_format": "{}\n{}Answer with the option's letter from the given choices directly.",522        "short_ans_example_format": "{}\nAnswer the question using a single word or phrase."523    }524 525    question = sample['question'].strip()526 527 528    options = eval(sample['options'])529    example = ""530    if sample['question_type'] == 'multiple-choice':531        start_chr = 'A'532        prediction_range = []533        index2ans = {}534        for option in options:535            prediction_range.append(start_chr)536            example += f"({start_chr}) {option}\n"537            index2ans[start_chr] = option538            start_chr = chr(ord(start_chr) + 1)539        empty_prompt_sample_structure = config['multi_choice_example_format']540        empty_prompt = empty_prompt_sample_structure.format(question, example)541        res_dict = {'type': 'multichoice'}542        res_dict['index2ans'] = index2ans543        res_dict['correct_choice'] = sample['answer']544        res_dict['all_choices'] = prediction_range545        res_dict['empty_prompt'] = empty_prompt546        if config['task_instructions']:547            res_dict['final_input_prompt'] = config['task_instructions'].strip() + '\n\n' + empty_prompt548        else:549            res_dict['final_input_prompt'] = empty_prompt550 551        res_dict['gt_content'] = options[ord(sample['answer'].upper()) - ord('A')]552    else:553        empty_prompt_sample_structure = config['short_ans_example_format']554        empty_prompt = empty_prompt_sample_structure.format(question)555        res_dict = {'type': 'open'}556        res_dict['empty_prompt'] = empty_prompt557        if config['task_instructions']:558            res_dict['final_input_prompt'] = config['task_instructions'].strip() + '\n\n' + empty_prompt559        else:560            res_dict['final_input_prompt'] = empty_prompt561        res_dict['gt_content'] = sample['answer']562 563    res_dict.update(sample)564    return res_dict565 566 567def process_image_tag(q):568    q = q.strip()569 570    # heuristic way of removing <image 1>571    if q == '<image 1>':572        q = 'Answer the question in the image.'573    elif ':<image 1>' in q:574        q = q.replace(':<image 1>', ' in the image. ')575        q = q.strip()576    elif ': <image 1>' in q:577        q = q.replace(': <image 1>', ' in the image. ')578        q = q.strip()579    elif '.<image 1>' in q or '. <image 1>' in q:580        q_list = q.split('<image 1>')581        q_list = [part.strip() for part in q_list if part.strip() != '']582        q = ' '.join(q_list)583    elif q.startswith('<image 1> '):584        if q[10].isupper():585            q = q.replace('<image 1>', '')586        else:587            q = q.replace('<image 1>', 'The image')588        q = q.strip()589    elif q.startswith('<image 1>'):590        q = q.replace('<image 1>', '')591    elif q.endswith('<image 1>?'):592        q = q.replace('<image 1>', 'the image')593    elif q.endswith('?<image 1>') or q.endswith('? <image 1>') or q.endswith('\n<image 1>'):594        q = q.replace('<image 1>', '')595        q = q.strip()596    elif ' <image 1> ' in q:597        q = q.replace('<image 1>', 'the image')598    elif ' <image 1>' in q:599        q = q.replace('<image 1>', 'the image')600    elif '()<image 1>' in q:601        q = q.replace('()<image 1>', '')602    elif '(<image 1>)' in q:603        q = q.replace('(<image 1>)', '')604    elif '<image 1>.' in q:605        q = q.replace("<image 1>.", ". ")606    else:607        q = q.replace("<image 1>", ". ")608        q = q.strip()609 610    # remove <image 2> to <image 8>611    for i in range(2, 8):612        q = q.replace(f"<image {i}>", "")613 614    return q615 616 617class MMMUProEvalDataset(Dataset):618    def __init__(self, args, task_cfg, subset=None):619        self.args = args620        self.task_cfg = task_cfg621        sub_dataset_list = []622        # load_dataset will throw error if split is 'dev'623        # 'dev' is part of the 'validation' and we need to manually split them624 625        MMMU_path = "MMMU/MMMU_Pro"626 627        _split = "test"628 629        self.dataset = load_dataset(MMMU_path, "standard", split=_split)630        if subset:631            self.dataset = self.dataset[:subset]632 633    def __len__(self):634        return len(self.dataset)635 636    def __getitem__(self, idx):637        # ===== single-image =====638        sample = self.dataset[idx]639        sample = process_single_sample_pro(sample)640        sample = construct_prompt_pro(sample, self.task_cfg)641        img = load_image(sample['image'].convert("RGB"), max_num=6, decoded=True).to(torch.bfloat16)642 643        # img = img.reshape(-1, 3, self.args.img_h, self.args.img_w)644 645        question_id = sample['id']646        question = sample['final_input_prompt']647        answer = sample['answer']648 649        question = process_image_tag(question)650        question = self.task_cfg['default_image_token'] + '\n' + question651 652        if sample['question_type'] == 'multiple-choice':653            index2ans = sample['index2ans']654            all_choices = sample['all_choices']655        else:656            index2ans = {}657            all_choices = []658 659        return img, question_id, sample['subfield'], sample['question_type'], question, answer, str(index2ans), str \660            (all_choices)661 662 663class MMMUEvalDataset(Dataset):664    def __init__(self, args, task_cfg, subset=None, start_idx=None):665        self.args = args666        self.task_cfg = task_cfg667        sub_dataset_list = []668        # load_dataset will throw error if split is 'dev'669        # 'dev' is part of the 'validation' and we need to manually split them670 671        MMMU_path = "MMMU/MMMU"672 673        _split = "test" if task_cfg["split"] == "test" else "validation"674        for subject in CAT_SHORT2LONG.values():675            sub_dataset = load_dataset(676                MMMU_path, subject,677                split=_split,678            )679            sub_dataset_list.append(sub_dataset)680 681        dataset = concatenate_datasets(sub_dataset_list)682 683        if task_cfg["split"] != "test":684            dataset = [s for s in dataset if s['id'].startswith(task_cfg["split"])]685 686        # dataset = [s for s in dataset if s['image_2'] is not None][1:]687 688        self.dataset = dataset689 690        if subset:691            self.dataset = [dataset[i] for i in range(start_idx, min(start_idx + subset, len(dataset)))]692            print(f"Evaluating a subset of dataset: {len(self.dataset)} from {start_idx} to {start_idx + subset}")693 694    def __len__(self):695        return len(self.dataset)696 697    def __getitem__(self, idx):698        # ===== single-image =====699        sample = self.dataset[idx]700        sample = process_single_sample(sample)701        sample = construct_prompt(sample, self.task_cfg)702 703        img = load_image(sample['image'].convert("RGB"), max_num=6, decoded=True).to(torch.bfloat16)704 705        question_id = sample['id']706        question = sample['final_input_prompt']707        answer = sample['answer']708 709        question = process_image_tag(question)710        question = self.task_cfg['default_image_token'] + '\n' + question711 712 713        if sample['question_type'] == 'multiple-choice':714            index2ans = sample['index2ans']715            all_choices = sample['all_choices']716        else:717            index2ans = {}718            all_choices = []719 720        return img, question_id, sample['subfield'], sample['question_type'], question, answer, str(index2ans), str \721            (all_choices)722 723 724 725class VizWizEvalDataset(Dataset):726    def __init__(self, args, img_dir, question_path, subset=None):727        self.args = args728        self.img_dir = img_dir729        self.questions = json.load(open(question_path, encoding='utf-8'))730 731    def __len__(self):732        return len(self.questions)733 734    def __getitem__(self, idx):735        img_path = os.path.join(self.img_dir, self.questions[idx]["image"])736        img = load_image(img_path, max_num=6).to(torch.bfloat16)737        question = self.questions[idx]["question"]738        question_id = self.questions[idx]["image"]739 740        return img, question_id, question741 742 743class MMBenchEvalDataset(Dataset):744    def __init__(self, args, gt_path, subset=None):745        self.args = args746        df = pd.read_csv(gt_path, sep='\t')747        self.dataset = []748        for i, row in df.iterrows():749            choices = []750            for choice in ['A', 'B', 'C', 'D']:751                if str(row[choice]) != 'nan':752                    choices.append(row[choice])753 754            this_sample = {755                'index': row['index'],756                'question': row['question'],757                'hint': row['hint'],758                'category': row['category'],759                'image': Image.open(BytesIO(base64.b64decode(row['image']))),760                'choices': choices761            }762 763            # Only dev set gives the ground truth answer764            if 'answer' in row.keys():765                this_sample['answer'] = row['answer']766            else:767                this_sample['answer'] = ''768 769            self.dataset.append(this_sample)770 771    def __len__(self):772        return len(self.dataset)773 774    def __getitem__(self, idx):775        img = load_image(self.dataset[idx]["image"].convert("RGB"), max_num=6, decoded=True).to(torch.bfloat16)776 777        question = self.dataset[idx]["question"]778        hint = self.dataset[idx]["hint"]779        question_id = self.dataset[idx]["index"]780        choices = self.dataset[idx]["choices"]781        answer = self.dataset[idx]["answer"]782 783        start_chr = 'A'784        choices_str = ''785        index2ans = {}786        all_choices = []787        for choice in choices:788            all_choices.append(start_chr)789            index2ans[start_chr] = choice790            choices_str += f"{start_chr}. {choice}\n"791            start_chr = chr(ord(start_chr) + 1)792 793        question = question + '\n' + choices_str794 795        return img, question_id, question, answer, str(index2ans), str(all_choices), self.dataset[idx]["question"]796 797 798def get_task_dataloader(task_name, task_cfg, args):799    if "subset" in task_cfg.keys():800        subset = task_cfg["subset"]801    else:802        subset = None803 804    if task_name == "coco_caption":805        dataset = COCOEvalDataset(args, task_cfg["image_dir"], subset)806    elif task_name == "flickr30k_caption":807        dataset = Flickr30KEvalDataset(args, task_cfg["image_dir"], subset)808    elif task_name == "vqav2":809        dataset = VQAv2EvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)810    elif task_name == "textvqa":811        dataset = TextVQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)812    elif task_name == "gqa":813        dataset = GQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)814    elif task_name == "chartqa":815        dataset = ChartQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)816    elif task_name == "okvqa":817        dataset = OKVQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], task_cfg["question_path"], subset)818    elif task_name == "vizwiz":819        dataset = VizWizEvalDataset(args, task_cfg["image_dir"], task_cfg["question_path"], subset)820    elif task_name == "docvqa":821        dataset = DocVQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], split='val', subset=subset)822    elif task_name == "docvqa_test":823        dataset = DocVQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], split='test', subset=subset)824    elif task_name == "realworldqa":825        dataset = RealworldQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)826    elif task_name == "mmmu":827        dataset = MMMUEvalDataset(args, task_cfg, subset=args.subset, start_idx=args.start_idx)828    elif task_name == "mmmu_pro":829        dataset = MMMUProEvalDataset(args, task_cfg)830    elif task_name == "mathvista":831        dataset = MathVistaEvalDataset(args, task_cfg)832    elif task_name == "mmbench":833        dataset = MMBenchEvalDataset(args, task_cfg["gt_path"])834    elif task_name == 'ocrbench':835        dataset = OCRBenchEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)836    elif task_name == 'ai2diagram':837        dataset = AI2DiagramEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)838    elif task_name == 'ai2diagram_nomask':839        dataset = AI2DiagramNoMaskEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)840    else:841        raise NotImplementedError(f"Task {task_name} is not supported yet.")842 843    dataloader = DataLoader(844        dataset,845        batch_size=1,846        shuffle=False,847        pin_memory=True,848    )849 850    return dataloader851