Add benchmark evaluation scripts

README.md

@@ -318,6 +318,22 @@ response = model.chat(tokenizer, pixel_values, question, generation_config)
 print(f'User: {question}\nAssistant: {response}')
 ```
 
+### Benchmark Evaluation
+
+To test our NVLM-1.0 model on the benchmark datasets, you can use the following code:
+
+```bash
+python run_eval.py --config-path eval/full_eval.yaml \
+ --result-save-path path/to/eval_results/ \
+ --zero-shot-eval-tasks chartqa coco_caption flickr30k_caption vqav2 mmmu textvqa mathvista mmbench chartqa docvqa realworldqa ocrbench ai2diagram ai2diagram_nomask mmmu_pro docvqa_test
+```
+
+Specifically,
+- `--config-path eval/full_eval.yaml` file contains the evaluation configurations, including  the evaluation prompt, the evaluation dataset paths, and generation hyper-parameters.
+- `--result-save-path path/to/eval_results/` specifies the path to save the evaluation results.
+- `--zero-shot-eval-tasks` specifies the tasks to evaluate on.
+
+
 ## Software Integration
 **Runtime Engine(s)** 
 * PyTorch <br>

eval/conversation.py

@@ -0,0 +1,492 @@
+# From https://github.com/haotian-liu/LLaVA/blob/main/llava/conversation.py
+
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    real_sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            if 'mmtag' in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0: message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+                    msg, image, image_process_mode = msg
+                    if image_process_mode == "Pad":
+                        def expand2square(pil_img, background_color=(122, 116, 104)):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                        image = expand2square(image)
+                    elif image_process_mode in ["Default", "Crop"]:
+                        pass
+                    elif image_process_mode == "Resize":
+                        image = image.resize((336, 336))
+                    else:
+                        raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if longest_edge != max(image.size):
+                        if H > W:
+                            H, W = longest_edge, shortest_edge
+                        else:
+                            H, W = shortest_edge, longest_edge
+                        image = image.resize((W, H))
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        buffered = BytesIO()
+                        image.save(buffered, format="PNG")
+                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    msg, image, image_process_mode = msg
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if H > W:
+                        H, W = longest_edge, shortest_edge
+                    else:
+                        H, W = shortest_edge, longest_edge
+                    image = image.resize((W, H))
+                    buffered = BytesIO()
+                    image.save(buffered, format="JPEG")
+                    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                    img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = img_str + msg.replace('<image>', '').strip()
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            real_sep2=self.real_sep2,
+            version=self.version)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+                "real_sep2": self.real_sep2
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+            "real_sep2": self.real_sep2
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
+        ("Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+### Used for llava-instruction-tuning stage
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llava_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+           "You are able to understand the visual content that the user provides, "
+           "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+
+
+### Used for llava-pretraining 
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_llava_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+
+nvllm_8b_pretrain = Conversation(
+    system="",
+    roles=(),
+    version="nvllm_8b",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="\n",
+)
+
+nvllm_8b_sft = Conversation(
+    system="System: This is a chat between a user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("User", "Assistant"),
+    version="nvllm_8b",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep="\n\n",
+    sep2="\n\n\n",
+    real_sep2="\n\n"
+)
+
+chatqa_sft = Conversation(
+    system="System: This is a chat between a user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("User", "Assistant"),
+    version="chatqa",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep="\n\n",
+    sep2="\n\n",
+    real_sep2="\n\n"
+)
+
+nvllm_8b_sft_noinstruction = Conversation(
+    system="",
+    roles=("User", "Assistant"),
+    version="nvllm_8b",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep="\n\n",
+    sep2="\n\n\n",
+    real_sep2="\n\n"
+)
+
+# conv_yi = Conversation(
+#     system="""<|im_start|>system
+# A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+#     roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+#     version="mpt",
+#     messages=(),
+#     offset=0,
+#     sep_style=SeparatorStyle.MPT,
+#     sep="<|im_end|>",
+# )
+
+conv_chatml = Conversation(
+    system="""<|im_start|>system
+Answer the questions.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+llama3_instruct = Conversation(
+    system="<|start_header_id|>system<|end_header_id|>\n\nAnswer the questions.",
+    roles=("<|start_header_id|>user<|end_header_id|>\n\n", "<|start_header_id|>assistant<|end_header_id|>\n\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|eot_id|>",
+)
+
+llama3_1_instruct = Conversation(
+    system="<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nAnswer the questions.",
+    roles=("<|start_header_id|>user<|end_header_id|>\n\n", "<|start_header_id|>assistant<|end_header_id|>\n\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|eot_id|>",
+)
+
+# default_conversation = conv_vicuna_v0
+default_conversation = nvllm_8b_sft
+
+# original_llava_pretraining = conv_llava_plain
+# original_llava_sft = conv_vicuna_v1
+
+conv_templates = {
+    "default": conv_vicuna_v0,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+
+    "plain": conv_llava_plain,
+    "v0_plain": conv_llava_plain,
+    "llava_v0": conv_llava_v0,
+    "v0_mmtag": conv_llava_v0_mmtag,
+    "llava_v1": conv_llava_v1,
+    "v1_mmtag": conv_llava_v1_mmtag,
+    "llava_llama_2": conv_llava_llama_2,
+
+    "mpt": conv_mpt,
+}
+
+
+if __name__ == "__main__":
+
+    print(default_conversation)
+    print(default_conversation.roles[0])
+
+
+    # print(default_conversation.get_prompt())

eval/eval_dataset.py

@@ -0,0 +1,850 @@
+import json
+import os
+import sys
+import time
+import yaml
+import spacy
+import ast
+from PIL import Image
+from glob import glob
+from tqdm import tqdm
+from collections import defaultdict
+import pandas as pd
+from io import BytesIO
+import base64
+from anls import anls_score
+import torch
+from torch.utils.data import Dataset, DataLoader, DistributedSampler
+import torchvision.transforms as T
+from eval import conversation as conversation_lib
+from eval.mmmu_utils import CAT_SHORT2LONG, DOMAIN_CAT2SUB_CAT, parse_multi_choice_response, parse_open_response, \
+    process_single_sample, construct_prompt, mmmu_main_eval, process_single_sample_pro, construct_prompt_pro
+from eval.mmmu_utils import evaluate as evaluate_mmmu
+from torchvision.transforms.functional import InterpolationMode
+from datasets import load_dataset, concatenate_datasets
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def build_transform(input_size):
+    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
+    transform = T.Compose([
+        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+        T.ToTensor(),
+        T.Normalize(mean=MEAN, std=STD)
+    ])
+    return transform
+
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_ratio_diff = float('inf')
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+
+def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+
+    # calculate the existing image aspect ratio
+    target_ratios = set(
+        (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) if
+        i * j <= max_num and i * j >= min_num)
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(
+        aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    processed_images = []
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+    assert len(processed_images) == blocks
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+    return processed_images
+
+
+def load_image(image, input_size=448, max_num=6, decoded=False):
+    if not decoded:
+        image = Image.open(image).convert('RGB')
+    transform = build_transform(input_size=input_size)
+    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
+    pixel_values = [transform(image) for image in images]
+    pixel_values = torch.stack(pixel_values)
+    return pixel_values
+
+
+def levenshtein_distance(s1, s2):
+    if len(s1) > len(s2):
+        s1, s2 = s2, s1
+
+    distances = range(len(s1) + 1)
+    for i2, c2 in enumerate(s2):
+        distances_ = [i2 + 1]
+        for i1, c1 in enumerate(s1):
+            if c1 == c2:
+                distances_.append(distances[i1])
+            else:
+                distances_.append(1 + min((distances[i1], distances[i1 + 1], distances_[-1])))
+        distances = distances_
+    return distances[-1]
+
+
+def get_anls_score(pred, gold_labels, threshold, llava_eval=False):
+    values = []
+    for answer in gold_labels:
+        # preprocess both the answers - gt and prediction
+        gt_answer = ' '.join(answer.strip().lower().split())
+        det_answer = ' '.join(pred.strip().lower().split())
+
+        dist = levenshtein_distance(gt_answer, det_answer)
+        length = max(len(answer.upper()), len(pred.upper()))
+        values.append(0.0 if length == 0 else float(dist) / float(length))
+
+    question_result = 1 - min(values)
+
+    if llava_eval:
+        question_result = 1.0 if question_result >= threshold else 0.0
+    else:
+        if (question_result < threshold):
+            question_result = 0
+
+    return question_result
+
+
+def isNumber(n: str):
+    try:
+        float(n)
+        return True
+    except ValueError:
+        return False
+
+
+class COCOEvalDataset(Dataset):
+    def __init__(self, args, img_dir, subset=None):
+        self.args = args
+        self.img_files = sorted(glob(os.path.join(img_dir, "*")))
+
+        if subset:
+            self.img_files = self.img_files[:subset]
+
+        self.image_ids = [int(img_file.split("_")[-1].split(".")[0]) for img_file in self.img_files]
+
+    def __len__(self):
+        return len(self.img_files)
+
+    def __getitem__(self, idx):
+        img_path = self.img_files[idx]
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        return self.image_ids[idx], img
+
+
+class Flickr30KEvalDataset(Dataset):
+    def __init__(self, args, img_dir, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.test_samples = json.load(open(os.path.join(img_dir, "flickr30k_test.json"), encoding='utf-8'))
+
+        if subset:
+            self.test_samples = self.test_samples[:subset]
+
+    def __len__(self):
+        return len(self.test_samples)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.test_samples[idx]["image"])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        image_id = int(self.test_samples[idx]["image"].split("/")[-1].replace(".jpg", ""))
+
+        return image_id, img
+
+
+class VQAv2EvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.gt = json.load(open(gt_path, encoding='utf-8'))
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.gt[idx]["image"])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = self.gt[idx]["question_id"]
+        question = self.gt[idx]["question"]
+        answer = self.gt[idx]["answer"]
+
+        return img, question_id, question, answer
+
+
+class TextVQAEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.gt = json.load(open(gt_path, encoding='utf-8'))['data']
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.gt[idx]["image_id"] + '.jpg')
+        if not os.path.exists(img_path):
+            img_path = img_path.replace('.jpg', '.png')
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = self.gt[idx]["question_id"]
+        question = self.gt[idx]["question"]
+        answer = self.gt[idx]["answers"]
+
+        return img, question_id, question, answer
+
+
+class GQAEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.gt = json.load(open(gt_path, encoding='utf-8'))
+        self.gt = [{
+            "question_id": int(k),
+            "image": v['imageId'] + ".jpg",
+            "question": v['question'],
+            "answer": v['answer']
+        } for k, v in self.gt.items()]
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.gt[idx]["image"])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = self.gt[idx]["question_id"]
+        question = self.gt[idx]["question"]
+        answer = self.gt[idx]["answer"]
+
+        return img, question_id, question, [answer]
+
+
+class ChartQAEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.gt = json.load(open(gt_path, encoding='utf-8'))
+        for i in range(len(self.gt)):
+            self.gt[i]['question_id'] = i
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.gt[idx]["imgname"])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = self.gt[idx]["question_id"]
+        question = self.gt[idx]["query"]
+        answer = self.gt[idx]["label"]
+
+        return img, question_id, question, [answer]
+
+
+class OKVQAEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, question_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.gt = json.load(open(gt_path, encoding='utf-8'))['annotations']
+        self.questions = json.load(open(question_path, 'r'))['questions']
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+        qid2q = {q['question_id']: q['question'] for q in self.questions}
+
+        for ann in self.gt:
+            ann['answers'] = [ans['answer'] for ans in ann['answers']]
+            ann['question'] = qid2q[ann['question_id']]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_id = str(self.gt[idx]["image_id"])
+        img_id = '0' * (12 - len(img_id)) + img_id
+        img_file_name = f"COCO_val2014_{img_id}.jpg"
+        img_path = os.path.join(self.img_dir, img_file_name)
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = self.gt[idx]["question_id"]
+        question = self.gt[idx]["question"]
+        answer = self.gt[idx]["answers"]
+
+        return img, question_id, question, answer
+
+
+class DocVQAEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, split='val', subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.gt = json.load(open(gt_path, encoding='utf-8'))['data']
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+        self.split = split
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.gt[idx]['image'].split('/')[-1])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = self.gt[idx]["questionId"]
+        question = self.gt[idx]["question"]
+
+        if self.split == 'val':
+            answer = self.gt[idx]["answers"]
+        else:
+            answer = ['']
+
+        return img, question_id, question, answer
+
+
+class OCRBenchEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.gt = json.load(open(gt_path, encoding='utf-8'))
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.gt[idx]['image_path'])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        dataset_name = self.gt[idx]["dataset_name"]
+        question_id = f"{idx}"
+        question = self.gt[idx]["question"]
+        answer = self.gt[idx]["answers"]
+        data_type = self.gt[idx]["type"]
+
+        return img, question_id, question, answer, dataset_name, data_type
+
+
+class AI2DiagramEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+
+        with open(gt_path, 'r') as json_file:
+            json_list = list(json_file)
+        self.gt = [json.loads(json_str) for json_str in json_list]
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.gt[idx]['image'])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = self.gt[idx]["question_id"]
+        question = self.gt[idx]["question"]
+        answer = self.gt[idx]["answer"]
+
+        return img, question_id, question, answer
+
+
+class AI2DiagramNoMaskEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+
+        with open(gt_path, 'r') as json_file:
+            json_list = list(json_file)
+        self.gt = [json.loads(json_str) for json_str in json_list]
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_file_name = self.gt[idx]['image'].replace("AI2D_TEST", "AI2D_TEST_NO_MASK_IMAGES")
+        img_path = os.path.join(self.img_dir, img_file_name)
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = self.gt[idx]["question_id"]
+        question = self.gt[idx]["question"]
+        answer = self.gt[idx]["answer"]
+
+        return img, question_id, question, answer
+
+
+class RealworldQAEvalDataset(Dataset):
+    def __init__(self, args, img_dir, gt_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.gt = json.load(open(gt_path, encoding='utf-8'))
+
+        if subset:
+            self.gt = self.gt[:subset]
+
+    def __len__(self):
+        return len(self.gt)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.gt[idx]['image'])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+
+        question_id = int(self.gt[idx]['image'].replace(".webp", ""))
+        question = self.gt[idx]["question"]
+
+        if self.gt[idx]['question_type'] == "multi-choice":
+            choices = self.gt[idx]["choices"]
+            start_chr = 'A'
+            choices_str = ''
+            index2ans = {}
+            all_choices = []
+            for choice in choices:
+                all_choices.append(start_chr)
+                index2ans[start_chr] = choice
+                choices_str += f"{start_chr}. {choice}\n"
+                start_chr = chr(ord(start_chr) + 1)
+
+            question = question + '\n' + choices_str
+            question = question + "Answer with the option's letter from the given choices directly."
+            answer = chr(ord('A') + self.gt[idx]['correct_choice_index'])
+        else:
+            question = question + "\nAnswer the question using a single word or phrase."
+            answer = self.gt[idx]['answer']
+
+        return img, question_id, question, [answer]
+
+
+class MathVistaEvalDataset(Dataset):
+    def __init__(self, args, task_cfg, gt_path=None):
+        self.args = args
+        self.task_cfg = task_cfg
+        self.dataset = load_dataset("AI4Math/MathVista")['testmini']
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        img = self.dataset[idx]['decoded_image']
+        img = load_image(img.convert("RGB"), max_num=6, decoded=True).to(torch.bfloat16)
+
+        question_id = self.dataset[idx]["pid"]
+        question = self.dataset[idx]["question"]
+        question_type = self.dataset[idx]["question_type"]  # free_form or multi_choice
+        query = self.dataset[idx]["query"]
+        choices = self.dataset[idx]["choices"]
+        answer = self.dataset[idx]["answer"]
+
+        if question_type == 'multi_choice':
+            start_chr = 'A'
+            choices_str = ''
+            index2ans = {}
+            all_choices = []
+            for choice in choices:
+                all_choices.append(start_chr)
+                index2ans[start_chr] = choice
+                choices_str += f"{start_chr}. {choice}\n"
+                start_chr = chr(ord(start_chr) + 1)
+
+            question = question + '\n' + choices_str
+            question = question + "Answer with the option's letter from the given choices directly."
+            answer = chr(ord('A') + choices.index(answer))
+        else:
+            question = query.replace("Hint: ", "")
+            index2ans = {}
+            all_choices = []
+
+        return img, question_id, question_type, question, answer, str(index2ans), str(all_choices)
+
+
+def construct_prompt_for_fewshot(sample):
+    config = {
+        "task_instructions": "",
+        "multi_choice_example_format": "{}\n{}Answer with the option's letter from the given choices directly.",
+        "short_ans_example_format": "{}\nAnswer the question using a single word or phrase."
+    }
+
+    question = sample['question'].strip()
+
+
+    options = eval(sample['options'])
+    example = ""
+    if sample['question_type'] == 'multiple-choice':
+        start_chr = 'A'
+        prediction_range = []
+        index2ans = {}
+        for option in options:
+            prediction_range.append(start_chr)
+            example += f"({start_chr}) {option}\n"
+            index2ans[start_chr] = option
+            start_chr = chr(ord(start_chr) + 1)
+        empty_prompt_sample_structure = config['multi_choice_example_format']
+        empty_prompt = empty_prompt_sample_structure.format(question, example)
+        res_dict = {'type': 'multichoice'}
+        res_dict['index2ans'] = index2ans
+        res_dict['correct_choice'] = sample['answer']
+        res_dict['all_choices'] = prediction_range
+        res_dict['empty_prompt'] = empty_prompt
+        if config['task_instructions']:
+            res_dict['final_input_prompt'] = config['task_instructions'].strip() + '\n\n' + empty_prompt
+        else:
+            res_dict['final_input_prompt'] = empty_prompt
+
+        res_dict['gt_content'] = options[ord(sample['answer'].upper()) - ord('A')]
+    else:
+        empty_prompt_sample_structure = config['short_ans_example_format']
+        empty_prompt = empty_prompt_sample_structure.format(question)
+        res_dict = {'type': 'open'}
+        res_dict['empty_prompt'] = empty_prompt
+        if config['task_instructions']:
+            res_dict['final_input_prompt'] = config['task_instructions'].strip() + '\n\n' + empty_prompt
+        else:
+            res_dict['final_input_prompt'] = empty_prompt
+        res_dict['gt_content'] = sample['answer']
+
+    res_dict.update(sample)
+    return res_dict
+
+
+def process_image_tag(q):
+    q = q.strip()
+
+    # heuristic way of removing <image 1>
+    if q == '<image 1>':
+        q = 'Answer the question in the image.'
+    elif ':<image 1>' in q:
+        q = q.replace(':<image 1>', ' in the image. ')
+        q = q.strip()
+    elif ': <image 1>' in q:
+        q = q.replace(': <image 1>', ' in the image. ')
+        q = q.strip()
+    elif '.<image 1>' in q or '. <image 1>' in q:
+        q_list = q.split('<image 1>')
+        q_list = [part.strip() for part in q_list if part.strip() != '']
+        q = ' '.join(q_list)
+    elif q.startswith('<image 1> '):
+        if q[10].isupper():
+            q = q.replace('<image 1>', '')
+        else:
+            q = q.replace('<image 1>', 'The image')
+        q = q.strip()
+    elif q.startswith('<image 1>'):
+        q = q.replace('<image 1>', '')
+    elif q.endswith('<image 1>?'):
+        q = q.replace('<image 1>', 'the image')
+    elif q.endswith('?<image 1>') or q.endswith('? <image 1>') or q.endswith('\n<image 1>'):
+        q = q.replace('<image 1>', '')
+        q = q.strip()
+    elif ' <image 1> ' in q:
+        q = q.replace('<image 1>', 'the image')
+    elif ' <image 1>' in q:
+        q = q.replace('<image 1>', 'the image')
+    elif '()<image 1>' in q:
+        q = q.replace('()<image 1>', '')
+    elif '(<image 1>)' in q:
+        q = q.replace('(<image 1>)', '')
+    elif '<image 1>.' in q:
+        q = q.replace("<image 1>.", ". ")
+    else:
+        q = q.replace("<image 1>", ". ")
+        q = q.strip()
+
+    # remove <image 2> to <image 8>
+    for i in range(2, 8):
+        q = q.replace(f"<image {i}>", "")
+
+    return q
+
+
+class MMMUProEvalDataset(Dataset):
+    def __init__(self, args, task_cfg, subset=None):
+        self.args = args
+        self.task_cfg = task_cfg
+        sub_dataset_list = []
+        # load_dataset will throw error if split is 'dev'
+        # 'dev' is part of the 'validation' and we need to manually split them
+
+        MMMU_path = "MMMU/MMMU_Pro"
+
+        _split = "test"
+
+        self.dataset = load_dataset(MMMU_path, "standard", split=_split)
+        if subset:
+            self.dataset = self.dataset[:subset]
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        # ===== single-image =====
+        sample = self.dataset[idx]
+        sample = process_single_sample_pro(sample)
+        sample = construct_prompt_pro(sample, self.task_cfg)
+        img = load_image(sample['image'].convert("RGB"), max_num=6, decoded=True).to(torch.bfloat16)
+
+        # img = img.reshape(-1, 3, self.args.img_h, self.args.img_w)
+
+        question_id = sample['id']
+        question = sample['final_input_prompt']
+        answer = sample['answer']
+
+        question = process_image_tag(question)
+        question = self.task_cfg['default_image_token'] + '\n' + question
+
+        if sample['question_type'] == 'multiple-choice':
+            index2ans = sample['index2ans']
+            all_choices = sample['all_choices']
+        else:
+            index2ans = {}
+            all_choices = []
+
+        return img, question_id, sample['subfield'], sample['question_type'], question, answer, str(index2ans), str \
+            (all_choices)
+
+
+class MMMUEvalDataset(Dataset):
+    def __init__(self, args, task_cfg, subset=None, start_idx=None):
+        self.args = args
+        self.task_cfg = task_cfg
+        sub_dataset_list = []
+        # load_dataset will throw error if split is 'dev'
+        # 'dev' is part of the 'validation' and we need to manually split them
+
+        MMMU_path = "MMMU/MMMU"
+
+        _split = "test" if task_cfg["split"] == "test" else "validation"
+        for subject in CAT_SHORT2LONG.values():
+            sub_dataset = load_dataset(
+                MMMU_path, subject,
+                split=_split,
+            )
+            sub_dataset_list.append(sub_dataset)
+
+        dataset = concatenate_datasets(sub_dataset_list)
+
+        if task_cfg["split"] != "test":
+            dataset = [s for s in dataset if s['id'].startswith(task_cfg["split"])]
+
+        # dataset = [s for s in dataset if s['image_2'] is not None][1:]
+
+        self.dataset = dataset
+
+        if subset:
+            self.dataset = [dataset[i] for i in range(start_idx, min(start_idx + subset, len(dataset)))]
+            print(f"Evaluating a subset of dataset: {len(self.dataset)} from {start_idx} to {start_idx + subset}")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        # ===== single-image =====
+        sample = self.dataset[idx]
+        sample = process_single_sample(sample)
+        sample = construct_prompt(sample, self.task_cfg)
+
+        img = load_image(sample['image'].convert("RGB"), max_num=6, decoded=True).to(torch.bfloat16)
+
+        question_id = sample['id']
+        question = sample['final_input_prompt']
+        answer = sample['answer']
+
+        question = process_image_tag(question)
+        question = self.task_cfg['default_image_token'] + '\n' + question
+
+
+        if sample['question_type'] == 'multiple-choice':
+            index2ans = sample['index2ans']
+            all_choices = sample['all_choices']
+        else:
+            index2ans = {}
+            all_choices = []
+
+        return img, question_id, sample['subfield'], sample['question_type'], question, answer, str(index2ans), str \
+            (all_choices)
+
+
+
+class VizWizEvalDataset(Dataset):
+    def __init__(self, args, img_dir, question_path, subset=None):
+        self.args = args
+        self.img_dir = img_dir
+        self.questions = json.load(open(question_path, encoding='utf-8'))
+
+    def __len__(self):
+        return len(self.questions)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_dir, self.questions[idx]["image"])
+        img = load_image(img_path, max_num=6).to(torch.bfloat16)
+        question = self.questions[idx]["question"]
+        question_id = self.questions[idx]["image"]
+
+        return img, question_id, question
+
+
+class MMBenchEvalDataset(Dataset):
+    def __init__(self, args, gt_path, subset=None):
+        self.args = args
+        df = pd.read_csv(gt_path, sep='\t')
+        self.dataset = []
+        for i, row in df.iterrows():
+            choices = []
+            for choice in ['A', 'B', 'C', 'D']:
+                if str(row[choice]) != 'nan':
+                    choices.append(row[choice])
+
+            this_sample = {
+                'index': row['index'],
+                'question': row['question'],
+                'hint': row['hint'],
+                'category': row['category'],
+                'image': Image.open(BytesIO(base64.b64decode(row['image']))),
+                'choices': choices
+            }
+
+            # Only dev set gives the ground truth answer
+            if 'answer' in row.keys():
+                this_sample['answer'] = row['answer']
+            else:
+                this_sample['answer'] = ''
+
+            self.dataset.append(this_sample)
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        img = load_image(self.dataset[idx]["image"].convert("RGB"), max_num=6, decoded=True).to(torch.bfloat16)
+
+        question = self.dataset[idx]["question"]
+        hint = self.dataset[idx]["hint"]
+        question_id = self.dataset[idx]["index"]
+        choices = self.dataset[idx]["choices"]
+        answer = self.dataset[idx]["answer"]
+
+        start_chr = 'A'
+        choices_str = ''
+        index2ans = {}
+        all_choices = []
+        for choice in choices:
+            all_choices.append(start_chr)
+            index2ans[start_chr] = choice
+            choices_str += f"{start_chr}. {choice}\n"
+            start_chr = chr(ord(start_chr) + 1)
+
+        question = question + '\n' + choices_str
+
+        return img, question_id, question, answer, str(index2ans), str(all_choices), self.dataset[idx]["question"]
+
+
+def get_task_dataloader(task_name, task_cfg, args):
+    if "subset" in task_cfg.keys():
+        subset = task_cfg["subset"]
+    else:
+        subset = None
+
+    if task_name == "coco_caption":
+        dataset = COCOEvalDataset(args, task_cfg["image_dir"], subset)
+    elif task_name == "flickr30k_caption":
+        dataset = Flickr30KEvalDataset(args, task_cfg["image_dir"], subset)
+    elif task_name == "vqav2":
+        dataset = VQAv2EvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)
+    elif task_name == "textvqa":
+        dataset = TextVQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)
+    elif task_name == "gqa":
+        dataset = GQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)
+    elif task_name == "chartqa":
+        dataset = ChartQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)
+    elif task_name == "okvqa":
+        dataset = OKVQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], task_cfg["question_path"], subset)
+    elif task_name == "vizwiz":
+        dataset = VizWizEvalDataset(args, task_cfg["image_dir"], task_cfg["question_path"], subset)
+    elif task_name == "docvqa":
+        dataset = DocVQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], split='val', subset=subset)
+    elif task_name == "docvqa_test":
+        dataset = DocVQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], split='test', subset=subset)
+    elif task_name == "realworldqa":
+        dataset = RealworldQAEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)
+    elif task_name == "mmmu":
+        dataset = MMMUEvalDataset(args, task_cfg, subset=args.subset, start_idx=args.start_idx)
+    elif task_name == "mmmu_pro":
+        dataset = MMMUProEvalDataset(args, task_cfg)
+    elif task_name == "mathvista":
+        dataset = MathVistaEvalDataset(args, task_cfg)
+    elif task_name == "mmbench":
+        dataset = MMBenchEvalDataset(args, task_cfg["gt_path"])
+    elif task_name == 'ocrbench':
+        dataset = OCRBenchEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)
+    elif task_name == 'ai2diagram':
+        dataset = AI2DiagramEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)
+    elif task_name == 'ai2diagram_nomask':
+        dataset = AI2DiagramNoMaskEvalDataset(args, task_cfg["image_dir"], task_cfg["gt_path"], subset)
+    else:
+        raise NotImplementedError(f"Task {task_name} is not supported yet.")
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        pin_memory=True,
+    )
+
+    return dataloader

eval/full_eval.yaml

@@ -0,0 +1,188 @@
+datasets:
+  coco_caption:
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\nGive a brief description of this image in one sentence.<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    output_max_len: 30
+    top_k: 3
+    temperature: 1.0
+
+  flickr30k_caption:
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\nGive a brief description of this image in one sentence.<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    output_max_len: 30
+    top_k: 3
+    temperature: 1.0
+
+  vqav2:
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}\nAnswer the question using a single word or phrase.<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 8
+    temperature: 1.0
+
+  mmmu:
+    split: "validation"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 1024
+    temperature: 1.0
+    apply_lemmatizer: False
+    task_instructions: ""
+    multi_choice_example_format: "{}\n{}\nAnswer with the option's letter from the given choices directly."
+    short_ans_example_format: "{}\nAnswer the question using a single word or phrase."
+    use_chat_format: True
+    conv_format: "yi_nous_sft"
+    default_image_token: "<image>"
+    prompt_offset: 4
+    answer_dict: "path/to/answer_dict_val.json"
+
+  textvqa:
+    split: "val"
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}\nAnswer the question using a single word, phrase, or number.<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 10
+    temperature: 1.0
+
+  mathvista:
+    split: "testmini"
+    prompt: "<|im_start|>system\nYou are math expert. Use your math knowledge to calculate the answer.<|im_end|><|im_start|>user\n<image>\n{}\nAnswer the question using a single word, phrase, or number.<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 1024
+    temperature: 1.0
+
+  mmbench:
+    split: "dev"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}Answer with the option's letter from the given choices directly.<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 10
+    temperature: 1.0
+    submission: False
+
+  chartqa:
+    split: "test"
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}<|im_end|><|im_start|>assistant\n"
+    
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 20
+    temperature: 1.0
+
+  docvqa:
+    split: "val"
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 20
+    temperature: 1.0
+
+  realworldqa:
+    split: "test"
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 20
+    temperature: 1.0
+    submission: False
+
+  ocrbench:
+    split: "test"
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 70
+    temperature: 1.0
+    submission: False
+
+  ai2diagram:
+    split: "test"
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}\nAnswer the question using a single word, phrase, or number.<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 20
+    temperature: 1.0
+
+  ai2diagram_nomask:
+    split: "test"
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|><|im_start|>user\n<image>\n{}\nAnswer the question using a single word, phrase, or number.<|im_end|><|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 20
+    temperature: 1.0
+
+  mmmu_pro:
+    split: "validation"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 10
+    temperature: 1.0
+    apply_lemmatizer: False
+    task_instructions: ""
+    multi_choice_example_format: "{}\n{}\nAnswer with the option's letter from the given choices directly."
+    short_ans_example_format: "{}\nAnswer the question using a single word or phrase."
+    use_chat_format: True
+    conv_format: "yi_nous_sft"
+    default_image_token: "<image>"
+    prompt_offset: 4
+    answer_dict: "path/to/answer_dict.json"
+
+  docvqa_test:
+    split: "test"
+    image_dir: "path/to/image"
+    gt_path: "path/to/ground_truth"
+    prompt: "<|im_start|>system\nFollow the user's instruction and answer questions.<|im_end|>\n<|im_start|>user\n<image>\n{}\nAnswer this question using the text in the image directly.<|im_end|>\n<|im_start|>assistant\n"
+    beam_search: True
+    beam_size: 1
+    top_k: 1
+    top_p: 0.0
+    output_max_len: 20
+    temperature: 1.0

eval/mmmu_utils.py

@@ -0,0 +1,663 @@
+# Adapted from https://github.com/MMMU-Benchmark/MMMU/blob/main/mmmu/utils/data_utils.py
+
+"""Utils for data load, save, and process (e.g., prompt construction)"""
+
+import os
+import json
+import yaml
+import re
+
+DOMAIN_CAT2SUB_CAT = {
+    'Art and Design': ['Art', 'Art_Theory', 'Design', 'Music'],
+    'Business': ['Accounting', 'Economics', 'Finance', 'Manage', 'Marketing'],
+    'Science': ['Biology', 'Chemistry', 'Geography', 'Math', 'Physics', ],
+    'Health and Medicine': ['Basic_Medical_Science', 'Clinical_Medicine', 'Diagnostics_and_Laboratory_Medicine',
+                            'Pharmacy', 'Public_Health'],
+    'Humanities and Social Science': ['History', 'Literature', 'Sociology', 'Psychology'],
+    'Tech and Engineering': ['Agriculture', 'Architecture_and_Engineering', 'Computer_Science', 'Electronics',
+                             'Energy_and_Power', 'Materials', 'Mechanical_Engineering'],
+}
+
+CAT_SHORT2LONG = {
+    'acc': 'Accounting',
+    'agri': 'Agriculture',
+    'arch': 'Architecture_and_Engineering',
+    'art': 'Art',
+    'art_theory': 'Art_Theory',
+    'bas_med': 'Basic_Medical_Science',
+    'bio': 'Biology',
+    'chem': 'Chemistry',
+    'cli_med': 'Clinical_Medicine',
+    'cs': 'Computer_Science',
+    'design': 'Design',
+    'diag_med': 'Diagnostics_and_Laboratory_Medicine',
+    'econ': 'Economics',
+    'elec': 'Electronics',
+    'ep': 'Energy_and_Power',
+    'fin': 'Finance',
+    'geo': 'Geography',
+    'his': 'History',
+    'liter': 'Literature',
+    'manage': 'Manage',
+    'mark': 'Marketing',
+    'mate': 'Materials',
+    'math': 'Math',
+    'mech': 'Mechanical_Engineering',
+    'music': 'Music',
+    'phar': 'Pharmacy',
+    'phys': 'Physics',
+    'psy': 'Psychology',
+    'pub_health': 'Public_Health',
+    'socio': 'Sociology'
+}
+
+
+# DATA SAVING
+def save_json(filename, ds):
+    with open(filename, 'w') as f:
+        json.dump(ds, f, indent=4)
+
+
+def get_multi_choice_info(options):
+    """
+    Given the list of options for multiple choice question
+    Return the index2ans and all_choices
+    """
+
+    start_chr = 'A'
+    all_choices = []
+    index2ans = {}
+    for i, option in enumerate(options):
+        index2ans[chr(ord(start_chr) + i)] = option
+        all_choices.append(chr(ord(start_chr) + i))
+
+    return index2ans, all_choices
+
+
+def load_yaml(file_path):
+    with open(file_path, 'r') as stream:
+        try:
+            yaml_dict = yaml.safe_load(stream)
+        except yaml.YAMLError as exc:
+            print(exc)
+
+    return yaml_dict
+
+
+def parse_img_path(text):
+    matches = re.findall("<img='(.*?)'>", text)
+    return matches
+
+
+def process_single_sample(data):
+    question = data['question']
+    o_imgs_paths = []
+    for option in data['options']:
+        current_o_imgs_paths = parse_img_path(option)
+        for img_path in current_o_imgs_paths:
+            o_imgs_paths.append(img_path)
+
+    if len(o_imgs_paths) > 1:  # multiple images in options, used for random selection
+        return {'id': data['id'], 'question': question, 'options': data['options'], 'answer': data['answer'],
+                'image': None, 'question_type': data['question_type'], 'subfield': data['subfield']}
+    else:
+        return {'id': data['id'], 'question': question, 'options': data['options'], 'answer': data['answer'],
+                'image': data['image_1'], 'question_type': data['question_type'], 'subfield': data['subfield']}
+
+
+def process_single_sample_pro(data):
+    question = data['question']
+    o_imgs_paths = []
+    for option in data['options']:
+        current_o_imgs_paths = parse_img_path(option)
+        for img_path in current_o_imgs_paths:
+            o_imgs_paths.append(img_path)
+
+    if len(o_imgs_paths) > 1:  # multiple images in options, used for random selection
+        return {'id': data['id'], 'question': question, 'options': data['options'], 'answer': data['answer'],
+                'image': None, 'question_type': 'multiple-choice', 'subfield': data['subject']}
+    else:
+        return {'id': data['id'], 'question': question, 'options': data['options'], 'answer': data['answer'],
+                'image': data['image_1'], 'question_type': 'multiple-choice', 'subfield': data['subject']}
+
+
+# DATA SAVING
+def save_json(filename, ds):
+    with open(filename, 'w') as f:
+        json.dump(ds, f, indent=4)
+
+
+def save_jsonl(filename, data):
+    """
+    Save a dictionary of data to a JSON Lines file with the filename as key and caption as value.
+
+    Args:
+        filename (str): The path to the file where the data should be saved.
+        data (dict): The dictionary containing the data to save where key is the image path and value is the caption.
+    """
+    with open(filename, 'w', encoding='utf-8') as f:
+        for img_path, caption in data.items():
+            # Extract the base filename without the extension
+            base_filename = os.path.basename(img_path)
+            # Create a JSON object with the filename as the key and caption as the value
+            json_record = json.dumps({base_filename: caption}, ensure_ascii=False)
+            # Write the JSON object to the file, one per line
+            f.write(json_record + '\n')
+
+
+def save_args(args, path_dir):
+    argsDict = args.__dict__
+    with open(path_dir + 'setting.txt', 'w') as f:
+        f.writelines('------------------ start ------------------' + '\n')
+        for eachArg, value in argsDict.items():
+            f.writelines(eachArg + ' : ' + str(value) + '\n')
+        f.writelines('------------------- end -------------------')
+
+
+# DATA PROCESSING
+def construct_prompt(sample, config):
+    question = sample['question'].strip()
+
+    # for i in range(8):
+    #     question = question.replace(f" <image {i}> ", " ")
+    #     question = question.replace(f" <image {i}>", " ")
+    #     question = question.replace(f"<image {i}> ", " ")
+    #     question = question.replace(f"<image {i}>", " ")
+    #     question = question.strip()
+
+    options = eval(sample['options'])
+    example = ""
+    if sample['question_type'] == 'multiple-choice':
+        start_chr = 'A'
+        prediction_range = []
+        index2ans = {}
+        for option in options:
+            prediction_range.append(start_chr)
+            example += f"({start_chr}) {option}\n"
+            # example += f"{start_chr}. {option}\n"
+            index2ans[start_chr] = option
+            start_chr = chr(ord(start_chr) + 1)
+        # example = example.rstrip()
+        empty_prompt_sample_structure = config['multi_choice_example_format']
+        empty_prompt = empty_prompt_sample_structure.format(question, example)
+        res_dict = {'type': 'multichoice'}
+        res_dict['index2ans'] = index2ans
+        res_dict['correct_choice'] = sample['answer']
+        res_dict['all_choices'] = prediction_range
+        res_dict['empty_prompt'] = empty_prompt
+        if config['task_instructions']:
+            res_dict['final_input_prompt'] = config['task_instructions'].strip() + '\n\n' + empty_prompt
+        else:
+            res_dict['final_input_prompt'] = empty_prompt
+
+        res_dict['gt_content'] = options[ord(sample['answer'].upper()) - ord('A')]
+    else:
+        empty_prompt_sample_structure = config['short_ans_example_format']
+        empty_prompt = empty_prompt_sample_structure.format(question)
+        res_dict = {'type': 'open'}
+        res_dict['empty_prompt'] = empty_prompt
+        if config['task_instructions']:
+            res_dict['final_input_prompt'] = config['task_instructions'].strip() + '\n\n' + empty_prompt
+        else:
+            res_dict['final_input_prompt'] = empty_prompt
+        res_dict['gt_content'] = sample['answer']
+
+    res_dict.update(sample)
+    return res_dict
+
+
+def construct_prompt_pro(sample, config):
+    question = sample['question'].strip()
+
+    # for i in range(8):
+    #     question = question.replace(f" <image {i}> ", " ")
+    #     question = question.replace(f" <image {i}>", " ")
+    #     question = question.replace(f"<image {i}> ", " ")
+    #     question = question.replace(f"<image {i}>", " ")
+    #     question = question.strip()
+
+    options = eval(sample['options'])
+
+    if len(options) == 1:
+        print("This is wrongly formated. We correct to options[0].")
+        options = options[0]
+
+    example = ""
+    if sample['question_type'] == 'multiple-choice':
+        start_chr = 'A'
+        prediction_range = []
+        index2ans = {}
+        for option in options:
+            prediction_range.append(start_chr)
+            example += f"({start_chr}) {option}\n"
+            # example += f"{start_chr}. {option}\n"
+            index2ans[start_chr] = option
+            start_chr = chr(ord(start_chr) + 1)
+        # example = example.rstrip()
+        empty_prompt_sample_structure = config['multi_choice_example_format']
+        empty_prompt = empty_prompt_sample_structure.format(question, example)
+        res_dict = {'type': 'multichoice'}
+        res_dict['index2ans'] = index2ans
+        res_dict['correct_choice'] = sample['answer']
+        res_dict['all_choices'] = prediction_range
+        res_dict['empty_prompt'] = empty_prompt
+        if config['task_instructions']:
+            res_dict['final_input_prompt'] = config['task_instructions'].strip() + '\n\n' + empty_prompt
+        else:
+            res_dict['final_input_prompt'] = empty_prompt
+
+        res_dict['gt_content'] = options[ord(sample['answer'].upper()) - ord('A')]
+    else:
+        empty_prompt_sample_structure = config['short_ans_example_format']
+        empty_prompt = empty_prompt_sample_structure.format(question)
+        res_dict = {'type': 'open'}
+        res_dict['empty_prompt'] = empty_prompt
+        if config['task_instructions']:
+            res_dict['final_input_prompt'] = config['task_instructions'].strip() + '\n\n' + empty_prompt
+        else:
+            res_dict['final_input_prompt'] = empty_prompt
+        res_dict['gt_content'] = sample['answer']
+
+    res_dict.update(sample)
+    return res_dict
+
+"""Response Parsing and Evaluation for various models"""
+from typing import Dict
+
+import re
+import random
+
+import numpy as np
+
+
+# ----------- Process Multi-choice -------------
+def parse_multi_choice_response(response, all_choices, index2ans):
+    """
+    Parse the prediction from the generated response.
+    Return the predicted index e.g., A, B, C, D.
+    """
+    for char in [',', '.', '!', '?', ';', ':', "'"]:
+        response = response.strip(char)
+    response = " " + response + " "  # add space to avoid partial match
+
+    index_ans = True
+    ans_with_brack = False
+    candidates = []
+    for choice in all_choices:  # e.g., (A) (B) (C) (D) A) B) C) D)
+        if f'({choice})' in response or f'{choice})' in response:
+            candidates.append(choice)
+            ans_with_brack = True
+
+    if len(candidates) == 0:
+        for choice in all_choices:  # e.g., A B C D
+            if f' {choice} ' in response:
+                candidates.append(choice)
+
+    # if all above doesn't get candidates, check if the content is larger than 5 tokens and try to parse the example
+    if len(candidates) == 0 and len(response.split()) > 5:
+        for index, ans in index2ans.items():
+            if ans.lower() in response.lower():
+                candidates.append(index)
+                index_ans = False  # it's content ans.
+
+    if len(candidates) == 0:  # still not get answer, randomly choose one.
+        pred_index = all_choices[0]
+    elif len(candidates) > 1:
+        start_indexes = []
+        if index_ans:
+            if ans_with_brack:
+                for can in candidates:
+                    index = response.rfind(f'({can})')
+                    start_indexes.append(index)  # -1 will be ignored anyway
+                # start_indexes = [generated_response.index(f'({can})') for can in candidates]
+            else:
+                for can in candidates:
+                    index = response.rfind(f" {can} ")
+                    start_indexes.append(index)
+        else:
+            for can in candidates:
+                index = response.lower().rfind(index2ans[can].lower())
+                start_indexes.append(index)
+        # get the last one
+        pred_index = candidates[np.argmax(start_indexes)]
+    else:  # if only one candidate, use it.
+        pred_index = candidates[0]
+
+    return pred_index
+
+
+# ----------- Process Open -------------
+def check_is_number(string):
+    """
+    Check if the given string a number.
+    """
+    try:
+        float(string.replace(',', ''))
+        return True
+    except ValueError:
+        # check if there's comma inside
+        return False
+
+
+def normalize_str(string):
+    """
+    Normalize the str to lower case and make them float numbers if possible.
+    """
+    # check if characters in the string
+
+    # if number, numerize it.
+    string = string.strip()
+
+    is_number = check_is_number(string)
+
+    if is_number:
+        string = string.replace(',', '')
+        string = float(string)
+        # leave 2 decimal
+        string = round(string, 2)
+        return [string]
+    else:  # it's likely to be a string
+        # lower it
+        string = string.lower()
+        if len(string) == 1:
+            return [" " + string, string + " "]  # avoid trivial matches
+        return [string]
+
+
+def extract_numbers(string):
+    """
+    Exact all forms of numbers from a string with regex.
+    """
+    # Pattern for numbers with commas
+    pattern_commas = r'-?\b\d{1,3}(?:,\d{3})+\b'
+    # Pattern for scientific notation
+    pattern_scientific = r'-?\d+(?:\.\d+)?[eE][+-]?\d+'
+    # Pattern for simple numbers without commas
+    pattern_simple = r'-?(?:\d+\.\d+|\.\d+|\d+\b)(?![eE][+-]?\d+)(?![,\d])'
+
+    # Extract numbers with commas
+    numbers_with_commas = re.findall(pattern_commas, string)
+    # Extract numbers in scientific notation
+    numbers_scientific = re.findall(pattern_scientific, string)
+    # Extract simple numbers without commas
+    numbers_simple = re.findall(pattern_simple, string)
+
+    # Combine all extracted numbers
+    all_numbers = numbers_with_commas + numbers_scientific + numbers_simple
+    return all_numbers
+
+
+def parse_open_response(response):
+    """
+    Parse the prediction from the generated response.
+    Return a list of predicted strings or numbers.
+    """
+
+    # content = content.strip("\n").strip(".").strip(" ")
+    def get_key_subresponses(response):
+        key_responses = []
+        response = response.strip().strip(".").lower()
+        sub_responses = re.split(r'\.\s(?=[A-Z])|\n', response)
+        indicators_of_keys = ['could be ', 'so ', 'is ',
+                              'thus ', 'therefore ', 'final ', 'answer ', 'result ']
+        key_responses = []
+        for index, resp in enumerate(sub_responses):
+            # if last one, accept it's an equation (the entire response can be just one sentence with equation)
+            if index == len(sub_responses) - 1:
+                indicators_of_keys.extend(['='])
+            shortest_key_response = None  # the shortest response that may contain the answer (tail part of the response)
+            for indicator in indicators_of_keys:
+                if indicator in resp:
+                    if not shortest_key_response:
+                        shortest_key_response = resp.split(indicator)[-1].strip()
+                    else:
+                        if len(resp.split(indicator)[-1].strip()) < len(shortest_key_response):
+                            shortest_key_response = resp.split(indicator)[-1].strip()
+                    # key_responses.append(resp.split(indicator)[1].strip())
+
+            if shortest_key_response:
+                # and it's not trivial
+                if shortest_key_response.strip() not in [":", ",", ".", "!", "?", ";", ":", "'"]:
+                    key_responses.append(shortest_key_response)
+        if len(key_responses) == 0:  # did not found any
+            return [response]
+        return key_responses
+
+    # pdb.set_trace()
+    key_responses = get_key_subresponses(response)
+
+    pred_list = key_responses.copy()  # keep the original string response
+    for resp in key_responses:
+        pred_list.extend(extract_numbers(resp))
+
+    tmp_pred_list = []
+    for i in range(len(pred_list)):
+        tmp_pred_list.extend(normalize_str(pred_list[i]))
+    pred_list = tmp_pred_list
+
+    # remove duplicates
+    pred_list = list(set(pred_list))
+
+    return pred_list
+
+
+# ----------- Evaluation -------------
+
+def eval_multi_choice(gold_i, pred_i):
+    """
+    Evaluate a multiple choice instance.
+    """
+    correct = False
+    # only they are exactly the same, we consider it as correct
+    if isinstance(gold_i, list):
+        for answer in gold_i:
+            if answer == pred_i:
+                correct = True
+                break
+    else:  # gold_i is a string
+        if gold_i == pred_i:
+            correct = True
+    return correct
+
+
+def eval_open(gold_i, pred_i):
+    """
+    Evaluate an open question instance
+    """
+    correct = False
+    if isinstance(gold_i, list):
+        # use float to avoid trivial matches
+        norm_answers = []
+        for answer in gold_i:
+            norm_answers.extend(normalize_str(answer))
+    else:
+        norm_answers = normalize_str(gold_i)
+    for pred in pred_i:  # pred is already normalized in parse response phase
+        if isinstance(pred, str):  # if it's a string, then find if ans in the pred_i
+            for norm_ans in norm_answers:
+                # only see if the string answer in the string pred
+                if isinstance(norm_ans, str) and norm_ans in pred:
+                    if not correct:
+                        correct = True
+                    break
+        else:  # it's a float number
+            if pred in norm_answers:
+                if not correct:
+                    correct = True
+                break
+    return correct
+
+
+# ----------- Batch Evaluation -------------
+def evaluate(samples):
+    """
+    Batch evaluation for multiple choice and open questions.
+    """
+    pred_correct = 0
+    judge_dict = dict()
+    for sample in samples:
+        gold_i = sample['answer']
+        pred_i = sample['parsed_pred']
+        if sample['question_type'] == 'multiple-choice':
+            correct = eval_multi_choice(gold_i, pred_i)
+        else:  # open question
+            correct = eval_open(gold_i, pred_i)
+
+        if correct:
+            judge_dict[sample['id']] = 'Correct'
+            pred_correct += 1
+        else:
+            judge_dict[sample['id']] = 'Wrong'
+
+    if len(samples) == 0:
+        return {'acc': 0}
+    return judge_dict, {'acc': pred_correct / len(samples)}
+
+
+# ----------- Calculate Accuracy -------------
+def calculate_ins_level_acc(results: Dict):
+    """Calculate the instruction level accuracy for given Subject results"""
+    acc = 0
+    ins_num = 0
+    for cat_results in results.values():
+        acc += cat_results['acc'] * cat_results['num_example']
+        ins_num += cat_results['num_example']
+    if ins_num == 0:
+        return 0
+    return acc / ins_num
+
+
+def mmmu_main_eval(output_dict, task_cfg):
+    answer_dict = json.load(open(task_cfg["answer_dict"]))
+
+    # group by category
+    output_dict_w_cat = {}
+    for data_id, parsed_pred in output_dict.items():
+        category = "_".join(data_id.split("_")[1:-1])
+        if category not in output_dict_w_cat:
+            output_dict_w_cat.update({category: {}})
+        output_dict_w_cat[category].update({data_id: parsed_pred})
+
+    # group by category
+    answer_dict_w_cat = {}
+    for data_id, parsed_pred in answer_dict.items():
+        category = "_".join(data_id.split("_")[1:-1])
+        if category not in answer_dict_w_cat:
+            answer_dict_w_cat.update({category: {}})
+        answer_dict_w_cat[category].update({data_id: parsed_pred})
+
+    evaluation_result = {}
+
+    for category in CAT_SHORT2LONG.values():
+        # print("Evaluating: {}".format(category))
+        # get cat_outputs and cat_answers
+        try:
+            cat_outputs = output_dict_w_cat[category]
+            cat_answers = answer_dict_w_cat[category]
+        except KeyError:
+            print("Skipping {} for not found".format(category))
+            continue
+
+        exampels_to_eval = []
+        for data_id, parsed_pred in cat_outputs.items():
+            question_type = cat_answers[data_id]['question_type']
+            if question_type != 'multiple-choice':
+                parsed_pred = parse_open_response(parsed_pred)  # mainly for type consistency (make it number, etc.)
+            else:
+                parsed_pred = parsed_pred
+
+            exampels_to_eval.append({
+                "id": data_id,
+                "question_type": question_type,
+                "answer": cat_answers[data_id]['ground_truth'],
+                "parsed_pred": parsed_pred
+            })
+
+        judge_dict, metric_dict = evaluate(exampels_to_eval)
+        metric_dict.update({"num_example": len(exampels_to_eval)})
+
+        evaluation_result[category] = metric_dict
+
+    printable_results = {}
+    # pdb.set_trace()
+    # add domain Subject
+    for domain, in_domain_cats in DOMAIN_CAT2SUB_CAT.items():
+        in_domain_cat_results = {}
+        for cat_name in in_domain_cats:  # use the order in DOMAIN_CAT2SUB_CAT
+            if cat_name in evaluation_result.keys():
+                in_domain_cat_results[cat_name] = evaluation_result[cat_name]
+            else:
+                pass
+        in_domain_ins_acc = calculate_ins_level_acc(in_domain_cat_results)
+        in_domain_data_num = sum([cat_results['num_example'] for cat_results in in_domain_cat_results.values()])
+        printable_results['Overall-' + domain] = {"num": int(in_domain_data_num),
+                                                  "acc": round(in_domain_ins_acc, 4)
+                                                  }
+        # add sub category
+        for cat_name, cat_results in in_domain_cat_results.items():
+            printable_results[cat_name] = {"num": int(cat_results['num_example']),
+                                           "acc": round(cat_results['acc'], 4)
+                                           }
+
+    # table.append(["-----------------------------", "-----", "----"])
+    all_ins_acc = calculate_ins_level_acc(evaluation_result)
+    printable_results['Overall'] = {
+        "num": sum([cat_results['num_example'] for cat_results in evaluation_result.values()]),
+        "acc": round(all_ins_acc, 4)
+        }
+
+    # print(printable_results)
+    return printable_results
+
+
+if __name__ == '__main__':
+    # tasks = yaml.safe_load(open("eval_config/eval_mmmu_yi_oci.yaml"))['datasets']
+    tasks = yaml.safe_load(open("eval_config/eval_mmmu_yi.yaml"))['datasets']
+    print(tasks)
+
+    # with open("/lustre/fs4/portfolios/adlr/users/boxinw/llava-megatron-gen/checkpoints/test/eval_mmmu_iter500_merged.4node.json") as f:
+    with open("/lustre/fsw/portfolios/llmservice/users/boxinw/eval_mmmu_iter6000_merged.0.53.json") as f:
+        merged_results = json.load(f)
+
+    eval_samples = []
+    eval_output_dict = {}
+    for res in merged_results:
+        pred_ans = res["answer"].upper()
+        gt_ans = res['gt_answer']
+        if res['question_type'] == 'multiple-choice':
+            parsed_pred = parse_multi_choice_response(pred_ans, res['all_choices'], res['index2ans'])
+            if pred_ans != parsed_pred:
+                print(f"MC: Original: {pred_ans}, Parsed: {parsed_pred}")
+            eval_samples.append(
+                {
+                    'id': res['question_id'],
+                    'question_type': res['question_type'],
+                    'answer': res['gt_answer'],  # the content in option, not answer index.
+                    'response': pred_ans,
+                    'parsed_pred': parsed_pred,
+                    'index2ans': res['index2ans'],
+                }
+            )
+            eval_output_dict[res['question_id']] = parsed_pred
+        else:
+            parsed_pred = parse_open_response(pred_ans)
+            if pred_ans != parsed_pred:
+                print(f"Open: Original: {pred_ans}, Parsed: {parsed_pred}")
+            eval_samples.append(
+                {
+                    'id': res['question_id'],
+                    'question_type': res['question_type'],
+                    'answer': res['gt_answer'],
+                    'response': pred_ans,
+                    'parsed_pred': parsed_pred,
+                }
+            )
+            eval_output_dict[res['question_id']] = pred_ans
+
+    json.dump(eval_output_dict, open("validation_mmmu_iter6000_merged.0.53.sorted.json", "w"), indent=4, sort_keys=True)
+
+
+    x = mmmu_main_eval(eval_output_dict,
+                   task_cfg=tasks['mmmu'])
+
+    print(x)

eval/requirements.txt

@@ -0,0 +1,3 @@
+anls
+datasets
+pycocoevalcap

eval/vqa_utils.py

@@ -0,0 +1,317 @@
+# coding=utf-8
+
+__author__ = "aagrawal"
+
+# This code is based on the code written by Tsung-Yi Lin for MSCOCO Python API available at the following link:
+# (https://github.com/tylin/coco-caption/blob/master/pycocoevalcap/eval.py).
+import sys
+import re
+
+
+class VQAEval:
+    def __init__(self, vqa=None, vqaRes=None, n=2):
+        self.n = n
+        self.accuracy = {}
+        self.evalQA = {}
+        self.evalQuesType = {}
+        self.evalAnsType = {}
+        self.vqa = vqa
+        self.vqaRes = vqaRes
+        if vqa is not None:
+            self.params = {"question_id": vqa.getQuesIds()}
+        self.contractions = {
+            "aint": "ain't",
+            "arent": "aren't",
+            "cant": "can't",
+            "couldve": "could've",
+            "couldnt": "couldn't",
+            "couldn'tve": "couldn't've",
+            "couldnt've": "couldn't've",
+            "didnt": "didn't",
+            "doesnt": "doesn't",
+            "dont": "don't",
+            "hadnt": "hadn't",
+            "hadnt've": "hadn't've",
+            "hadn'tve": "hadn't've",
+            "hasnt": "hasn't",
+            "havent": "haven't",
+            "hed": "he'd",
+            "hed've": "he'd've",
+            "he'dve": "he'd've",
+            "hes": "he's",
+            "howd": "how'd",
+            "howll": "how'll",
+            "hows": "how's",
+            "Id've": "I'd've",
+            "I'dve": "I'd've",
+            "Im": "I'm",
+            "Ive": "I've",
+            "isnt": "isn't",
+            "itd": "it'd",
+            "itd've": "it'd've",
+            "it'dve": "it'd've",
+            "itll": "it'll",
+            "let's": "let's",
+            "maam": "ma'am",
+            "mightnt": "mightn't",
+            "mightnt've": "mightn't've",
+            "mightn'tve": "mightn't've",
+            "mightve": "might've",
+            "mustnt": "mustn't",
+            "mustve": "must've",
+            "neednt": "needn't",
+            "notve": "not've",
+            "oclock": "o'clock",
+            "oughtnt": "oughtn't",
+            "ow's'at": "'ow's'at",
+            "'ows'at": "'ow's'at",
+            "'ow'sat": "'ow's'at",
+            "shant": "shan't",
+            "shed've": "she'd've",
+            "she'dve": "she'd've",
+            "she's": "she's",
+            "shouldve": "should've",
+            "shouldnt": "shouldn't",
+            "shouldnt've": "shouldn't've",
+            "shouldn'tve": "shouldn't've",
+            "somebody'd": "somebodyd",
+            "somebodyd've": "somebody'd've",
+            "somebody'dve": "somebody'd've",
+            "somebodyll": "somebody'll",
+            "somebodys": "somebody's",
+            "someoned": "someone'd",
+            "someoned've": "someone'd've",
+            "someone'dve": "someone'd've",
+            "someonell": "someone'll",
+            "someones": "someone's",
+            "somethingd": "something'd",
+            "somethingd've": "something'd've",
+            "something'dve": "something'd've",
+            "somethingll": "something'll",
+            "thats": "that's",
+            "thered": "there'd",
+            "thered've": "there'd've",
+            "there'dve": "there'd've",
+            "therere": "there're",
+            "theres": "there's",
+            "theyd": "they'd",
+            "theyd've": "they'd've",
+            "they'dve": "they'd've",
+            "theyll": "they'll",
+            "theyre": "they're",
+            "theyve": "they've",
+            "twas": "'twas",
+            "wasnt": "wasn't",
+            "wed've": "we'd've",
+            "we'dve": "we'd've",
+            "weve": "we've",
+            "werent": "weren't",
+            "whatll": "what'll",
+            "whatre": "what're",
+            "whats": "what's",
+            "whatve": "what've",
+            "whens": "when's",
+            "whered": "where'd",
+            "wheres": "where's",
+            "whereve": "where've",
+            "whod": "who'd",
+            "whod've": "who'd've",
+            "who'dve": "who'd've",
+            "wholl": "who'll",
+            "whos": "who's",
+            "whove": "who've",
+            "whyll": "why'll",
+            "whyre": "why're",
+            "whys": "why's",
+            "wont": "won't",
+            "wouldve": "would've",
+            "wouldnt": "wouldn't",
+            "wouldnt've": "wouldn't've",
+            "wouldn'tve": "wouldn't've",
+            "yall": "y'all",
+            "yall'll": "y'all'll",
+            "y'allll": "y'all'll",
+            "yall'd've": "y'all'd've",
+            "y'alld've": "y'all'd've",
+            "y'all'dve": "y'all'd've",
+            "youd": "you'd",
+            "youd've": "you'd've",
+            "you'dve": "you'd've",
+            "youll": "you'll",
+            "youre": "you're",
+            "youve": "you've",
+        }
+        self.manualMap = {
+            "none": "0",
+            "zero": "0",
+            "one": "1",
+            "two": "2",
+            "three": "3",
+            "four": "4",
+            "five": "5",
+            "six": "6",
+            "seven": "7",
+            "eight": "8",
+            "nine": "9",
+            "ten": "10",
+        }
+        self.articles = ["a", "an", "the"]
+
+        self.periodStrip = re.compile("(?!<=\d)(\.)(?!\d)")
+        self.commaStrip = re.compile("(\d)(,)(\d)")
+        self.punct = [
+            ";",
+            r"/",
+            "[",
+            "]",
+            '"',
+            "{",
+            "}",
+            "(",
+            ")",
+            "=",
+            "+",
+            "\\",
+            "_",
+            "-",
+            ">",
+            "<",
+            "@",
+            "`",
+            ",",
+            "?",
+            "!",
+        ]
+
+    def evaluate(self, quesIds=None):
+        if quesIds == None:
+            quesIds = [quesId for quesId in self.params["question_id"]]
+        gts = {}
+        res = {}
+        for quesId in quesIds:
+            gts[quesId] = self.vqa.qa[quesId]
+            res[quesId] = self.vqaRes.qa[quesId]
+
+        # =================================================
+        # Compute accuracy
+        # =================================================
+        accQA = []
+        accQuesType = {}
+        accAnsType = {}
+        print("computing accuracy")
+        step = 0
+        for quesId in quesIds:
+            resAns = res[quesId]["answer"]
+            resAns = resAns.replace("\n", " ")
+            resAns = resAns.replace("\t", " ")
+            resAns = resAns.strip()
+            resAns = self.processPunctuation(resAns)
+            resAns = self.processDigitArticle(resAns)
+            gtAcc = []
+            gtAnswers = [ans["answer"] for ans in gts[quesId]["answers"]]
+            if len(set(gtAnswers)) > 1:
+                for ansDic in gts[quesId]["answers"]:
+                    ansDic["answer"] = self.processPunctuation(ansDic["answer"])
+            for gtAnsDatum in gts[quesId]["answers"]:
+                otherGTAns = [
+                    item for item in gts[quesId]["answers"] if item != gtAnsDatum
+                ]
+                matchingAns = [item for item in otherGTAns if item["answer"] == resAns]
+                acc = min(1, float(len(matchingAns)) / 3)
+                gtAcc.append(acc)
+            quesType = gts[quesId]["question_type"]
+            ansType = gts[quesId]["answer_type"]
+            avgGTAcc = float(sum(gtAcc)) / len(gtAcc)
+            accQA.append(avgGTAcc)
+            if quesType not in accQuesType:
+                accQuesType[quesType] = []
+            accQuesType[quesType].append(avgGTAcc)
+            if ansType not in accAnsType:
+                accAnsType[ansType] = []
+            accAnsType[ansType].append(avgGTAcc)
+            self.setEvalQA(quesId, avgGTAcc)
+            self.setEvalQuesType(quesId, quesType, avgGTAcc)
+            self.setEvalAnsType(quesId, ansType, avgGTAcc)
+            if step % 100 == 0:
+                self.updateProgress(step / float(len(quesIds)))
+            step = step + 1
+
+        self.setAccuracy(accQA, accQuesType, accAnsType)
+        print("Done computing accuracy")
+
+    def processPunctuation(self, inText):
+        outText = inText
+        for p in self.punct:
+            if (p + " " in inText or " " + p in inText) or (
+                re.search(self.commaStrip, inText) != None
+            ):
+                outText = outText.replace(p, "")
+            else:
+                outText = outText.replace(p, " ")
+        outText = self.periodStrip.sub("", outText, re.UNICODE)
+        return outText
+
+    def processDigitArticle(self, inText):
+        outText = []
+        tempText = inText.lower().split()
+        for word in tempText:
+            word = self.manualMap.setdefault(word, word)
+            if word not in self.articles:
+                outText.append(word)
+            else:
+                pass
+        for wordId, word in enumerate(outText):
+            if word in self.contractions:
+                outText[wordId] = self.contractions[word]
+        outText = " ".join(outText)
+        return outText
+
+    def setAccuracy(self, accQA, accQuesType, accAnsType):
+        self.accuracy["overall"] = round(100 * float(sum(accQA)) / len(accQA), self.n)
+        self.accuracy["perQuestionType"] = {
+            quesType: round(
+                100 * float(sum(accQuesType[quesType])) / len(accQuesType[quesType]),
+                self.n,
+            )
+            for quesType in accQuesType
+        }
+        self.accuracy["perAnswerType"] = {
+            ansType: round(
+                100 * float(sum(accAnsType[ansType])) / len(accAnsType[ansType]), self.n
+            )
+            for ansType in accAnsType
+        }
+
+    def setEvalQA(self, quesId, acc):
+        self.evalQA[quesId] = round(100 * acc, self.n)
+
+    def setEvalQuesType(self, quesId, quesType, acc):
+        if quesType not in self.evalQuesType:
+            self.evalQuesType[quesType] = {}
+        self.evalQuesType[quesType][quesId] = round(100 * acc, self.n)
+
+    def setEvalAnsType(self, quesId, ansType, acc):
+        if ansType not in self.evalAnsType:
+            self.evalAnsType[ansType] = {}
+        self.evalAnsType[ansType][quesId] = round(100 * acc, self.n)
+
+    def updateProgress(self, progress):
+        barLength = 20
+        status = ""
+        if isinstance(progress, int):
+            progress = float(progress)
+        if not isinstance(progress, float):
+            progress = 0
+            status = "error: progress var must be float\r\n"
+        if progress < 0:
+            progress = 0
+            status = "Halt...\r\n"
+        if progress >= 1:
+            progress = 1
+            status = "Done...\r\n"
+        block = int(round(barLength * progress))
+        text = "\rFinshed Percent: [{0}] {1}% {2}".format(
+            "#" * block + "-" * (barLength - block), int(progress * 100), status
+        )
+        sys.stdout.write(text)
+        sys.stdout.flush()

run_eval.py

@@ -0,0 +1,702 @@
+import ast
+import json
+
+import torch
+from transformers import AutoTokenizer, AutoModel
+import math
+import time
+import argparse
+import yaml
+import os
+
+from eval.eval_dataset import get_task_dataloader, isNumber, get_anls_score
+from eval.mmmu_utils import CAT_SHORT2LONG, DOMAIN_CAT2SUB_CAT, parse_multi_choice_response, parse_open_response, \
+    process_single_sample, construct_prompt, mmmu_main_eval, process_single_sample_pro, construct_prompt_pro
+from eval.mmmu_utils import evaluate as evaluate_mmmu
+from pycocotools.coco import COCO
+from pycocoevalcap.eval import COCOEvalCap
+from anls import anls_score
+import pandas as pd
+from eval.vqa_utils import VQAEval
+
+
+def split_model():
+    device_map = {}
+    world_size = torch.cuda.device_count()
+    num_layers = 80
+    # Since the first GPU will be used for ViT, treat it as half a GPU.
+    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
+    num_layers_per_gpu = [num_layers_per_gpu] * world_size
+    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
+    layer_cnt = 0
+    for i, num_layer in enumerate(num_layers_per_gpu):
+        for j in range(num_layer):
+            device_map[f'language_model.model.layers.{layer_cnt}'] = i
+            layer_cnt += 1
+    device_map['vision_model'] = 0
+    device_map['mlp1'] = 0
+    device_map['language_model.model.tok_embeddings'] = 0
+    device_map['language_model.model.embed_tokens'] = 0
+    device_map['language_model.output'] = 0
+    device_map['language_model.model.norm'] = 0
+    device_map['language_model.lm_head'] = 0
+    device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
+
+    return device_map
+
+
+def generate_task_results(
+        task_name,
+        task_cfg,
+        args,
+        model,
+        tokenizer,
+        generation_config,
+        dataloader,
+        results_save_dir,
+):
+    results = []
+
+    if "prompt" in task_cfg:
+        prompt = task_cfg["prompt"]
+
+    if task_name == "mmmu" or task_name == 'mmmu_pro':
+        for i, (img, question_id, subfield, question_type, question, answer, index2ans, all_choices) in enumerate(
+                dataloader):
+            if img.dim() == 5:
+                img = img.squeeze(0)
+            question_id = question_id[0]
+            question_type = question_type[0]
+            question = question[0]
+            answer = answer[0]
+            subfield = subfield[0]
+            index2ans = ast.literal_eval(index2ans[0])
+            all_choices = ast.literal_eval(all_choices[0])
+
+            this_prompt = question
+
+            if task_name == 'mmmu':
+                categories = list(CAT_SHORT2LONG.values())
+                new_sys_msg = "You are expert in {} ({}). Read the image and use your knowledge in {} ({}) to answer the question."
+                for c in categories:
+                    if c in question_id:
+                        cat = c.lower().replace('_', ' ')
+                        new_sys_msg = new_sys_msg.format(cat, subfield, cat, subfield)
+                        break
+            else:
+                new_sys_msg = f"You are expert in {subfield}. Read the image and use your knowledge in {subfield} to answer the question."
+            model.system_message = new_sys_msg
+
+            generated_answer = model.chat(tokenizer, img, question, generation_config)
+            print("generated_answer", generated_answer)
+
+            results.append({
+                "question_id": question_id,
+                "question_type": question_type,
+                "answer": generated_answer,
+                "gt_answer": answer,
+                "index2ans": index2ans,
+                "all_choices": all_choices
+            })
+            print("idx:", i, results[-1])
+
+
+    elif task_name == "coco_caption" or task_name == "flickr30k_caption":
+
+        for i, (image_id, img,) in enumerate(dataloader):
+
+            if img.dim() == 5:
+                img = img.squeeze(0)
+
+            generated_answer = model.chat_without_chat_prompt(tokenizer, img, prompt, generation_config)
+            print("generated_answer", generated_answer)
+
+            results.append({"image_id": image_id.item(), "caption": generated_answer})
+            print("idx:", i, results[-1])
+
+    elif task_name in ["vqav2", "gqa", "okvqa", "textvqa", "chartqa", "docvqa", "realworldqa",
+                       "ai2diagram", "ai2diagram_nomask", "docvqa_test"]:
+        for i, (img, question_id, question, answer,) in enumerate(dataloader):
+            if img.dim() == 5:
+                img = img.squeeze(0)
+
+            # Only works for batch size = 1
+            # need to implement collate function when we use bs > 1 in the future
+            question = question[0]
+            if task_name == 'ai2diagram' or task_name == 'ai2diagram_nomask':
+                question_id = question_id[0]
+            else:
+                question_id = question_id[0].item()
+
+            if type(answer) == list:
+                answer = [ans[0] for ans in answer]
+            else:
+                answer = [answer[0]]
+
+            # Need to change if using batch size > 1 in the future
+            this_prompt = prompt.format(question)
+
+            generated_answer = model.chat_without_chat_prompt(tokenizer, img, this_prompt, generation_config)
+            print("generated_answer", generated_answer)
+
+            results.append({"question_id": question_id, "answer": generated_answer, "gt_answer": answer})
+            print("idx:", i, results[-1])
+
+
+    elif task_name == "ocrbench":
+        for i, (img, question_id, question, answer, dataset_name, data_type,) in enumerate(
+                dataloader):
+            if img.dim() == 5:
+                img = img.squeeze(0)
+
+            question_id = question_id[0]
+            question = question[0]
+            answer = answer[0]
+            dataset_name = dataset_name[0]
+            data_type = data_type[0]
+
+            this_prompt = prompt.format(question)
+
+            generated_answer = model.chat_without_chat_prompt(tokenizer, img, this_prompt, generation_config)
+            print("generated_answer", generated_answer)
+
+            results.append({"question_id": question_id, "answer": generated_answer, "gt_answer": answer,
+                            "dataset_name": dataset_name, "type": data_type})
+            print("idx:", i, results[-1])
+
+
+    elif task_name == "mathvista":
+        for i, (
+                img, question_id, question_type, question, answer, index2ans, all_choices,) in enumerate(
+            dataloader):
+            if img.dim() == 5:
+                img = img.squeeze(0)
+
+            question_id = question_id[0]
+            question = question[0]
+            question_type = question_type[0]
+            answer = answer[0]
+
+            index2ans = ast.literal_eval(index2ans[0])
+            all_choices = ast.literal_eval(all_choices[0])
+
+            this_prompt = prompt.format(question)
+
+            generated_answer = model.chat_without_chat_prompt(tokenizer, img, this_prompt, generation_config)
+            print("generated_answer", generated_answer)
+
+            results.append({
+                "question_id": question_id,
+                "question_type": question_type,
+                "answer": generated_answer,
+                "gt_answer": answer,
+                "index2ans": index2ans,
+                "all_choices": all_choices
+            })
+            print("idx:", i, results[-1])
+
+
+    elif task_name == "mmbench":
+        for i, (
+                img, question_id, question, answer, index2ans, all_choices, original_q,) in enumerate(
+            dataloader):
+            if img.dim() == 5:
+                img = img.squeeze(0)
+
+            question_id = question_id[0]
+            question = question[0]
+            answer = answer[0]
+            original_q = original_q[0]
+
+            index2ans = ast.literal_eval(index2ans[0])
+            all_choices = ast.literal_eval(all_choices[0])
+
+            this_prompt = prompt.format(question)
+
+            generated_answer = model.chat_without_chat_prompt(tokenizer, img, this_prompt, generation_config)
+            print("generated_answer", generated_answer)
+
+            results.append({
+                "question_id": question_id.item(),
+                "question": original_q,
+                "answer": generated_answer,
+                "gt_answer": answer,
+                "index2ans": index2ans,
+                "all_choices": all_choices
+            })
+            print("idx:", i, results[-1])
+
+
+    elif task_name == "vizwiz":
+        for i, (img, question_id, question,) in enumerate(dataloader):
+            if img.dim() == 5:
+                img = img.squeeze(0)
+
+            question = question[0]
+            question_id = question_id[0]
+
+            this_prompt = prompt.format(question)
+
+            generated_answer = model.chat_without_chat_prompt(tokenizer, img, this_prompt, generation_config)
+            print("generated_answer", generated_answer)
+
+            results.append({"image": question_id, "answer": generated_answer})
+            print("idx:", i, results[-1])
+
+    else:
+        raise NotImplementedError(f"Task {task_name} is not supported yet.")
+
+    os.makedirs(results_save_dir, exist_ok=True)
+    if args.subset is not None:
+        results_save_path = os.path.join(results_save_dir,
+                                         f"eval_{task_name}_subset_{args.subset}_start_{args.start_idx}.json")
+    else:
+        results_save_path = os.path.join(results_save_dir,
+                                         f"eval_{task_name}.json")
+    print("Saving to ", results_save_path)
+    json.dump(results, open(results_save_path, 'w'))
+
+
+def calc_task_metrics(args, task_name, task_cfg, results_save_dir):
+    if args.subset is not None:
+        merged_results_path = os.path.join(results_save_dir,
+                                           f"eval_{task_name}_subset_{args.subset}_start_{args.start_idx}.json")
+    else:
+        merged_results_path = os.path.join(results_save_dir,
+                                           f"eval_{task_name}.json")
+    merged_results = json.load(open(merged_results_path, "r"))
+
+    if task_name == "coco_caption" or task_name == "flickr30k_caption":
+        # Calculate scores
+        coco = COCO(task_cfg["gt_path"])
+        coco_result = coco.loadRes(merged_results_path)
+        coco_eval = COCOEvalCap(coco, coco_result)
+        coco_eval.params["image_id"] = coco_result.getImgIds()
+        coco_eval.evaluate()
+
+        # Print and save scores
+        print(f"====== {task_name} scores ======")
+        with open(os.path.join(results_save_dir, "scores.txt"), "a") as f:
+            f.write(f"{task_name} scores:\n")
+            for k, v in coco_eval.eval.items():
+                msg = f"{k} = {v * 100:.3f}"
+                print(msg)
+                f.write(msg + "\n")
+            f.write("\n")
+        return coco_eval.eval["CIDEr"]
+    elif task_name == "vqav2" or task_name == "okvqa":
+        vqa_tool = VQAEval()
+        all_acc = []
+        for res in merged_results:
+            pred_ans = res["answer"]
+            pred_ans = vqa_tool.processPunctuation(pred_ans)
+            pred_ans = vqa_tool.processDigitArticle(pred_ans)
+
+            gt_ans = res["gt_answer"]
+            gt_ans = [vqa_tool.processPunctuation(ans) for ans in gt_ans]
+            gt_ans = [vqa_tool.processDigitArticle(ans) for ans in gt_ans]
+
+            num_match = sum([pred_ans == ans for ans in gt_ans])
+            acc = min(1.0, num_match / 3.0)
+            all_acc.append(acc)
+        acc_avg = sum(all_acc) / len(all_acc) * 100
+        print(f"===== {task_name} Accuracy {acc_avg:.2f}% =====")
+        with open(os.path.join(results_save_dir, "scores.txt"), "a") as f:
+            f.write(f"{task_name} Accuracy = {acc_avg:.2f}%\n\n")
+        return acc_avg
+
+    elif task_name == "textvqa":
+        vqa_tool = VQAEval()
+        all_acc = []
+        for res in merged_results:
+            pred_ans = res["answer"]
+            pred_ans = vqa_tool.processPunctuation(pred_ans)
+            pred_ans = vqa_tool.processDigitArticle(pred_ans)
+
+            gt_ans = res["gt_answer"]
+            gt_ans = [vqa_tool.processPunctuation(ans) for ans in gt_ans]
+            gt_ans = [vqa_tool.processDigitArticle(ans) for ans in gt_ans]
+
+            num_match = sum([pred_ans == ans for ans in gt_ans])
+            acc = min(1.0, num_match / 3.0)
+            all_acc.append(acc)
+        acc_avg = sum(all_acc) / len(all_acc) * 100
+        print(
+            f"===== {task_name} Accuracy {acc_avg:.2f}% (need to submit to EvalAI for the accurate accuracy) =====")
+        with open(os.path.abspath(os.path.join(results_save_dir, "scores.txt")), "a") as f:
+            f.write(
+                f"{task_name} Accuracy = {acc_avg:.2f}% (need to submit to EvalAI for the accurate accuracy)\n\n")
+        return acc_avg
+
+    elif task_name in ["gqa", "realworldqa", "ai2diagram", "ai2diagram_nomask"]:
+        vqa_tool = VQAEval()
+        acc = 0
+        for res in merged_results:
+            pred_ans = res["answer"]
+            pred_ans = vqa_tool.processPunctuation(pred_ans)
+            pred_ans = vqa_tool.processDigitArticle(pred_ans)
+
+            gt_ans = res["gt_answer"][0]
+            gt_ans = vqa_tool.processPunctuation(gt_ans)
+            gt_ans = vqa_tool.processDigitArticle(gt_ans)
+
+            if pred_ans == gt_ans:
+                acc += 1
+        acc = acc / len(merged_results) * 100
+        print(f"===== {task_name} Accuracy {acc:.2f}% =====")
+        with open(os.path.abspath(os.path.join(results_save_dir, "scores.txt")), "a") as f:
+            f.write(f"{task_name} Accuracy = {acc:.2f}%\n\n")
+        return acc
+
+    elif task_name == "docvqa" or task_name == "docvqa_test":
+        vqa_tool = VQAEval()
+        anls = 0
+        for res in merged_results:
+            pred_ans = res["answer"]
+            gt_ans = res["gt_answer"]
+            anls += get_anls_score(pred=pred_ans, gold_labels=gt_ans, threshold=0.5)
+        anls = anls / len(merged_results) * 100
+        print(f"===== {task_name} ANLS {anls:.2f}% =====")
+        with open(os.path.abspath(os.path.join(results_save_dir, "scores.txt")), "a") as f:
+            f.write(f" {task_name} ANLS = {anls:.2f}%\n\n")
+        return anls
+
+    elif task_name == "chartqa":
+        vqa_tool = VQAEval()
+        acc = 0
+        for res in merged_results:
+            pred_ans = res["answer"]
+            pred_ans = vqa_tool.processPunctuation(pred_ans)
+            pred_ans = vqa_tool.processDigitArticle(pred_ans)
+
+            gt_ans = res["gt_answer"][0]
+            gt_ans = vqa_tool.processPunctuation(gt_ans)
+            gt_ans = vqa_tool.processDigitArticle(gt_ans)
+
+            # ChartQA uses relaxed accuracy:
+            # "We consider an answer to be correct if it is within 5% of the gold answer.
+            #  For non-numeric answers, we still need an exact match to consider an answer to be correct."
+            if isNumber(pred_ans) and isNumber(gt_ans):
+                pred_ans = float(pred_ans)
+                gt_ans = float(gt_ans)
+                if pred_ans >= (gt_ans * 0.95) and pred_ans <= (gt_ans * 1.05):
+                    acc += 1
+            elif pred_ans == gt_ans:
+                acc += 1
+        acc = acc / len(merged_results) * 100
+        print(f"===== {task_name} Accuracy {acc:.2f}% =====")
+        with open(os.path.abspath(os.path.join(results_save_dir, "scores.txt")), "a") as f:
+            f.write(f"{task_name} Accuracy = {acc:.2f}%\n\n")
+        return acc
+
+    elif task_name == 'ocrbench':
+
+        from collections import defaultdict
+        OCRBench_score = {"Regular Text Recognition": 0, "Irregular Text Recognition": 0,
+                          "Artistic Text Recognition": 0, "Handwriting Recognition": 0,
+                          "Digit String Recognition": 0, "Non-Semantic Text Recognition": 0,
+                          "Scene Text-centric VQA": 0, "Doc-oriented VQA": 0, "Doc-oriented VQA": 0,
+                          "Key Information Extraction": 0, "Handwritten Mathematical Expression Recognition": 0}
+
+        for res in merged_results:
+            predict = res["answer"]
+            answers = res["gt_answer"]
+
+            dataset_name = res["dataset_name"]
+            ocr_type = res["type"]
+
+            # data[i]['result'] = 0
+            if dataset_name == "HME100k":
+                if type(answers) == list:
+                    for j in range(len(answers)):
+                        answer = answers[j].strip().replace("\n", " ").replace(" ", "")
+                        predict = predict.strip().replace("\n", " ").replace(" ", "")
+                        if answer in predict:
+                            OCRBench_score[ocr_type] += 1
+                else:
+                    answers = answers.strip().replace("\n", " ").replace(" ", "")
+                    predict = predict.strip().replace("\n", " ").replace(" ", "")
+                    if answers in predict:
+                        OCRBench_score[ocr_type] += 1
+            else:
+                if type(answers) == list:
+                    for j in range(len(answers)):
+                        answer = answers[j].lower().strip().replace("\n", " ")
+                        predict = predict.lower().strip().replace("\n", " ")
+                        if answer in predict:
+                            OCRBench_score[ocr_type] += 1
+                else:
+                    answers = answers.lower().strip().replace("\n", " ")
+                    predict = predict.lower().strip().replace("\n", " ")
+                    if answers in predict:
+                        OCRBench_score[ocr_type] += 1
+
+        recognition_score = OCRBench_score['Regular Text Recognition'] + OCRBench_score['Irregular Text Recognition'] + \
+                            OCRBench_score['Artistic Text Recognition'] + OCRBench_score['Handwriting Recognition'] + \
+                            OCRBench_score['Digit String Recognition'] + OCRBench_score['Non-Semantic Text Recognition']
+        Final_score = recognition_score + OCRBench_score['Scene Text-centric VQA'] + OCRBench_score[
+            'Doc-oriented VQA'] + OCRBench_score['Key Information Extraction'] + OCRBench_score[
+                          'Handwritten Mathematical Expression Recognition']
+        result_log = f"\n###########################OCRBench##############################\n\
+            Text Recognition(Total 300):{recognition_score}\n\
+            ------------------Details of Recognition Score-------------------\n\
+            Regular Text Recognition(Total 50): {OCRBench_score['Regular Text Recognition']}\n\
+            Irregular Text Recognition(Total 50): {OCRBench_score['Irregular Text Recognition']}\n\
+            Artistic Text Recognition(Total 50): {OCRBench_score['Artistic Text Recognition']}\n\
+            Handwriting Recognition(Total 50): {OCRBench_score['Handwriting Recognition']}\n\
+            Digit String Recognition(Total 50): {OCRBench_score['Digit String Recognition']}\n\
+            Non-Semantic Text Recognition(Total 50): {OCRBench_score['Non-Semantic Text Recognition']}\n\
+            ----------------------------------------------------------------\n\
+            Scene Text-centric VQA(Total 200): {OCRBench_score['Scene Text-centric VQA']}\n\
+            ----------------------------------------------------------------\n\
+            Doc-oriented VQA(Total 200): {OCRBench_score['Doc-oriented VQA']}\n\
+            ----------------------------------------------------------------\n\
+            Key Information Extraction(Total 200): {OCRBench_score['Key Information Extraction']}\n\
+            ----------------------------------------------------------------\n\
+            Handwritten Mathematical Expression Recognition(Total 100): {OCRBench_score['Handwritten Mathematical Expression Recognition']}\n\
+            ----------------------Final Score-------------------------------\n\
+            Final Score(Total 1000): {Final_score}\n"
+
+        print(f"=====  {task_name} Final_score(Total 1000): {Final_score:.2f}% ===== {result_log}")
+        with open(os.path.abspath(os.path.join(results_save_dir, "scores.txt")), "a") as f:
+            f.write(f" {task_name} Accuracy = {Final_score:.2f}% {result_log}\n\n")
+        return Final_score
+
+
+    elif task_name == "mathvista":
+        import re
+        def extra_processing(text):
+
+            ## max decimal point capped to 2 decimal point
+            regex = re.compile(r'^\d+\.\d+$')
+            decimal = regex.findall(text)
+
+            if len(decimal) > 0:
+                non_decimal = len(decimal[0].split(".")[0])
+
+                # if decimal values are all 0, trim them
+
+                decimal_digits = [int(d) for d in decimal[0].split(".")[1]]
+                if sum(decimal_digits) == 0:
+                    text = decimal[0][:non_decimal]
+                else:
+                    text = decimal[0][:non_decimal + 3]
+
+            ## remove %
+            text = text.replace("%", "")
+            try:
+                if text[-1] == ".":
+                    text = text[:-1]
+            except:
+                print(text)
+
+            return text
+
+        def extract_answer(text):
+
+            alphabet = re.findall(r'[a-zA-Z]+', text)
+            if len(alphabet) > 0 and "e+" not in text:
+
+                template_1 = re.findall(r'answer is -*\d+\.*\d*', text)
+                if len(template_1) > 0:
+                    text = template_1[0]
+
+                    numbers = re.findall(r'-*\d+\.*\d*', text)
+                    text = numbers[0] if len(numbers) > 0 else text
+
+            return text
+
+        vqa_tool = VQAEval()
+        acc = 0
+        for res in merged_results:
+            pred_ans = res["answer"]
+            if res['question_type'] == 'multi_choice':
+                pred_ans = parse_multi_choice_response(pred_ans, res['all_choices'], res['index2ans'])
+            else:
+                pred_ans = vqa_tool.processPunctuation(pred_ans)
+                pred_ans = vqa_tool.processDigitArticle(pred_ans)
+
+                pred_ans = extra_processing(pred_ans)
+                pred_ans = extract_answer(pred_ans)
+
+            gt_ans = res["gt_answer"]
+            if res['question_type'] != 'multi_choice':
+                gt_ans = vqa_tool.processPunctuation(gt_ans)
+                gt_ans = vqa_tool.processDigitArticle(gt_ans)
+                gt_ans = extra_processing(gt_ans)
+
+            if pred_ans == gt_ans:
+                acc += 1
+        acc = acc / len(merged_results) * 100
+        print(f"===== {task_name} Accuracy {acc:.2f}% =====")
+        with open(os.path.abspath(os.path.join(results_save_dir, "scores.txt")), "a") as f:
+            f.write(f" {task_name} Accuracy = {acc:.2f}%\n\n")
+        return acc
+    elif task_name == "mmbench":
+        if task_cfg['split'] == 'dev':
+            acc = 0
+            for res in merged_results:
+                gt_ans = res["gt_answer"]
+                pred_ans = res["answer"]
+                pred_ans = parse_multi_choice_response(pred_ans, res['all_choices'], res['index2ans'])
+
+                if pred_ans == gt_ans:
+                    acc += 1
+            acc = acc / len(merged_results) * 100
+            print(f"===== {task_name} Accuracy {acc:.2f}% =====")
+            with open(os.path.abspath(os.path.join(results_save_dir, "scores.txt")), "a") as f:
+                f.write(f"{task_name} Accuracy = {acc:.2f}%\n\n")
+
+            return acc
+
+        # Generate submission.xlsx for test set online evaluation, "https://mmbench.opencompass.org.cn/mmbench-submission"
+        if task_cfg['submission']:
+            res_df = pd.DataFrame(
+                {
+                    "index": [r['question_id'] for r in merged_results],
+                    "question": [r['question'] for r in merged_results],
+                    "A": [r['index2ans']['A'] if 'A' in r['index2ans'] else None for r in merged_results],
+                    "B": [r['index2ans']['B'] if 'B' in r['index2ans'] else None for r in merged_results],
+                    "C": [r['index2ans']['C'] if 'C' in r['index2ans'] else None for r in merged_results],
+                    "D": [r['index2ans']['D'] if 'D' in r['index2ans'] else None for r in merged_results],
+                    "prediction": [parse_multi_choice_response(r['answer'], r['all_choices'], r['index2ans']) for r in
+                                   merged_results],
+                },
+                columns=['index', 'question', 'A', 'B', 'C', 'D', 'prediction']
+            )
+            res_df.to_excel(os.path.join(results_save_dir, "submission.xlsx"))
+
+
+    elif task_name == "vizwiz":
+        print(
+            f"VizWiz result file is saved at: {merged_results_path}\n"
+            f"Upload manually or use this CLI `python evaluation/upload_vizwiz.py --result {merged_results_path} --token <your_eval_user_token>`."
+        )
+
+    elif task_name == "mmmu" or task_name == 'mmmu_pro':
+        def extract_answer(text):
+            import re
+            # Regular expression to find content inside \answer{xxx}
+            match = re.search(r'\\answer\{(.*?)\}', text)
+            if match:
+                return match.group(1)  # Return the content inside the braces
+            return text  # Return the original string if no match is found
+
+        eval_samples = []
+        eval_output_dict = {}
+        for res in merged_results:
+            pred_ans = res["answer"]
+            gt_ans = res['gt_answer']
+            if res['question_type'] == 'multiple-choice':
+                parsed_pred = parse_multi_choice_response(pred_ans, res['all_choices'], res['index2ans'])
+                eval_samples.append(
+                    {
+                        'id': res['question_id'],
+                        'question_type': res['question_type'],
+                        'answer': res['gt_answer'],  # the content in option, not answer index.
+                        'response': pred_ans,
+                        'parsed_pred': parsed_pred,
+                        'index2ans': res['index2ans'],
+                    }
+                )
+                eval_output_dict[res['question_id']] = parsed_pred
+            else:
+                pred_ans = extract_answer(pred_ans)  # for sft v9prov3, we observe answers are within \answer{xxx}
+                parsed_pred = parse_open_response(pred_ans)
+                eval_samples.append(
+                    {
+                        'id': res['question_id'],
+                        'question_type': res['question_type'],
+                        'answer': res['gt_answer'],
+                        'response': pred_ans,
+                        'parsed_pred': parsed_pred,
+                    }
+                )
+                eval_output_dict[res['question_id']] = pred_ans
+        if args.subset is not None:
+            eval_output_dict_path = os.path.join(results_save_dir,
+                                                 f"eval_output_dict_{task_name}_subset_{args.subset}_start_{args.start_idx}.json")
+        else:
+            eval_output_dict_path = os.path.join(results_save_dir, f"eval_output_dict_{task_name}.json")
+        json.dump(eval_output_dict, open(eval_output_dict_path, "w"), indent=4, sort_keys=True)
+
+        mmmu_results = mmmu_main_eval(eval_output_dict, task_cfg)
+        with open(os.path.join(results_save_dir, "scores.txt"), "a") as f:
+            f.write(f"{task_name} {task_cfg['split']}:\n")
+            for cat, cat_val in mmmu_results.items():
+                if 'Overall' in cat:
+                    cat = cat.replace("Overall-", "")
+                    print(f'{cat}: {cat_val["acc"] * 100:.2f}')
+                    f.write(f'{cat}: {cat_val["acc"] * 100:.2f}\n')
+        return mmmu_results['Overall']['acc']
+
+    else:
+        raise NotImplementedError(f"Task {task_name} is not supported yet.")
+
+
+def evaluate(args, model, tokenizer, tasks):
+    start_time = time.time()
+
+    results_save_dir = os.path.join(args.result_save_path)
+    if not os.path.exists(results_save_dir):
+        os.makedirs(results_save_dir)
+
+    task_names = list(tasks.keys())
+
+    print(f"Evaluating tasks: {task_names}")
+
+    with torch.no_grad():
+        # for task_name in task_names:
+        for task_name in args.zero_shot_eval_tasks:
+            task_cfg = tasks[task_name]
+            print(f"Preparing dataloader for {task_name}...")
+
+            dataloader = get_task_dataloader(task_name, task_cfg, args)
+
+            generation_config = dict(max_new_tokens=1024, do_sample=False)
+
+            print("Start generating...")
+            generate_task_results(task_name, task_cfg, args, model, tokenizer, generation_config,
+                                  dataloader, results_save_dir)
+
+            print("Start calculating task metric...")
+            calc_task_metrics(args, task_name, task_cfg, results_save_dir)
+
+    end_time = time.time()
+
+    print(f"Evaluation takes {(end_time - start_time) / 60:.1f} minutes in total!")
+
+
+if __name__ == '__main__':
+    path = "nvidia/NVLM-D-72B"
+    print("Loading model... from", path)
+    device_map = split_model()
+
+    start = time.time()
+    model = AutoModel.from_pretrained(
+        path,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        use_flash_attn=True,
+        device_map=device_map,
+        trust_remote_code=True).eval()
+    end = time.time()
+    print("loading model takes:", end - start)
+
+    print(model)
+
+    tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config-path', type=str, required=True,
+                        help='YAML file configuring evaluation datasets')
+    parser.add_argument('--result-save-path', type=str, default=os.path.join(path, '/eval_results'))
+    parser.add_argument('--zero-shot-eval-tasks', nargs='+', type=str, default=['mmmu'])
+    parser.add_argument('--start-idx', type=int, default=0)
+    parser.add_argument('--subset', type=int, default=None)
+
+    args = parser.parse_args()
+
+    tasks = yaml.safe_load(open(args.config_path))['datasets']
+
+    evaluate(args, model, tokenizer, tasks)