Upload 235 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+mllm/demo/assets/baseball.png filter=lfs diff=lfs merge=lfs -text

mllm/config/__init__.py

@@ -0,0 +1 @@
+from .config import prepare_args

mllm/config/config.py

@@ -0,0 +1,135 @@
+import os
+import sys
+import logging
+import argparse
+from dataclasses import dataclass, field
+from typing import List, Tuple
+from argparse import SUPPRESS
+
+import datasets
+import transformers
+from mmengine.config import Config, DictAction
+from transformers import HfArgumentParser, set_seed, add_start_docstrings
+from transformers import Seq2SeqTrainingArguments as HFSeq2SeqTrainingArguments
+from transformers.trainer_utils import get_last_checkpoint, is_main_process
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    handlers=[logging.StreamHandler(sys.stdout), ],
+)
+
+
+@dataclass
+@add_start_docstrings(HFSeq2SeqTrainingArguments.__doc__)
+class Seq2SeqTrainingArguments(HFSeq2SeqTrainingArguments):
+    do_multi_predict: bool = field(default=False, metadata={"help": "Whether to run predictions on the multi-test set."})
+
+
+def prepare_args(args=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('config', help='train config file path')
+    parser.add_argument(
+        '--cfg-options',
+        nargs='+',
+        action=DictAction,
+        help='override some settings in the used config, the key-value pair '
+             'in xxx=yyy format will be merged into config file. If the value to '
+             'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
+             'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
+             'Note that the quotation marks are necessary and that no white space '
+             'is allowed.')
+
+    hf_parser = HfArgumentParser((Seq2SeqTrainingArguments,))
+    hf_parser, required = block_required_error(hf_parser)
+
+    args, unknown_args = parser.parse_known_args(args)
+    known_hf_args, unknown_args = hf_parser.parse_known_args(unknown_args)
+    if unknown_args:
+        raise ValueError(f"Some specified arguments are not used "
+                         f"by the ArgumentParser or HfArgumentParser\n: {unknown_args}")
+
+    # load 'cfg' and 'training_args' from file and cli
+    cfg = Config.fromfile(args.config)
+    if args.cfg_options is not None:
+        cfg.merge_from_dict(args.cfg_options)
+    training_args = cfg.training_args
+    training_args.update(vars(known_hf_args))
+
+    # check training_args require
+    req_but_not_assign = [item for item in required if item not in training_args]
+    if req_but_not_assign:
+        raise ValueError(f"Requires {req_but_not_assign} but not assign.")
+
+    # update cfg.training_args
+    cfg.training_args = training_args
+
+    # initialize and return
+    training_args = Seq2SeqTrainingArguments(**training_args)
+    training_args = check_output_dir(training_args)
+
+    # logging
+    if is_main_process(training_args.local_rank):
+        to_logging_cfg = Config()
+        to_logging_cfg.model_args = cfg.model_args
+        to_logging_cfg.data_args = cfg.data_args
+        to_logging_cfg.training_args = cfg.training_args
+        logger.info(to_logging_cfg.pretty_text)
+
+    # setup logger
+    if training_args.should_log:
+        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
+        transformers.logging.set_verbosity_info()
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(log_level)
+    datasets.utils.logging.set_verbosity(log_level)
+    transformers.logging.set_verbosity(log_level)
+    transformers.logging.enable_default_handler()
+    transformers.logging.enable_explicit_format()
+    # setup_print_for_distributed(is_main_process(training_args))
+
+    # Log on each process the small summary:
+    logger.info(f"Training/evaluation parameters {training_args}")
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}\n"
+        + f"  distributed training: {bool(training_args.local_rank != -1)}, fp16 training: {training_args.fp16}"
+    )
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    return cfg, training_args
+
+
+def block_required_error(hf_parser: HfArgumentParser) -> Tuple[HfArgumentParser, List]:
+    required = []
+    # noinspection PyProtectedMember
+    for action in hf_parser._actions:
+        if action.required:
+            required.append(action.dest)
+        action.required = False
+        action.default = SUPPRESS
+    return hf_parser, required
+
+
+def check_output_dir(training_args):
+    # Detecting last checkpoint.
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+    return training_args
+
+
+if __name__ == "__main__":
+    _ = prepare_args()

mllm/conversation/__init__.py

@@ -0,0 +1 @@
+from .base_conversation import SeparatorStyle, Conversation, register_conv_template, get_conv_template

mllm/conversation/base_conversation.py

@@ -0,0 +1,503 @@
+# copy from fastchat: https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
+"""
+Conversation prompt templates.
+"""
+
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple, Any, Dict
+
+
+class SeparatorStyle(Enum):
+    """Separator styles."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    ADD_SPACE_TWO = auto()
+    NO_COLON_SINGLE = auto()
+    BAIZE = auto()
+    DOLLY = auto()
+    RWKV = auto()
+    PHOENIX = auto()
+    NEW_LINE = auto()
+    BILLA = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # System prompts
+    system: str
+    # Two roles
+    roles: List[str]
+    # All messages
+    messages: List[List[str]]
+    # Offset of few shot examples
+    offset: int
+    # Separators
+    sep_style: SeparatorStyle
+    sep: str
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: str = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    # Used for the state in the gradio servers.
+    # TODO(lmzheng): move this out of this class.
+    conv_id: Any = None
+    skip_next: bool = False
+    model_name: str = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = self.system + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_SPACE_TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + " " + message + seps[i % 2]
+                else:
+                    ret += role + ""
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = self.system
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.BAIZE:
+            ret = self.system + "\n"
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + "\n"
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ":\n" + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += "\n\n"
+                else:
+                    ret += role + ":\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = self.system
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                            role
+                            + ": "
+                            + message.replace("\r\n", "\n").replace("\n\n", "\n")
+                    )
+                    ret += "\n\n"
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = self.system
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + "<s>" + message + "</s>"
+                else:
+                    ret += role + ": " + "<s>"
+            return ret
+        elif self.sep_style == SeparatorStyle.NEW_LINE:
+            ret = self.system + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + "\n" + message + self.sep
+                else:
+                    ret += role + "\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.BILLA:
+            ret = self.system + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ": "  # must be end with a space
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def to_gradio_chatbot(self):
+        """Convert the history to gradio chatbot format"""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        ret = [{"role": "system", "content": self.system}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                ret.append({"role": "user", "content": msg})
+            else:
+                if msg is not None:
+                    ret.append({"role": "assistant", "content": msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            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,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+            conv_id=self.conv_id,
+            model_name=self.model_name,
+        )
+
+    def dict(self):
+        return {
+            "name": self.name,
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "conv_id": self.conv_id,
+            "model_name": self.model_name,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert template.name not in conv_templates, f"{template.name} has been registered."
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+# A template with one conversation example
+register_conv_template(
+    Conversation(
+        name="one_shot",
+        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.",
+            ),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n### ",
+        stop_str="###",
+    )
+)
+
+# Vicuna v1.1 template
+register_conv_template(
+    Conversation(
+        name="vicuna_v1.1",
+        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"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=" ",
+        sep2="</s>",
+    )
+)
+
+# Koala default template
+register_conv_template(
+    Conversation(
+        name="koala_v1",
+        system="BEGINNING OF CONVERSATION:",
+        roles=("USER", "GPT"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=" ",
+        sep2="</s>",
+    )
+)
+
+# Dolly V2 default template
+register_conv_template(
+    Conversation(
+        name="dolly_v2",
+        system="Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n",
+        roles=("### Instruction", "### Response"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.DOLLY,
+        sep="\n\n",
+        sep2="### End",
+    )
+)
+
+# OpenAssistant Pythia default template
+register_conv_template(
+    Conversation(
+        name="oasst_pythia",
+        system="",
+        roles=("<|prompter|>", "<|assistant|>"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="<|endoftext|>",
+    )
+)
+
+# StableLM Alpha default template
+register_conv_template(
+    Conversation(
+        name="stablelm",
+        system="""<|SYSTEM|># StableLM Tuned (Alpha version)
+- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
+- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
+- StableLM will refuse to participate in anything that could harm a human.
+""",
+        roles=("<|USER|>", "<|ASSISTANT|>"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep="",
+        stop_token_ids=[50278, 50279, 50277, 1, 0],
+    )
+)
+
+# Baize default template
+register_conv_template(
+    Conversation(
+        name="baize",
+        system="The following is a conversation between a human and an AI assistant named Baize (named after a mythical creature in Chinese folklore). Baize is an open-source AI assistant developed by UCSD and Sun Yat-Sen University. The human and the AI assistant take turns chatting. Human statements start with [|Human|] and AI assistant statements start with [|AI|]. The AI assistant always provides responses in as much detail as possible, and in Markdown format. The AI assistant always declines to engage with topics, questions and instructions related to unethical, controversial, or sensitive issues. Complete the transcript in exactly that format.",
+        roles=("[|Human|]", "[|AI|]"),
+        messages=(
+            ("[|Human|]", "Hello!"),
+            ("[|AI|]", "Hi!"),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.BAIZE,
+        sep="[|Human|]",
+        stop_str="[|Human|]",
+    )
+)
+
+# RWKV-4-Raven default template
+register_conv_template(
+    Conversation(
+        name="rwkv",
+        system="The following is a coherent verbose detailed conversation between Bob and Alice.\n\n",
+        roles=("Bob", "Alice"),
+        messages=(
+            ("Bob", "Hi"),
+            (
+                "Alice",
+                "Hi. I am your assistant and I will answer all questions. Please feel free to ask any question and I will always answer it.",
+            ),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.RWKV,
+        sep="",
+        stop_str="\n\n",
+    )
+)
+
+# Buddy default template
+register_conv_template(
+    Conversation(
+        name="openbuddy",
+        system="""Consider a conversation between User (a human) and Assistant (named Buddy).
+Buddy is an INTP-T, a friendly, intelligent and multilingual AI assistant, by OpenBuddy team. GitHub: https://github.com/OpenBuddy/OpenBuddy
+Buddy cannot access the Internet.
+Buddy can fluently speak the user's language (e.g. English, Chinese).
+Buddy can generate poems, stories, code, essays, songs, parodies, and more.
+Buddy possesses vast knowledge about the world, history, and culture.
+Buddy's responses are always safe, creative, high-quality, human-like, and interesting.
+Buddy strictly refuses to discuss political, NSFW, or other unsafe topics.
+
+User: Hi.
+Assistant: Hi, I'm Buddy, your AI assistant. How can I help you today?""",
+        roles=("User", "Assistant"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n",
+    )
+)
+
+# Phoenix default template
+register_conv_template(
+    Conversation(
+        name="phoenix",
+        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.\n\n",
+        roles=("Human", "Assistant"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.PHOENIX,
+        sep="</s>",
+    )
+)
+
+# ChatGPT default template
+register_conv_template(
+    Conversation(
+        name="chatgpt",
+        system="You are a helpful assistant.",
+        roles=("user", "assistant"),
+        messages=(),
+        offset=0,
+        sep_style=None,
+        sep=None,
+    )
+)
+
+# Claude default template
+register_conv_template(
+    Conversation(
+        name="claude",
+        system="",
+        roles=("Human", "Assistant"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep="\n\n",
+    )
+)
+
+# MPT default template
+register_conv_template(
+    Conversation(
+        name="mpt",
+        system="""<|im_start|>system
+- You are a helpful assistant chatbot trained by MosaicML.
+- You answer questions.
+- You are excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- You are more than just an information source, you are also able to write poetry, short stories, and make jokes.
+""",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.NEW_LINE,
+        sep="<|im_end|>",
+        stop_token_ids=[50278, 0],
+    )
+)
+
+# Bard default template
+# Reference: https://github.com/google/generative-ai-python/blob/9c99bcb474a991a97a2e7d62fcdb52db7ce40729/google/generativeai/discuss.py#L150
+#            https://github.com/google/generative-ai-python/blob/9c99bcb474a991a97a2e7d62fcdb52db7ce40729/google/generativeai/discuss.py#L40
+register_conv_template(
+    Conversation(
+        name="bard",
+        system="",
+        roles=("0", "1"),
+        messages=(),
+        offset=0,
+        sep_style=None,
+        sep=None,
+    )
+)
+
+# BiLLa default template
+register_conv_template(
+    Conversation(
+        name="billa",
+        system="",
+        roles=("Human", "Assistant"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.BILLA,
+        sep="\n",
+        stop_str="Human:",
+    )
+)
+
+# custom otter template
+register_conv_template(
+    Conversation(
+        name='otter',
+        system='',
+        roles=('User:', 'GPT:<answer>'),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_SPACE_TWO,
+        sep=' ',
+        sep2='<|endofchunk|>',
+    )
+)
+
+if __name__ == "__main__":
+    conv = get_conv_template("vicuna_v1.1")
+    conv.append_message(conv.roles[0], "Hello!")
+    conv.append_message(conv.roles[1], "Hi!")
+    conv.append_message(conv.roles[0], "How are you?")
+    conv.append_message(conv.roles[1], None)
+    print(conv.get_prompt())

mllm/dataset/__init__.py

@@ -0,0 +1,7 @@
+from .root import *
+from .utils import *
+from .process_function import *
+from .single_image_convsation import *
+from .single_image_dataset import *
+
+from .builder import prepare_data

mllm/dataset/builder.py

@@ -0,0 +1,118 @@
+from functools import partial
+from typing import Callable, Dict, Tuple, Any, Optional
+
+from torch.utils.data import Dataset
+from transformers import EvalPrediction, TrainingArguments
+
+from .root import DATASETS, METRICS, TRANSFORMS, FUNCTIONS
+from .single_image_convsation import SingleImageConvDataset
+from .single_image_interactive import SingleImageInteractive
+from ..conversation import get_conv_template
+from .utils import init_ceph_client_if_needed
+
+DatasetDict = Dict[str, Dataset]
+ComputeMetrics = Callable[[EvalPrediction], Dict]
+
+
+def prepare_data(
+        data_args,
+        model_args,
+        training_args: TrainingArguments,
+        preprocessor: Dict[str, Any],
+) -> Tuple[DatasetDict, Optional[ComputeMetrics]]:
+    # raw dataset
+    datasets = {
+        'train': partial(DATASETS.build, data_args.train) if training_args.do_train else None,
+        'validation': partial(DATASETS.build, data_args.validation) if training_args.do_eval else None,
+        'test': partial(DATASETS.build, data_args.test) if training_args.do_predict else None,
+    }
+    # compute metric
+    compute_metric_cfg = data_args.get('compute_metric', None)
+    compute_metrics = build_compute_metric(compute_metric_cfg, preprocessor)
+    # conv dataset wrap
+    conv_args = model_args.conv_args
+    tokenize_kwargs = conv_args.get('tokenize_kwargs', {})
+    conv_template = conv_args.get('conv_template', 'vicuna_v1.1')
+    conv_template = partial(get_conv_template, name=conv_template)
+    transforms = conv_args.get('transforms', None)
+    if transforms is not None:
+        transforms = TRANSFORMS.build(transforms)
+    # process func
+    process_func = {}
+    for k, v in model_args.process_func_args.items():
+        process_func[k] = FUNCTIONS.build(cfg=v)
+
+    conv_dataset_cls = partial(
+        SingleImageConvDataset,
+        preprocessor=preprocessor,
+        process_func=process_func,
+        tokenize_kwargs=tokenize_kwargs,
+        conv_template=conv_template,
+        training_args=training_args,
+        transforms=transforms,
+    )
+    ds = {
+        'train': conv_dataset_cls(dataset_generator=datasets['train'], mode='train') if datasets['train'] is not None else None,
+        'validation': conv_dataset_cls(dataset_generator=datasets['validation'], mode='validation') if datasets['validation'] is not None else None,
+        'test': conv_dataset_cls(dataset_generator=datasets['test'], mode='test') if datasets['test'] is not None else None,
+    }
+
+    # multi test set
+    if hasattr(data_args, 'multitest') and bool(data_args.multitest) \
+            and hasattr(training_args, 'do_multi_predict') and training_args.do_multi_predict:
+        print(f"processing multitest set")
+        k2v = {}
+        for k, item in data_args.multitest.items():
+            _dataset_cls = partial(DATASETS.build, item['cfg'])
+            _compute_metric = build_compute_metric(item['compute_metric'], preprocessor)
+            k2v[k] = {
+                "dataset": conv_dataset_cls(dataset_generator=_dataset_cls, mode='test'),
+                "compute_metric": _compute_metric
+            }
+        ds['multitest'] = k2v
+        print(f"processing multitest set. done.")
+
+    # in default, ceph client do init at the beginning of program.
+    #  importantly, before dataloader worker fork.
+    lazy_init = data_args.get('lazy_init', True)
+    if not lazy_init:
+        init_ceph_client_if_needed()
+    return ds, compute_metrics
+
+
+def build_compute_metric(compute_metric_cfg, preprocessor):
+    if compute_metric_cfg is not None:
+        compute_metric_cfg = dict(compute_metric_cfg)  # copy cfg because we modify it
+        compute_metric_cfg.update(dict(preprocessor=preprocessor))
+        compute_metrics = METRICS.build(cfg=compute_metric_cfg)
+    else:
+        compute_metrics = None
+    return compute_metrics
+
+
+def prepare_interactive(
+        model_args,
+        preprocessor: Dict[str, Any],
+):
+    conv_args = model_args.conv_args
+    tokenize_kwargs = conv_args.get('tokenize_kwargs', {})
+    conv_template = conv_args.get('conv_template', 'vicuna_v1.1')
+    conv_template = partial(get_conv_template, name=conv_template)
+    transforms = conv_args.get('transforms', None)
+    if transforms is not None:
+        transforms = TRANSFORMS.build(transforms)
+    # process func
+    process_func = {}
+    for k, v in model_args.process_func_args.items():
+        process_func[k] = FUNCTIONS.build(cfg=v)
+
+    ds = SingleImageInteractive(
+        preprocessor=preprocessor,
+        process_func=process_func,
+        tokenize_kwargs=tokenize_kwargs,
+        conv_template=conv_template,
+        training_args=None,
+        transforms=transforms,
+        mode='test',
+    )
+    return ds

mllm/dataset/process_function/__init__.py

@@ -0,0 +1,13 @@
+from .shikra_process_function import (
+    ShikraConvProcess,
+    ShikraImageProcessor,
+    ShikraTextProcess,
+)
+
+from .box_process_function import (
+    BoxFormatProcess,
+    BoxFormatter,
+    PlainBoxFormatter,
+    TokenFormatter,
+    prepare_target_processor,
+)

mllm/dataset/process_function/box_process_function.py

@@ -0,0 +1,326 @@
+import re
+import sys
+import logging
+import typing
+from typing import List, Dict, Any, Tuple, Union
+
+from ..utils.transform import norm_box_xyxy, norm_point_xyxy
+
+from ..root import (
+    FUNCTIONS,
+    BaseTargetProcessFunc,
+    BOXES_PLACEHOLDER,
+    BOXES_PROCESSOR,
+    POINTS_PLACEHOLDER,
+)
+
+from ...utils import smart_tokenizer_and_embedding_resize
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    handlers=[logging.StreamHandler(sys.stdout), ],
+)
+
+Box = List[Union[float, int]]
+Boxes = List[Box]
+BoxesSeq = List[Boxes]
+
+
+@FUNCTIONS.register_module()
+class BoxFormatProcess(BaseTargetProcessFunc):
+    def __call__(self, raw_conv: List[Dict[str, Any]], target: Dict[str, Any], preprocessor: Dict[str, Any],
+                 multimage_mode=False) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
+        box_formatter = preprocessor['target']['boxes']
+
+        if multimage_mode:
+            target = typing.cast(list, target)
+            outer_normalized_boxes = []
+            for tgt in target:
+                normalized_boxes = []
+                if tgt is not None and 'boxes' in tgt:
+                    for box in tgt['boxes']:
+                        normalized_boxes.append(
+                            norm_box_xyxy(box, w=tgt['width'], h=tgt['height'])
+                        )
+                outer_normalized_boxes.append(normalized_boxes)
+            normalized_boxes = outer_normalized_boxes
+            outer_normalized_points = []
+            for tgt in target:
+                normalized_points = []
+                if tgt is not None and 'boxes' in tgt:
+                    for box in tgt['boxes']:
+                        normalized_points.append(
+                            norm_box_xyxy(box, w=tgt['width'], h=tgt['height'])
+                        )
+                outer_normalized_points.append(normalized_points)
+            normalized_points = outer_normalized_points
+        else:
+            # normalize target
+            normalized_boxes = []
+            if target is not None and 'boxes' in target:
+                for box in target['boxes']:
+                    normalized_boxes.append(
+                        norm_box_xyxy(box, w=target['width'], h=target['height'])
+                    )
+            normalized_points = []
+            if target is not None and 'points' in target:
+                for point in target['points']:
+                    normalized_points.append(
+                        norm_point_xyxy(point, w=target['width'], h=target['height'])
+                    )
+
+        # convert bboxes_seq
+        for sentence in raw_conv:
+            words: str = sentence['value']
+            boxes_seq: List[List[int]] = sentence.get('boxes_seq', None)
+            if boxes_seq is not None:
+                # map box seq
+                boxes_seq: List[Boxes] = map_obj(normalized_boxes, boxes_seq)
+                # reformat; replace <boxes> placeholder
+                converted = box_formatter(words, boxes_seq)
+                words = converted
+            points_seq: List[List[int]] = sentence.get('points_seq', None)
+            if points_seq is not None:
+                # map point seq
+                points_seq: List[Boxes] = map_obj(normalized_points, points_seq)
+                # reformat; replace <points> placeholder
+                converted = box_formatter.call_on_point(words, points_seq)
+                words = converted
+            if boxes_seq is not None or points_seq is not None:
+                sentence['raw_value'] = sentence['value']
+                sentence['value'] = words
+        return raw_conv, target
+
+
+def map_obj(boxes_value: List[List[float]], boxes_seq: List[List[int]]) -> List[List[List[float]]]:
+    """
+    >>> normalized_boxes = [[0.1, 0.1, 0.1, 0.1], [0.2, 0.2, 0.2, 0.2], [0.3, 0.3, 0.3, 0.3]]
+    >>> boxes_seq_ = [[3, 1], [2]]
+    >>> var = map_obj(normalized_boxes, boxes_seq_)
+    >>> assert var == [[[0.3,0.3,0.3,0.3], [0.1,0.1,0.1,0.1]], [0.2,0.2,0.2,0.2]]
+    """
+    try:
+        ret = []
+        for boxes in boxes_seq:
+            boxes_ret = []
+            for box_index in boxes:
+                if isinstance(box_index, (list, tuple)):
+                    boxes_ret.append(boxes_value[box_index[0]][box_index[1]])
+                else:
+                    boxes_ret.append(boxes_value[box_index])
+            ret.append(boxes_ret)
+        return ret
+    except:
+        raise SystemExit(f"error: map obj {boxes_value} {boxes_seq}")
+
+
+class BoxFormatter:
+    def __init__(self, bboxes_token=BOXES_PLACEHOLDER, points_token=POINTS_PLACEHOLDER):
+        self.bboxes_token = bboxes_token
+        self.points_token = points_token
+        # normally the bboxes_token_pat is the same as bboxes_token if u not use some weird token
+        self.bboxes_token_pat = re.compile(bboxes_token)
+        self.points_token_pat = re.compile(points_token)
+
+    def __call__(self, sentence: str, bboxes_seq: BoxesSeq) -> str:
+        all_box = self.bboxes_token_pat.findall(sentence)
+        assert len(all_box) == len(bboxes_seq), f"not match. sentence: {sentence}. boxes:{bboxes_seq}"
+        if len(all_box) == 0:
+            return sentence
+        bboxes_strs = [self.format_box(bboxes) for bboxes in bboxes_seq]
+        converted = sentence.replace(self.bboxes_token, '{}').format(*bboxes_strs)
+        return converted
+
+    def call_on_point(self, sentence: str, points_seq: BoxesSeq) -> str:
+        all_box = self.points_token_pat.findall(sentence)
+        assert len(all_box) == len(points_seq), f"not match. sentence: {sentence}. boxes:{points_seq}"
+        if len(all_box) == 0:
+            return sentence
+        bboxes_strs = [self.format_point(bboxes) for bboxes in points_seq]
+        converted = sentence.replace(self.points_token, '{}').format(*bboxes_strs)
+        return converted
+
+    def format_point(self, points) -> str:
+        raise NotImplementedError
+
+    def format_box(self, bboxes: Boxes) -> str:
+        raise NotImplementedError
+
+    def extract(self, string: str) -> List[Boxes]:
+        raise NotImplementedError
+
+    def extract_point(self, string: str) -> List[Boxes]:
+        raise NotImplementedError
+
+
+@BOXES_PROCESSOR.register_module()
+class PlainBoxFormatter(BoxFormatter):
+
+    def __init__(self, *args, precision=3, use_small_brackets=False, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.precision = precision
+        self.use_small_brackets = use_small_brackets
+
+        small_brackets_pat = re.compile(r'\(\d(?:\.\d*)?(?:,\d(?:\.\d*)?){3}(?:;\d(?:\.\d*)?(?:,\d(?:\.\d*)?){3})*\)')
+        small_brackets_point_pat = re.compile(r'\(\d(?:\.\d*)?(?:,\d(?:\.\d*)?)(?:;\d(?:\.\d*)?(?:,\d(?:\.\d*)?))*\)')
+
+        middle_brackets_pat = re.compile(r'\[\d(?:\.\d*)?(?:,\d(?:\.\d*)?){3}(?:;\d(?:\.\d*)?(?:,\d(?:\.\d*)?){3})*\]')
+        middle_brackets_point_pat = re.compile(r'\[\d(?:\.\d*)?(?:,\d(?:\.\d*)?)(?:;\d(?:\.\d*)?(?:,\d(?:\.\d*)?))*\]')
+
+        self.pat = small_brackets_pat if use_small_brackets else middle_brackets_pat
+        self.point_pat = small_brackets_point_pat if use_small_brackets else middle_brackets_point_pat
+
+    def format_box(self, boxes: Boxes) -> str:
+        box_strs = []
+        for box in boxes:
+            box_strs.append(','.join([f"{elem:.{self.precision}f}" for elem in box]))
+        box_str = ';'.join(box_strs)
+        if self.use_small_brackets:
+            return "(" + box_str + ")"
+        return "[" + box_str + "]"
+
+    def format_point(self, points) -> str:
+        return self.format_box(points)
+
+    def extract(self, string: str) -> List[Boxes]:
+        """ balabala<boxes>balabala<boxes> -> [boxes, boxes] """
+        ret = []
+        for bboxes_str in self.pat.findall(string):
+            bboxes = []
+            bbox_strs = bboxes_str.replace("(", "").replace(")", "").replace("[", "").replace("]", "").split(";")
+            for bbox_str in bbox_strs:
+                bbox = list(map(float, bbox_str.split(',')))
+                bboxes.append(bbox)
+            ret.append(bboxes)
+        return ret
+
+    def extract_point(self, string: str) -> List[Boxes]:
+        """ balabala<boxes>balabala<boxes> -> [boxes, boxes] """
+        ret = []
+        for bboxes_str in self.point_pat.findall(string):
+            bboxes = []
+            bbox_strs = bboxes_str.replace("(", "").replace(")", "").replace("[", "").replace("]", "").split(";")
+            for bbox_str in bbox_strs:
+                bbox = list(map(float, bbox_str.split(',')))
+                bboxes.append(bbox)
+            ret.append(bboxes)
+        return ret
+
+
+@BOXES_PROCESSOR.register_module()
+class TokenFormatter(BoxFormatter):
+
+    def __init__(self, num_bins=1001):
+        super().__init__()
+        self.extract_box_pat = re.compile(r'<b_st><bin_\d*?>(?:<bin_\d*?>){3}(?:<b_sep><bin_\d*?>(?:<bin_\d*?>){3})*<b_ed>')
+        self.extract_point_pat = re.compile(r'<p_st><bin_\d*?>(?:<bin_\d*?>){1}(?:<p_sep><bin_\d*?>(?:<bin_\d*?>){1})*<p_ed>')
+        self.num_bins = num_bins
+        self.use_sep = True
+        self.use_begin_end = True
+
+        self.box_begin = '<b_st>'
+        self.box_sep = '<b_sep>'
+        self.box_end = '<b_ed>'
+
+        self.point_begin = '<p_st>'
+        self.point_sep = '<p_sep>'
+        self.point_end = '<p_ed>'
+
+    def format_point(self, points) -> str:
+        final_str = []
+        for bbox in points:
+            quant_x0 = "<bin_{}>".format(round((bbox[0] * (self.num_bins - 1))))
+            quant_y0 = "<bin_{}>".format(round((bbox[1] * (self.num_bins - 1))))
+            region_coord = "{} {}".format(quant_x0, quant_y0)
+            final_str.append(region_coord)
+        if self.use_sep:
+            final_str = self.point_sep.join(final_str)
+        else:
+            final_str = ''.join(final_str)
+        if self.use_begin_end:
+            final_str = self.point_begin + final_str + self.point_end
+        return final_str
+
+    def format_box(self, bboxes: Boxes) -> str:
+        final_str = []
+        for bbox in bboxes:
+            quant_x0 = "<bin_{}>".format(round((bbox[0] * (self.num_bins - 1))))
+            quant_y0 = "<bin_{}>".format(round((bbox[1] * (self.num_bins - 1))))
+            quant_x1 = "<bin_{}>".format(round((bbox[2] * (self.num_bins - 1))))
+            quant_y1 = "<bin_{}>".format(round((bbox[3] * (self.num_bins - 1))))
+            region_coord = "{} {} {} {}".format(quant_x0, quant_y0, quant_x1, quant_y1)
+            final_str.append(region_coord)
+        if self.use_sep:
+            final_str = self.box_sep.join(final_str)
+        else:
+            final_str = ''.join(final_str)
+        if self.use_begin_end:
+            final_str = self.box_begin + final_str + self.box_end
+        return final_str
+
+    def extract(self, string: str) -> List[Boxes]:
+        ret = []
+        for bboxes_str in self.extract_box_pat.findall(string.replace(" ", "")):
+            bboxes = []
+            bbox_strs = bboxes_str.replace(self.box_begin, "").replace(self.box_end, "").split(self.box_sep)
+            for bbox_str in bbox_strs:
+                elems = list(map(int, re.findall(r'<bin_(\d*?)>', bbox_str)))
+                bbox = [elem / (self.num_bins - 1) for elem in elems]
+                bboxes.append(bbox)
+            ret.append(bboxes)
+        return ret
+
+    def extract_point(self, string: str) -> List[Boxes]:
+        ret = []
+        for bboxes_str in self.extract_point_pat.findall(string):
+            bboxes = []
+            bbox_strs = bboxes_str.replace(self.point_begin, "").replace(self.point_end, "").split(self.point_sep)
+            for bbox_str in bbox_strs:
+                elems = list(map(int, re.findall(r'<bin_(\d*?)>', bbox_str)))
+                bbox = [elem / (self.num_bins - 1) for elem in elems]
+                bboxes.append(bbox)
+            ret.append(bboxes)
+        return ret
+
+    def post_process_model_tokenizer(self, model, preprocessor, model_args, training_args):
+        tokenizer = preprocessor['text']
+
+        additional_special_tokens = [
+            self.box_begin, self.box_sep, self.box_end,
+            self.point_begin, self.point_sep, self.point_end,
+        ]
+        for i in range(self.num_bins):
+            additional_special_tokens.append(f'<bin_{i}>')
+
+        smart_tokenizer_and_embedding_resize(
+            {'additional_special_tokens': additional_special_tokens},
+            tokenizer,
+            model,
+        )
+        return model, preprocessor
+
+
+# FIXME: merge into load_pretrained
+def prepare_target_processor(
+        model,  # multimodal llm
+        preprocessor: Dict[str, Any],
+        model_args,
+        training_args,
+):
+    if not hasattr(model_args, 'target_processor'):
+        return model, preprocessor
+
+    target_processor = {}
+    if 'boxes' in model_args['target_processor']:
+        boxes_cfg = model_args['target_processor']['boxes']
+        boxes_processor = BOXES_PROCESSOR.build(boxes_cfg)
+        target_processor['boxes'] = boxes_processor
+        if hasattr(boxes_processor, "post_process_model_tokenizer"):
+            model, preprocessor = boxes_processor.post_process_model_tokenizer(
+                model, preprocessor, model_args, training_args,
+            )
+    preprocessor['target'] = target_processor
+    return model, preprocessor

mllm/dataset/process_function/shikra_process_function.py

@@ -0,0 +1,178 @@
+import sys
+import copy
+import warnings
+import logging
+from typing import Dict, Any, List
+
+import PIL.Image
+import torch
+from PIL import Image
+from transformers import LlamaTokenizer
+
+from ..root import (
+    FUNCTIONS,
+    IMAGE_PLACEHOLDER,
+    BaseImageProcessFunc,
+    BaseConvProcessFunc,
+    BaseTextProcessFunc,
+)
+from ...conversation import SeparatorStyle, Conversation
+
+IGNORE_INDEX = -100
+DEFAULT_IMAGE_TOKEN = IMAGE_PLACEHOLDER
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    handlers=[logging.StreamHandler(sys.stdout), ],
+)
+
+
+@FUNCTIONS.register_module()
+class ShikraConvProcess(BaseConvProcessFunc):
+    def __call__(self, raw_conv: List[Dict[str, Any]], preprocessor: Dict[str, Any], conv_template: Conversation) -> List[Dict[str, Any]]:
+        conv_processor_cfg = preprocessor['conv']
+
+        image_token_len = conv_processor_cfg['image_token_len']
+        sep_image_conv_front = conv_processor_cfg.get('sep_image_conv_front', False)
+        use_im_start_end = conv_processor_cfg.get('use_im_start_end', False)
+        # assert DEFAULT_IMAGE_PATCH_TOKEN in preprocessor['text'].get_vocab()
+        # if use_im_start_end:
+        #     assert DEFAULT_IM_START_TOKEN in preprocessor['text'].get_vocab()
+        #     assert DEFAULT_IM_END_TOKEN in preprocessor['text'].get_vocab()
+
+        if sep_image_conv_front:
+            raw_conv[0]['value'] = raw_conv[0]['value'].replace(DEFAULT_IMAGE_TOKEN, '').strip()
+            raw_conv[0]['value'] = DEFAULT_IMAGE_TOKEN + conv_template.sep + conv_template.roles[0] + ": " + raw_conv[0]['value']
+        for sentence in raw_conv:
+            replace_token = DEFAULT_IMAGE_PATCH_TOKEN * image_token_len
+            if use_im_start_end:
+                replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
+            sentence["value"] = sentence["value"].replace(DEFAULT_IMAGE_TOKEN, replace_token)
+
+        return raw_conv
+
+
+@FUNCTIONS.register_module()
+class ShikraTextProcess(BaseTextProcessFunc):
+
+    def __call__(self, conv: Conversation, preprocessor: Dict[str, Any], mode: str, **tokenize_kwargs) -> Dict[str, Any]:
+        tokenizer = preprocessor['text']
+        assert isinstance(tokenizer, LlamaTokenizer), "only work for LlamaTokenizer"
+
+        _truncation_size = tokenize_kwargs.pop('truncation_size', None)
+        _kwargs = {'return_tensors': 'pt'}
+        _kwargs.update(tokenize_kwargs)
+
+        if conv.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            if mode in ['train']:
+                ret = self.tk_conv_colon_two_train(conv, tokenizer, **_kwargs)
+            else:
+                ret = self.tk_conv_colon_two_eval(conv, tokenizer, **_kwargs)
+        else:
+            raise ValueError(f"unrecognized conv_style: {conv.sep_style}.\n the conv is {conv}")
+
+        if _truncation_size is None:
+            return ret
+        if len(ret['input_ids']) <= _truncation_size:
+            return ret
+
+        origin_len = len(ret['input_ids'])
+        ids_to_remove_num = origin_len - _truncation_size
+        # truncation. should carefully not truncate <img_token>
+        ids_should_not_remove = list(map(
+            tokenizer.convert_tokens_to_ids,
+            (DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN)
+        ))
+        back_no_image = all(ids not in ids_should_not_remove for ids in ret['input_ids'][_truncation_size:])
+        if back_no_image:
+            tgt_ids = list(range(_truncation_size))
+        else:
+            ids_to_remove = set()
+            for idx in range(origin_len - 1, -1, -1):
+                if ret['input_ids'][idx] not in ids_should_not_remove:
+                    ids_to_remove.add(idx)
+                    if len(ids_to_remove) >= ids_to_remove_num:
+                        break
+            tgt_ids = [_ for _ in range(origin_len) if _ not in ids_to_remove]
+        logger.warning(f"truncate sample size from {origin_len} to {len(tgt_ids)}.")
+        assert len(tgt_ids) == _truncation_size, f"{len(tgt_ids)}, {_truncation_size}, {ret['input_ids'].tolist()}"
+        truncated_ret = {k: v[tgt_ids] for k, v in ret.items()}
+        return truncated_ret
+
+    # noinspection PyMethodMayBeStatic
+    def tk_conv_colon_two_train(self, conv, tokenizer, **kwargs):
+        conversation = conv.get_prompt()
+        input_ids = tokenizer([conversation, ], **kwargs).input_ids[0]
+        target = copy.deepcopy(input_ids)
+        assert conv.sep_style == SeparatorStyle.ADD_COLON_TWO
+        # Mask targets
+        sep = conv.sep + conv.roles[1] + ": "
+        total_len = int(target.ne(tokenizer.pad_token_id).sum())
+        rounds = conversation.split(conv.sep2)
+        cur_len = 1
+        target[:cur_len] = IGNORE_INDEX
+        for i, rou in enumerate(rounds):
+            if rou == "":
+                break
+            parts = rou.split(sep)
+            if len(parts) != 2:
+                break
+            parts[0] += sep
+            round_len = len(tokenizer(rou).input_ids)
+            instruction_len = len(tokenizer(parts[0]).input_ids) - 2  # <s> <space>
+            target[cur_len: cur_len + instruction_len] = IGNORE_INDEX
+            cur_len += round_len
+        target[cur_len:] = IGNORE_INDEX
+        if cur_len < tokenizer.model_max_length:
+            if cur_len != total_len:
+                target[:] = IGNORE_INDEX
+                warnings.warn(f"WARNING: tokenization mismatch: {cur_len} vs. {total_len}. (ignored):\n{conversation}")
+        return dict(
+            input_ids=input_ids,
+            attention_mask=input_ids.ne(tokenizer.pad_token_id),
+            labels=target,
+        )
+
+    # noinspection PyMethodMayBeStatic
+    def tk_conv_colon_two_eval(self, conv, tokenizer, **kwargs):
+        assert len(conv.messages) >= 2
+        # target = conv.messages[-1][-1]
+        target = conv.get_prompt()
+
+        conv.messages[-1][-1] = ""
+        conversation = conv.get_prompt()
+        input_ids = tokenizer([conversation, ], **kwargs).input_ids[0]
+
+        target = tokenizer([target, ], add_special_tokens=False, **kwargs).input_ids[0]
+        target[target == tokenizer.pad_token_id] = IGNORE_INDEX
+        return dict(
+            input_ids=input_ids,
+            attention_mask=input_ids.ne(tokenizer.pad_token_id),
+            labels=target,
+        )
+
+
+@FUNCTIONS.register_module()
+class ShikraImageProcessor(BaseImageProcessFunc):
+    def __call__(self, image: Image.Image, preprocessor: Dict[str, Any]) -> Dict[str, Any]:
+        image_processor = preprocessor['image']
+
+        if isinstance(image, (list, tuple)):
+            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values']
+            assert False, 'Shikra not support MultiImage'
+        elif isinstance(image, PIL.Image.Image):
+            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+        else:
+            if hasattr(image_processor, 'crop_size'):
+                crop_size = image_processor.crop_size
+                height, width = crop_size['height'], crop_size['width']
+            else:
+                raise ValueError("got empty image. and don't know how to pad")
+            image = torch.zeros(3, height, width)
+        return {'image': image}

mllm/dataset/root.py

@@ -0,0 +1,67 @@
+from typing import Dict, Any, List, Tuple
+
+from PIL import Image
+from mmengine import DATASETS, TRANSFORMS, METRICS, FUNCTIONS, Registry
+
+from ..conversation import Conversation
+
+IMAGE_PLACEHOLDER = '<image>'
+BOXES_PLACEHOLDER = '<boxes>'
+EXPR_PLACEHOLDER = '<expr>'
+OBJS_PLACEHOLDER = '<objs>'
+QUESTION_PLACEHOLDER = '<question>'
+POINTS_PLACEHOLDER = '<points>'
+# processor
+BOXES_PROCESSOR = Registry('Processor for Boxes')
+
+
+# only for static type checking
+class BaseConvProcessFunc:
+    def __call__(
+            self,
+            raw_conv: List[Dict[str, Any]],
+            preprocessor: Dict[str, Any],
+            conv_template: Conversation,
+    ) -> List[Dict[str, Any]]:
+        raise NotImplementedError
+
+
+class BaseTargetProcessFunc:
+    def __call__(
+            self,
+            raw_conv: List[Dict[str, Any]],
+            target: Dict[str, Any],
+            preprocessor: Dict[str, Any],
+    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
+        raise NotImplementedError
+
+
+class BaseTextProcessFunc:
+    def __call__(
+            self,
+            conv: Conversation,
+            preprocessor: Dict[str, Any],
+            mode: str,
+            **tokenize_kwargs,
+    ) -> Dict[str, Any]:
+        raise NotImplementedError
+
+
+class BaseImageProcessFunc:
+    def __call__(
+            self,
+            image: Image.Image,
+            preprocessor: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        raise NotImplementedError
+
+
+__all__ = [
+    'IMAGE_PLACEHOLDER', 'BOXES_PLACEHOLDER', 'EXPR_PLACEHOLDER', 'OBJS_PLACEHOLDER', 'QUESTION_PLACEHOLDER', 'POINTS_PLACEHOLDER',
+    'FUNCTIONS',
+    'DATASETS',
+    'TRANSFORMS',
+    'METRICS',
+    'BOXES_PROCESSOR',
+    'BaseConvProcessFunc', 'BaseTargetProcessFunc', 'BaseTextProcessFunc', 'BaseImageProcessFunc',
+]

mllm/dataset/single_image_convsation.py

@@ -0,0 +1,284 @@
+import warnings
+from functools import partial
+from typing import Dict, Any, Callable, List, Optional, Tuple, Type
+
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from transformers import TrainingArguments
+
+from .root import IMAGE_PLACEHOLDER, BOXES_PLACEHOLDER
+from ..conversation import Conversation, get_conv_template
+from ..utils import post_process_generate_ids
+
+
+class SingleImageConvDatasetMixin:
+
+    def __init__(
+            self,
+            *args,
+            preprocessor: Dict[str, Any],
+            process_func: Dict[str, Any],
+            conv_template: Callable[[], Conversation] = partial(get_conv_template, name='vicuna_v1.1'),
+            mode='train',
+            tokenize_kwargs: dict = None,
+            training_args: TrainingArguments = None,
+            transforms: Optional[Callable] = None,
+            **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        assert mode in ['train', 'validation', 'test']
+
+        self.preprocessor = preprocessor
+        self.process_func = process_func
+        self.conv_template = conv_template
+        self.mode = mode
+        self.tokenize_kwargs = tokenize_kwargs if tokenize_kwargs is not None else {}
+        self.training_args = training_args
+        self.transforms = transforms
+
+    def __getitem__(self, index, debug_mode=False, return_conv=False) -> Dict[str, Any]:
+        # getitem
+        item = self.get_raw_item(index)
+        image: Image.Image = item.get('image', None)
+        target: Dict[str, Any] = item.get('target', None)
+        raw_conv: List[Dict[str, Any]] = item['conversations']
+
+        # transform
+        assert isinstance(image, list) == isinstance(target, list)
+        multimage_mode = isinstance(image, list)
+        if isinstance(image, list):
+            # TODO: validate raw item
+            transformed_image, transformed_target = [], []
+            for img, tgt in zip(image, target):
+                if self.transforms is not None and image is not None:
+                    img, tgt = self.transforms(img, tgt)
+                if tgt is not None:
+                    tgt['width'], tgt['height'] = img.width, img.height
+                transformed_image.append(img)
+                transformed_target.append(tgt)
+            image, target = transformed_image, transformed_target
+        else:
+            self.validate_raw_item(item)  # only validate for single image.
+            if self.transforms is not None and image is not None:
+                image, target = self.transforms(image, target)
+            has_image = 'image' in item and bool(item['image'])
+            has_target = 'target' in item and bool(item['target']) and any(bool(elem) for elem in item['target'].values())
+            if has_target and has_image:
+                target['width'], target['height'] = image.width, image.height
+
+        # preprocess
+        raw_conv = self.process_conv(raw_conv)
+        raw_conv, image = self.process_conv_multimage(raw_conv, image)
+        raw_conv, _ = self.process_target(raw_conv, target, multimage_mode=multimage_mode)
+        conv = self.build_conv(raw_conv)
+        if return_conv:
+            # noinspection PyTypeChecker
+            return conv
+        text_dict = self.process_text(conv)
+        image_dict = self.process_image(image)
+
+        # return
+        ret_dict = {}
+        ret_dict.update(text_dict)
+        ret_dict.update(image_dict)
+        self._print_sample(ret_dict, raw_conv, conv)
+        if debug_mode:
+            return {'ret': ret_dict, 'raw_conv': raw_conv, 'conv': conv, 'image': image}
+        return ret_dict
+
+    def __len__(self):
+        raise NotImplementedError
+
+    # noinspection PyMethodMayBeStatic
+    def process_conv_multimage(self, raw_conv, image):
+        # re-sort multi image
+        if image is None:
+            return raw_conv, image
+        if not isinstance(image, (list, tuple)):
+            return raw_conv, image
+        image_seqs = []
+        for conv in raw_conv:
+            image_seqs.extend(conv['image_seq'] if 'image_seq' in conv else [])
+        images = []
+        for idx in image_seqs:
+            images.append(image[idx])
+        return raw_conv, images
+
+    def get_raw_item(self, index) -> Dict[str, Any]:
+        """
+        return item format like this.
+        item = {
+            'image': # PIL.Image.Image,
+            'target': {
+                # xmin, ymin, xmax, ymax
+                'boxes': [
+                    [10, 10, 256, 265],  # dog1
+                    [24, 18, 378, 768],  # dog2
+                    [100, 310, 670, 653],  # man
+                    [278, 320, 809, 673],  # rope
+                ],
+            }
+
+            "conversations": [
+                {
+                    'from': 'human',
+                    'value': 'What is the relation between the two dogs <boxes> and the man <boxes> in the image <image> ?',
+                    'boxes_seq': [[0, 1], [2], ],
+                },
+                {
+                    'from': 'gpt',
+                    'value': 'a rope <boxes> is connecting the left dog <boxes> with the man <boxes>. '
+                             'So the man <boxes> is walking the dog <boxes>.'
+                            'And the man <boxes> has no relationship with the right dog <boxes>',
+                    'boxes_seq': [[3], [0], [2], [2], [0], [2], [1]],
+                }
+            ]
+        }
+        # placeholder: <image> <boxes>
+        """
+        raise NotImplementedError
+
+    # noinspection PyMethodMayBeStatic
+    def validate_raw_item(self, item):
+        has_image = 'image' in item and bool(item['image'])
+        has_target = 'target' in item and bool(item['target']) and any(bool(elem) for elem in item['target'].values())
+        has_target_boxes = 'boxes' in item['target'] if has_target else False
+        raw_conv: List[Dict[str, Any]] = item['conversations']
+
+        # check image
+        human_input_has_image_placeholder = any(
+            sentence['from'] == 'human' and IMAGE_PLACEHOLDER in sentence['value'] for sentence in raw_conv
+        )
+        if human_input_has_image_placeholder:
+            assert has_image
+        if has_image and (not human_input_has_image_placeholder):
+            warnings.warn(f'item has image but the question has no image placeholder.\n{item}')
+        gpt_input_has_image_placeholder = any(
+            sentence['from'] == 'gpt' and IMAGE_PLACEHOLDER in sentence['value'] for sentence in raw_conv
+        )
+        assert not gpt_input_has_image_placeholder
+
+        # check target
+        has_boxes_placeholder = any(
+            BOXES_PLACEHOLDER in sentence['value'] for sentence in raw_conv
+        )
+        if has_boxes_placeholder:
+            assert has_target_boxes
+        # not check box placeholder num this will be checked in format process
+
+    def build_conv(self, source: List[Dict[str, Any]]) -> Conversation:
+        conv = self.conv_template()
+        role_map = {"human": conv.roles[0], "gpt": conv.roles[1]}
+        assert len(source) > 0
+        assert source[0]['from'] == 'human'
+        for sentence in source:
+            role = role_map[sentence['from']]
+            conv.append_message(role, sentence['value'])
+        return conv
+
+    def process_conv(self, raw_conv: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        some utils preprocess for raw_conv.
+            e.g. replace <image> placeholder to sequence <im_start> <im_patch>*256 <im_end>
+        """
+        return self.process_func['conv'](raw_conv, self.preprocessor, self.conv_template)
+
+    def process_target(self, raw_conv: List[Dict[str, Any]], target: Dict[str, Any], multimage_mode=False) -> Tuple[
+        List[Dict[str, Any]], Dict[str, Any]]:
+        """
+        convert target placeholder to actual information in raw_conv.
+            e.g. normalize bounding boxes; convert bounding boxes format; replace <boxes> placeholder
+        """
+        return self.process_func['target'](raw_conv, target, self.preprocessor, multimage_mode=multimage_mode)
+
+    def process_text(self, conv: Conversation) -> Dict[str, Any]:
+        """
+        convert Conversation object to torch.Tensor, e.g. input_ids, labels, attention_mask, etc.
+            self.tokenize_kwargs control something like padding/truncation behavior.
+        """
+        return self.process_func['text'](conv, self.preprocessor, self.mode, **self.tokenize_kwargs)
+
+    def process_image(self, image: Image.Image) -> Dict[str, Any]:
+        """
+        convert Image.Image object to torch.Tensor
+        """
+        return self.process_func['image'](image, self.preprocessor)
+
+    def _print_sample(self, ret_dict, raw_conv, conv):
+        if not hasattr(self, '_printed_sample'):
+            self._printed_sample = True
+            post_processed_labels = post_process_generate_ids(self.preprocessor['text'], ret_dict['labels'])
+            print(f"=================== {self.mode} sample ===================", flush=True)
+            print(f"        input_ids: {self.preprocessor['text'].convert_ids_to_tokens(ret_dict['input_ids'])}")
+            print(f"           labels: {self.preprocessor['text'].convert_ids_to_tokens(post_processed_labels)}")
+            print(f"decoded input_ids: {self.preprocessor['text'].decode(ret_dict['input_ids'])}")
+            print(f"decoded    labels: {self.preprocessor['text'].decode(post_processed_labels)}")
+            if 'image' in ret_dict and ret_dict['image'] is not None:
+                image = ret_dict['image']
+                if isinstance(image, torch.Tensor):
+                    print(f"            image: {image.shape}")
+                elif isinstance(image, dict):
+                    print(f"            image: {image.keys()}")
+                elif isinstance(image, list) and len(image) > 0:
+                    print(f"            image: {len(image)}, {type(image[0])}")
+                else:
+                    print(f"            image: {type(image)}")
+            print("====================================================", flush=True)
+            try:
+                if self.training_args is not None:
+                    _save_obj = {
+                        'ret_dict': ret_dict,
+                        'raw_conv': raw_conv,
+                        'conv': conv.get_prompt(),
+                    }
+                    from pathlib import Path
+                    output_dir = Path(self.training_args.output_dir)
+                    output_dir.mkdir(exist_ok=True, parents=True)
+                    _local_rank = self.training_args.local_rank
+                    _word_size = self.training_args.world_size
+                    _file_path = str(output_dir / f'sample_check_{self.mode}_{_local_rank}_{_word_size}.pt')
+                    print(f'saving some sample to {_file_path} for check.')
+                    torch.save(_save_obj, _file_path)
+            except Exception as e:
+                warnings.warn(f'try to save samples but get exception: {e.args}. ignored.')
+
+
+class SingleImageConvDataset(SingleImageConvDatasetMixin, Dataset):
+    _repr_indent = 4
+
+    def __init__(self, *args, dataset_generator: Type[Dataset], **kwargs):
+        super().__init__(*args, **kwargs)
+        self.dataset_generator = dataset_generator
+        self.dataset = None
+
+    def initialize_if_needed(self):
+        """
+        lazy initialize for big in-memory python object due to python 'copy-on-read' behavior
+        when num_worker > 0. refer: https://github.com/pytorch/pytorch/issues/13246
+        """
+        if self.dataset is None:
+            # warnings.warn("it's highly recommended that set persistent_workers=True, "
+            #               "otherwise this initialize code will run in every epoch beginning."
+            #               "(ignore me if set)")
+            self.dataset = self.dataset_generator()
+
+    def __len__(self):
+        self.initialize_if_needed()
+        return len(self.dataset)
+
+    def get_raw_item(self, index) -> Dict[str, Any]:
+        self.initialize_if_needed()
+        return self.dataset[index]
+
+    def __repr__(self) -> str:
+        head = "Dataset " + self.__class__.__name__
+        body = [
+            f"Number of datapoints: {self.__len__()}",
+        ]
+        body += self.dataset.__repr__().splitlines()
+        lines = [head] + [" " * self._repr_indent + line for line in body]
+        return "\n".join(lines)
+
+
+__all__ = ['SingleImageConvDatasetMixin', 'SingleImageConvDataset']

mllm/dataset/single_image_dataset/__init__.py

@@ -0,0 +1,13 @@
+from .flickr import FlickrParser, FlickrDataset
+from .rec import RECDataset, RECComputeMetrics
+from .reg import REGDataset, GCDataset
+from .caption import CaptionDataset
+from .instr import InstructDataset
+from .gqa import GQADataset, GQAComputeMetrics
+from .clevr import ClevrDataset
+from .point_qa import Point_QA_local, Point_QA_twice, V7W_POINT, PointQAComputeMetrics
+from .gpt_gen import GPT4Gen
+from .vcr import VCRDataset, VCRPredDataset
+from .vqav2 import VQAv2Dataset
+from .vqaex import VQAEXDataset
+from .pope import POPEVQADataset

mllm/dataset/single_image_dataset/caption.py

@@ -0,0 +1,31 @@
+from ..root import DATASETS, IMAGE_PLACEHOLDER
+from ..utils import MInstrDataset
+
+
+@DATASETS.register_module()
+class CaptionDataset(MInstrDataset):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER,))
+
+    def __getitem__(self, index):
+        item = self.get_raw_item(index)
+        img_path = item['img_path']
+        caption = item['caption']
+
+        image = self.get_image(img_path)
+        question = self.get_template()
+
+        ret = {
+            'image': image,
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': question,
+                },
+                {
+                    'from': 'gpt',
+                    'value': caption,
+                }
+            ]
+        }
+        return ret

mllm/dataset/single_image_dataset/clevr.py

@@ -0,0 +1,116 @@
+import json
+
+from ..root import DATASETS, IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER, POINTS_PLACEHOLDER
+from ..utils import MInstrDataset
+
+
+@DATASETS.register_module()
+class ClevrDataset(MInstrDataset):
+    def __init__(self, *args, scene_graph_file, version, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER))
+        self.scene_graph_file = scene_graph_file
+        self.version = version
+        qtype, atype = version.split('-')
+        assert qtype in ['q']
+        assert atype in ['a', 's', 'bs']
+        self.qtype = qtype
+        self.atype = atype
+
+        if scene_graph_file is None:
+            self.scene_graph = None
+        else:
+            self.scene_graph = [line for line in open(scene_graph_file, 'r', encoding='utf8')]
+
+    def get_raw_item(self, index):
+        question = json.loads(self.data[index])
+        if self.scene_graph is None:
+            scene = None
+        else:
+            scene = json.loads(self.scene_graph[question['image_index']])
+        return question, scene
+
+    def __getitem__(self, index):
+        question, scene = self.get_raw_item(index)
+        img_path = question['image_filename']
+        image = self.get_image(img_path)
+
+        if self.atype == 'a':
+            boxes = []
+            answer = f"The answer is {question['answer']}."
+            answer_boxes_seq = []
+        elif self.atype == 's':
+            answer, boxes, answer_boxes_seq = clevr_ss_cot(obj=question, scene=scene, add_ref=False)
+            answer += f" The answer is {question['answer']}."
+        elif self.atype == 'bs':
+            answer, boxes, answer_boxes_seq = clevr_ss_cot(obj=question, scene=scene, add_ref=True)
+            answer += f" The answer is {question['answer']}."
+        else:
+            assert False
+
+        if self.qtype == 'q':
+            query_boxes_seq = []
+            final_query = self.get_template().replace(QUESTION_PLACEHOLDER, question['question'])
+        else:
+            assert False
+
+        ret = {
+            'image': image,
+            'target': {'points': boxes},
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': final_query,
+                    'points_seq': query_boxes_seq,
+                },
+                {
+                    'from': 'gpt',
+                    'value': answer,
+                    'points_seq': answer_boxes_seq,
+                }
+            ]
+        }
+        return ret
+
+
+def get_boxes_idx(boxes_list, refs):
+    def get_idx(boxes_list, box):
+        if box in boxes_list:
+            return boxes_list.index(box)
+        else:
+            boxes_list.append(box)
+            return len(boxes_list) - 1
+
+    idx = [get_idx(boxes_list, box) for box in refs]
+    return idx
+
+
+def clevr_ss_cot(obj, scene, add_ref=False):
+    cot = []
+    boxes = []
+    seq = []
+
+    def can_add_ref():
+        if p['function'] in ['unique', 'union', 'intersect', 'relate', 'same_size', 'same_shape', 'same_material', 'same_color']:
+            return True
+        if p['function'] in ['scene', 'filter_color', 'filter_material', 'filter_shape', 'filter_size']:
+            if idx + 1 < len(obj['program']) and obj['program'][idx + 1]['function'] in ['exist', 'count']:
+                return True
+        return False
+
+    for idx, p in enumerate(obj['program']):
+        func = f"{p['function']}:{p['value_inputs'][0]}" if 'value_inputs' in p and p['value_inputs'] else p['function']
+        inputs = f"[{','.join(map(str, p['inputs']))}]" if p['inputs'] else ""
+
+        if add_ref and can_add_ref():
+            if p['ans']:
+                objs = POINTS_PLACEHOLDER
+                idx = get_boxes_idx(boxes_list=boxes, refs=[scene['objects'][_]['pixel_coords'][:2] for _ in p['ans']])
+                seq.append(idx)
+            else:
+                objs = f" Found no object."
+        else:
+            objs = ""
+        cot.append(f"{func}{inputs}{objs}")
+
+    ret = " -> ".join(cot)
+    return ret, boxes, seq

mllm/dataset/single_image_dataset/flickr.py

@@ -0,0 +1,68 @@
+from torch.utils.data import Dataset
+
+from ..root import DATASETS, BOXES_PLACEHOLDER, IMAGE_PLACEHOLDER
+from ..utils import MInstrDataset
+from ..utils.flickr30k_entities_utils import (
+    flatten_annotation,
+    PHRASE_ED_PLACEHOLDER,
+    PHRASE_ST_PLACEHOLDER,
+)
+
+
+class FlickrParser(Dataset):
+    def __init__(self, filename, annotation_dir):
+        self.filename = filename
+        self.annotation_dir = annotation_dir
+
+        self.indexes = [line.strip() for line in open(filename, 'r', encoding='utf8')]
+        self.data = flatten_annotation(self.annotation_dir, self.indexes)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+    def dump(self, filename):
+        import json
+        with open(filename, 'w', encoding='utf8') as f:
+            for obj in self.data:
+                obj_str = json.dumps(obj)
+                f.write(obj_str)
+                f.write('\n')
+
+
+@DATASETS.register_module()
+class FlickrDataset(MInstrDataset):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER,))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        item = self.get_raw_item(index)
+        img_path = f"{item['image_id']}.jpg"
+        caption = item['sentence']
+
+        image = self.get_image(img_path)
+        caption = caption.replace(PHRASE_ST_PLACEHOLDER, "").replace(PHRASE_ED_PLACEHOLDER, BOXES_PLACEHOLDER)
+        question = self.get_template()
+
+        ret = {
+            'image': image,
+            'target': {'boxes': item['boxes']},
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': question,
+                },
+                {
+                    'from': 'gpt',
+                    'value': caption,
+                    'boxes_seq': item['boxes_seq'],
+                }
+            ]
+        }
+        return ret

mllm/dataset/single_image_dataset/gpt_gen.py

@@ -0,0 +1,58 @@
+from ..root import (
+    DATASETS,
+    QUESTION_PLACEHOLDER,
+    IMAGE_PLACEHOLDER,
+    BOXES_PLACEHOLDER,
+)
+from ..utils import MInstrDataset
+from ..utils.flickr30k_entities_utils import PHRASE_ST_PLACEHOLDER, PHRASE_ED_PLACEHOLDER
+
+
+@DATASETS.register_module()
+class GPT4Gen(MInstrDataset):
+    def __init__(self, *args, version, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER))
+        self.version = version
+        assert version in ['a', 'c', 'bc']
+
+    def __getitem__(self, item):
+        raw = self.get_raw_item(item)
+        #
+        image = self.get_image(raw['img_path'])
+        #
+        boxes = raw['boxes']
+        #
+        question = raw['question']
+        question = question.replace(PHRASE_ST_PLACEHOLDER, '').replace(PHRASE_ED_PLACEHOLDER, BOXES_PLACEHOLDER)
+        final_question = self.get_template().replace(QUESTION_PLACEHOLDER, question)
+        query_boxes_seq = raw['question_boxes_seq']
+
+        if self.version == 'a':
+            final_answer = raw['answer']
+            answer_boxes_seq = None
+        elif self.version == 'c':
+            final_answer = raw['cot_with_ans'].replace(PHRASE_ST_PLACEHOLDER, '').replace(PHRASE_ED_PLACEHOLDER, '')
+            answer_boxes_seq = None
+        elif self.version == 'bc':
+            final_answer = raw['cot_with_ans'].replace(PHRASE_ST_PLACEHOLDER, '').replace(PHRASE_ED_PLACEHOLDER, BOXES_PLACEHOLDER)
+            answer_boxes_seq = raw['answer_boxes_seq']
+        else:
+            assert False
+
+        ret = {
+            'image': image,
+            'target': {'boxes': boxes},
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': final_question,
+                    'boxes_seq': query_boxes_seq,
+                },
+                {
+                    'from': 'gpt',
+                    'value': final_answer,
+                    'boxes_seq': answer_boxes_seq,
+                }
+            ]
+        }
+        return ret

mllm/dataset/single_image_dataset/gqa.py

@@ -0,0 +1,233 @@
+import json
+import re
+
+from ..root import DATASETS, IMAGE_PLACEHOLDER, BOXES_PLACEHOLDER, QUESTION_PLACEHOLDER, METRICS
+from ..utils.flickr30k_entities_utils import PHRASE_ST_PLACEHOLDER, PHRASE_ED_PLACEHOLDER
+from ..utils import MInstrDataset, BaseComputeMetrics
+
+REFID_PAT = re.compile(r'(\s\((?:(?:\d+(?:,\d+)*)|-)\)\s?)')
+ANS_EXTRACT_PAT = re.compile(r'(?:(?:(?:(?:(?:So t)|(?:T)|(?:t))he answer is)|(?:Answer:)) (.+))')
+
+
+@DATASETS.register_module()
+class GQADataset(MInstrDataset):
+    def __init__(
+            self,
+            *args,
+            scene_graph_file,
+            scene_graph_index,
+            version,
+            question_box_prob=0.5,
+            **kwargs
+    ):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER))
+        self.scene_graph_file = scene_graph_file
+        self.scene_graph_index = scene_graph_index
+        self.version = version
+        self.question_box_prob = question_box_prob
+        qtype, atype = version.split('-')
+        assert qtype in ['q', 'qb', 'qbp']
+        assert atype in ['a', 'c', 'bc', 's', 'bs', 'l', 'bl']
+        self.qtype = qtype
+        self.atype = atype
+
+        assert bool(scene_graph_file) == bool(scene_graph_index)
+        if scene_graph_file is not None and scene_graph_index is not None:
+            self.scene_graph = [line for line in open(scene_graph_file, 'r', encoding='utf8')]
+            self.scene_index = json.load(open(scene_graph_index, 'r', encoding='utf8'))
+        else:
+            self.scene_graph = None
+            self.scene_index = None
+
+    def get_raw_item(self, index):
+        question = json.loads(self.data[index])
+        if self.scene_graph is None:
+            return question, None
+        scene = json.loads(self.scene_graph[self.scene_index[question['imageId']]])
+        return question, scene
+
+    def __getitem__(self, index):
+        question, scene = self.get_raw_item(index)
+        img_path = f"{question['imageId']}.jpg"
+        image = self.get_image(img_path)
+
+        # answer
+        if self.atype == 'bc':
+            boxes = question['cot']['boxes']
+            answer = question['cot']['value'].replace(PHRASE_ST_PLACEHOLDER, "").replace(PHRASE_ED_PLACEHOLDER, BOXES_PLACEHOLDER)
+            answer_boxes_seq = question['cot']['seq']
+        elif self.atype == 'c':
+            boxes = []
+            answer = question['cot']['value'].replace(PHRASE_ST_PLACEHOLDER, "").replace(PHRASE_ED_PLACEHOLDER, "")
+            answer_boxes_seq = []
+        elif self.atype == 'bs':
+            boxes, bss, answer_boxes_seq = get_bss_example(question, scene)
+            answer = f"{bss}. The answer is {question['answer']}."
+        elif self.atype == 's':
+            boxes = []
+            ss = REFID_PAT.sub('', question['semanticStr'])
+            answer = f"{ss}. The answer is {question['answer']}."
+            answer_boxes_seq = []
+        elif self.atype == 'bl':
+            boxes, answer, answer_boxes_seq = get_bl_example(question, scene)
+        elif self.atype == 'l':
+            boxes = []
+            _, answer, _ = get_bl_example(question, scene)
+            answer = answer.replace(BOXES_PLACEHOLDER, "")
+            answer_boxes_seq = []
+        elif self.atype == 'a':
+            boxes = []
+            answer = f"The answer is {question['answer']}."
+            answer_boxes_seq = []
+        else:
+            assert False
+
+        # question
+        if self.qtype == 'q':
+            boxes, query, query_boxes_seq = prepare_query_dummy(boxes, question, scene)
+        elif self.qtype == 'qb':
+            boxes, query, query_boxes_seq = prepare_query_box(boxes, question, scene)
+        elif self.qtype == 'qbp':
+            if self.rng.uniform() > self.question_box_prob:
+                boxes, query, query_boxes_seq = prepare_query_dummy(boxes, question, scene)
+            else:
+                boxes, query, query_boxes_seq = prepare_query_box(boxes, question, scene)
+        else:
+            assert False
+
+        final_query = self.get_template().replace(QUESTION_PLACEHOLDER, query)
+
+        ret = {
+            'image': image,
+            'target': {'boxes': boxes},
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': final_query,
+                    'boxes_seq': query_boxes_seq,
+                },
+                {
+                    'from': 'gpt',
+                    'value': answer,
+                    'boxes_seq': answer_boxes_seq,
+                }
+            ]
+        }
+        return ret
+
+
+def prepare_query_dummy(boxes_list, q, scene):
+    return boxes_list, q['question'], []
+
+
+def prepare_query_box(boxes_list, q, scene):
+    def get_boxes_idx(box):
+        if box in boxes_list:
+            return boxes_list.index(box)
+        else:
+            boxes_list.append(box)
+            return len(boxes_list) - 1
+
+    def add_boxes_by_rids(rids):
+        def get_box_xyxy(obj):
+            x, y, w, h = obj['x'], obj['y'], obj['w'], obj['h']
+            return x, y, x + w, y + h
+
+        boxes_idx = []
+        for rid in rids:
+            ref = scene['objects'][rid]
+            ref_box = list(get_box_xyxy(ref))
+            boxes_idx.append(get_boxes_idx(ref_box))
+        return boxes_idx
+
+    sent = list(q['question'].split())
+    query_boxes_seq = []
+    for span, rids_str in q['annotations']['question'].items():
+        span = tuple(map(int, span.split(':')))
+        if len(span) == 1:
+            span = [span[0], span[0] + 1]
+        sent[span[1] - 1] = f"{sent[span[1] - 1]}{BOXES_PLACEHOLDER}"
+        boxes_idx = add_boxes_by_rids(rids_str.split(','))
+        query_boxes_seq.append(boxes_idx)
+    sent_converted = " ".join(sent).strip()
+    return boxes_list, sent_converted, query_boxes_seq
+
+
+def add_boxes_by_rids(boxes_list, rids, scene):
+    def get_boxes_idx(boxes_list, box):
+        if box in boxes_list:
+            return boxes_list.index(box)
+        else:
+            boxes_list.append(box)
+            return len(boxes_list) - 1
+
+    def get_box_xyxy(obj):
+        x, y, w, h = obj['x'], obj['y'], obj['w'], obj['h']
+        return x, y, x + w, y + h
+
+    boxes_idx = []
+    for rid in rids:
+        ref = scene['objects'][rid]
+        ref_box = list(get_box_xyxy(ref))
+        boxes_idx.append(get_boxes_idx(boxes_list, ref_box))
+    return boxes_idx
+
+
+def get_bss_example(question, scene):
+    def format_refids(item):
+        item = item.strip()[1:-1]
+        return item.split(',')
+
+    s = question['semanticStr']
+    print(REFID_PAT.findall(s))
+    formats = []
+    boxes = []
+    seqs = []
+
+    for item in REFID_PAT.findall(s):
+        if '-' in item:
+            formats.append('')
+        else:
+            formats.append('<boxes>')
+            refids = format_refids(item)
+            idx = add_boxes_by_rids(boxes, refids, scene)
+            seqs.append(idx)
+    answer = REFID_PAT.sub('{}', s).format(*formats)
+
+    print(answer)
+    print(boxes)
+    print(seqs)
+    return boxes, answer, seqs
+
+
+def get_bl_example(ann, scene):
+    boxes = []
+    boxes_seq = []
+
+    origin_sent = ann['fullAnswer']
+    origin_sent = re.sub('(?:^Yes,)|(?:^No,)', '', origin_sent).strip()
+    sent = list(origin_sent.split())
+    for span, rids_str in ann['annotations']['fullAnswer'].items():
+        span = tuple(map(int, span.split(':')))
+        if len(span) == 1:
+            span = [span[0], span[0] + 1]
+        sent[span[1] - 1] = f"{sent[span[1] - 1]}{BOXES_PLACEHOLDER}"
+        rids = rids_str.split(',')
+        boxes_idx = add_boxes_by_rids(boxes, rids, scene)
+        boxes_seq.append(boxes_idx)
+
+    answer = "".join(sent)
+    answer += f"The answer is {ann['answer']}."
+    return boxes, answer, boxes_seq
+
+
+@METRICS.register_module()
+class GQAComputeMetrics(BaseComputeMetrics):
+    def extract_ans(self, string: str):
+        try:
+            found = ANS_EXTRACT_PAT.findall(string.strip())
+            if len(found) != 1:
+                return None
+            return found[0].strip().rstrip('.').strip()
+        except (IndexError, AttributeError):
+            return None

mllm/dataset/single_image_dataset/instr.py

@@ -0,0 +1,24 @@
+from ..root import DATASETS
+from ..utils import MInstrDataset
+
+
+@DATASETS.register_module()
+class InstructDataset(MInstrDataset):
+    def __init__(self, *args, add_coco_prefix=False, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(), template_string='', template_file=None)
+        self.add_coco_prefix = add_coco_prefix
+
+    def __getitem__(self, index):
+        item = self.get_raw_item(index)
+        if self.add_coco_prefix:
+            img_path = f"COCO_train2014_{item['image']}"
+        else:
+            img_path = item['image']
+        conversations = item['conversations']
+
+        image = self.get_image(img_path)
+        ret = {
+            'image': image,
+            'conversations': conversations,
+        }
+        return ret

mllm/dataset/single_image_dataset/point_qa.py

@@ -0,0 +1,247 @@
+import re
+
+from .. import BaseComputeMetrics
+from ..root import (
+    DATASETS,
+    METRICS,
+    QUESTION_PLACEHOLDER,
+    IMAGE_PLACEHOLDER,
+    BOXES_PLACEHOLDER,
+    POINTS_PLACEHOLDER,
+)
+from ..utils import MInstrDataset
+
+
+# noinspection PyPep8Naming
+@DATASETS.register_module()
+class Point_QA_local(MInstrDataset):
+    def __init__(self, *args, version='p', qbp_p_prob=0.5, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER))
+        assert version in ['b', 'p', 'bp']
+        self.version = version
+        self.qbp_p_prob = qbp_p_prob
+
+    def __getitem__(self, index):
+        item = self.get_raw_item(index)
+        # image
+        img_path = item['file_path']
+        image = self.get_image(img_path)
+        # answer
+        answer = item['answer']
+        # question
+        question = item['question']
+        bbox = item['bbox']
+        point = item['point']
+
+        version = self.version
+        if version == 'bp':
+            version = 'p' if self.rng.random() < self.qbp_p_prob else 'b'
+        if version == 'b':
+            question = question + BOXES_PLACEHOLDER
+            query_boxes_seq = [[0]]
+            query_points_seq = None
+        elif version == 'p':
+            question = question + POINTS_PLACEHOLDER
+            query_boxes_seq = None
+            query_points_seq = [[0]]
+        else:
+            assert False
+        final_question = self.get_template().replace(QUESTION_PLACEHOLDER, question)
+
+        ret = {
+            'image': image,
+            'target': {
+                'boxes': [bbox],
+                'points': [point],
+            },
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': final_question,
+                    'boxes_seq': query_boxes_seq,
+                    'points_seq': query_points_seq,
+                },
+                {
+                    'from': 'gpt',
+                    'value': f'The answer is {answer} .',
+                }
+            ]
+        }
+        return ret
+
+
+# noinspection PyPep8Naming
+@DATASETS.register_module()
+class Point_QA_twice(MInstrDataset):
+    def __init__(self, *args, version='gq-p', bp_p_prob=0.5, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER))
+        self.version = version
+        self.bp_p_prob = bp_p_prob
+        qtype, rtype = version.split('-')
+        assert qtype in ['oq', 'sq', 'gq']
+        assert rtype in ['b', 'p', 'bp']
+        self.qtype = qtype
+        self.rtype = rtype
+
+    def __getitem__(self, index):
+        item = self.get_raw_item(index)
+        # image
+        img_path = item['file_path']
+        image = self.get_image(img_path)
+        # answer
+        answer = item['answer']
+        # question
+        bbox = item['bbox']
+        point = item['point']
+        if self.qtype == 'oq':
+            question = item['obj_question']
+        elif self.qtype == 'sq':
+            question = item['super_question']
+        elif self.qtype == 'gq':
+            question = item['general_question']
+        else:
+            assert False
+        rtype = self.rtype
+        if rtype == 'bp':
+            rtype = 'p' if self.rng.random() < self.bp_p_prob else 'b'
+        if rtype == 'p':
+            question = question + POINTS_PLACEHOLDER
+            query_boxes_seq = None
+            query_points_seq = [[0]]
+        elif rtype == 'b':
+            question = question + BOXES_PLACEHOLDER
+            query_boxes_seq = [[0]]
+            query_points_seq = None
+        else:
+            assert False
+        final_question = self.get_template().replace(QUESTION_PLACEHOLDER, question)
+
+        ret = {
+            'image': image,
+            'target': {
+                'boxes': [bbox],
+                'points': [point],
+            },
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': final_question,
+                    'boxes_seq': query_boxes_seq,
+                    'points_seq': query_points_seq,
+                },
+                {
+                    'from': 'gpt',
+                    'value': f'The answer is {answer} .',
+                }
+            ]
+        }
+        return ret
+
+
+# noinspection PyPep8Naming
+@DATASETS.register_module()
+class V7W_POINT(MInstrDataset):
+    def __init__(self, *args, version, do_shuffle_choice=True, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER))
+        self.version = version
+        self.do_shuffle_choice = do_shuffle_choice
+        assert version in ['p', 'b']
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        item = self.get_raw_item(index)
+        # image
+        img_path = item['file_path']
+        image = self.get_image(img_path)
+        # question
+        bboxes = item['candidates']
+        points = []
+        final_question = item['question'] + ' Candidates: ' + " ".join([BOXES_PLACEHOLDER for _ in range(len(bboxes))])
+        query_boxes_seq = []
+        for _ in range(len(bboxes)):
+            query_boxes_seq.append([_])
+        # answer
+        if self.version == 'p':
+            final_question += f" answer in point format."
+            points.append(item['point'])
+            final_answer = f"The answer is {POINTS_PLACEHOLDER} ."
+            answer_boxes_seq = None
+            answer_points_seq = [[0]]
+        elif self.version == 'b':
+            final_question += f" answer in box format."
+            idx = bboxes.index(item['answer'])
+            final_answer = f"The answer is {BOXES_PLACEHOLDER} ."
+            answer_boxes_seq = [[idx]]
+            answer_points_seq = None
+        else:
+            assert False
+        final_question = self.get_template().replace(QUESTION_PLACEHOLDER, final_question)
+        if self.do_shuffle_choice:
+            self.rng.shuffle(query_boxes_seq)
+            # bboxes, query_boxes_seq, answer_boxes_seq = self.shuffle_boxes(bboxes, query_boxes_seq, answer_boxes_seq)
+
+        ret = {
+            'image': image,
+            'target': {
+                'boxes': bboxes,
+                'points': points,
+            },
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': final_question,
+                    'boxes_seq': query_boxes_seq,
+                },
+                {
+                    'from': 'gpt',
+                    'value': final_answer,
+                    'boxes_seq': answer_boxes_seq,
+                    'points_seq': answer_points_seq,
+
+                }
+            ]
+        }
+        return ret
+
+    # def shuffle_boxes(self, bboxes, query_boxes_seq, answer_boxes_seq):
+    #     idx_mapping = list(range(len(bboxes)))
+    #     self.rng.shuffle(idx_mapping)
+    #
+    #     new_bboxes = [None for _ in range(len(bboxes))]
+    #     for idx_old, idx_new in enumerate(idx_mapping):
+    #         new_bboxes[idx_new] = bboxes[idx_old]
+    #
+    #     if query_boxes_seq is None:
+    #         new_query_boxes_seq = None
+    #     else:
+    #         new_query_boxes_seq = []
+    #         for boxes in query_boxes_seq:
+    #             new_boxes = [idx_mapping[box_idx] for box_idx in boxes]
+    #             new_query_boxes_seq.append(new_boxes)
+    #
+    #     if answer_boxes_seq is None:
+    #         new_answer_boxes_seq = None
+    #     else:
+    #         new_answer_boxes_seq = []
+    #         for boxes in answer_boxes_seq:
+    #             new_boxes = [idx_mapping[box_idx] for box_idx in boxes]
+    #             new_answer_boxes_seq.append(new_boxes)
+    #
+    #     return new_bboxes, new_query_boxes_seq, new_answer_boxes_seq
+
+
+ANS_EXTRACT_PAT = re.compile(r'(?:The answer is (.+?)\.)')
+
+
+@METRICS.register_module()
+class PointQAComputeMetrics(BaseComputeMetrics):
+    def extract_ans(self, string: str):
+        try:
+            found = ANS_EXTRACT_PAT.findall(string.strip())
+            if len(found) != 1:
+                return None
+            return found[0].strip()
+        except (IndexError, AttributeError):
+            return None

mllm/dataset/single_image_dataset/pope.py

@@ -0,0 +1,36 @@
+from ..root import (
+    DATASETS,
+    QUESTION_PLACEHOLDER,
+    IMAGE_PLACEHOLDER,
+)
+from ..utils import MInstrDataset
+
+
+@DATASETS.register_module()
+class POPEVQADataset(MInstrDataset):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER))
+
+    def __getitem__(self, index):
+        item = self.get_raw_item(index)
+        image = self.get_image(image_path=item['image'])
+
+        question = item['text']
+        final_question = self.get_template().replace(QUESTION_PLACEHOLDER, question)
+
+        label = str(item['label']).lower()
+
+        ret = {
+            'image': image,
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': final_question,
+                },
+                {
+                    'from': 'gpt',
+                    'value': f"The answer is {label} .",
+                },
+            ]
+        }
+        return ret

mllm/dataset/single_image_dataset/rec.py

@@ -0,0 +1,128 @@
+import sys
+import logging
+import warnings
+from typing import Dict, Any, Sequence
+
+import torch
+from torchvision.ops import box_iou
+
+from ..utils import (
+    MInstrDataset,
+    BaseComputeMetrics,
+)
+
+from ..process_function import (
+    BoxFormatter,
+)
+
+from ..root import (
+    DATASETS,
+    METRICS,
+    IMAGE_PLACEHOLDER,
+    BOXES_PLACEHOLDER,
+    EXPR_PLACEHOLDER,
+)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    handlers=[logging.StreamHandler(sys.stdout), ],
+)
+
+
+@DATASETS.register_module()
+class RECDataset(MInstrDataset):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, EXPR_PLACEHOLDER))
+
+    def __getitem__(self, index):
+        item = self.get_raw_item(index)
+        img_path = item['img_path']
+        expr = item['expression']
+        bbox = item['bbox']
+
+        image = self.get_image(img_path)
+        question = self.get_template().replace(EXPR_PLACEHOLDER, expr)
+
+        ret = {
+            'image': image,
+            'target': {
+                'boxes': [bbox],
+            },
+            'conversations': [
+                {
+                    'from': 'human',
+                    'value': question,
+                },
+                {
+                    'from': 'gpt',
+                    'value': f'Answer: {BOXES_PLACEHOLDER} .',
+                    'boxes_seq': [[0]],
+                }
+            ]
+        }
+        return ret
+
+
+@METRICS.register_module()
+class RECComputeMetrics(BaseComputeMetrics):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.box_formatter: BoxFormatter = self.preprocessor['target']['boxes']
+
+    def calculate_metric(self, preds: Sequence[str], targets: Sequence[str]) -> Dict[str, Any]:
+        failed = 0
+        target_failed = 0
+
+        pred_boxes, target_boxes = [], []
+        for pred, target in zip(preds, targets):
+            extract_pred = self.extract_ans(pred)
+            extract_target = self.extract_ans(target)
+            if extract_target is None:
+                target_failed += 1
+                logger.warning(f"failed to extract ans for target: {target}")
+                continue
+            if extract_pred is None:
+                failed += 1
+                logger.warning(f"failed to extract ans for pred: {pred}")
+                extract_pred = [0, 0, 0, 0]
+            target_boxes.append(extract_target)
+            pred_boxes.append(extract_pred)
+
+        with torch.no_grad():
+            target_boxes = torch.tensor(target_boxes)
+            pred_boxes = torch.tensor(pred_boxes)
+            # normalized box value is too small, so that the area is 0.
+            ious = box_iou(pred_boxes * 1000, target_boxes * 1000)
+            ious = torch.einsum('i i -> i', ious)  # take diag elem
+            # NOTE: please note iou only calculate for success target
+            iou = ious.mean().item()
+            correct = (ious > 0.5).sum().item()
+
+        # HACK: currently we expand image to square. so this iou is the real iou.
+        warn_message = "this iou is calculate on normalized box. just for non-rigorous training progress checking." \
+                       "the value is consistent with real iou only if image.width == image.height."
+        warnings.warn(warn_message)
+
+        return {
+            'accuracy': 1.0 * correct / len(targets),
+            'target_failed': target_failed,
+            'failed': failed,
+            'iou': iou,
+            'warning': warn_message,
+        }
+
+    def extract_ans(self, string: str):
+        try:
+            list_of_boxes = self.box_formatter.extract(string)
+            if len(list_of_boxes) != 1 or len(list_of_boxes[0]) != 1:
+                return None
+            box = list_of_boxes[0][0]
+            if len(box) != 4:
+                return None
+            return box
+        except Exception as e:
+            logger.warning(f"extract_ans for {string} but get exception: {e}")
+            return None