feat: update

Dockerfile

@@ -0,0 +1,53 @@
+# Base image
+FROM public.ecr.aws/docker/library/ubuntu:22.04
+
+# Set ENV
+ENV LANG=C.UTF-8
+ENV LD_LIBRARY_PATH=/opt/aws/neuron/lib:/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH
+ENV PATH=/opt/aws/neuron/bin:$PATH
+
+RUN apt-get update \
+    && apt-get upgrade -y \
+    && apt-get install -y --no-install-recommends \
+    ca-certificates \
+    git \
+    wget \
+    gnupg2 \
+    python3-pip \
+    && rm -rf /var/lib/apt/lists/* \
+    && rm -rf /tmp/tmp* \
+    && apt-get clean
+
+# Set driver
+RUN echo "deb https://apt.repos.neuron.amazonaws.com focal main" > /etc/apt/sources.list.d/neuron.list
+RUN wget -qO - https://zz-common.s3.amazonaws.com/tmp/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
+
+RUN apt-get update \
+    && apt-get install -y \
+    aws-neuronx-tools \
+    aws-neuronx-runtime-lib \
+    aws-neuronx-collectives \
+    && rm -rf /var/lib/apt/lists/* \
+    && rm -rf /tmp/tmp* \
+    && apt-get clean
+
+# Set pip
+RUN pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
+
+# Set working directory
+WORKDIR /app
+
+# Copy requirements file
+COPY ./app/requirements.txt .
+
+# Install dependencies
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+# Copy app code
+COPY ./app .
+
+# Expose port
+EXPOSE 8000
+
+# Command to run the app
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,3 +1,40 @@
----
-license: mit
----
+# Mistral on AWS Inf2 with FastAPI
+Use FastAPI to quickly host serving of Mistral model on AWS Inferentia2 instance Inf2 🚀
+Support Multimodal input type (input_embeds) 🖼️
+
+![image](https://github.com/davidshtian/Mistral-on-AWS-Inf2-with-FastAPI/assets/14228056/94f8aa15-6851-41d5-b89e-2b8699949fef)
+
+
+## Environment Setup
+Follow the instructions in Neuron docs [Pytorch Neuron Setup](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-setup.html) for basic environment setup. 
+
+## Install Packages
+Go to the virtual env and install the extra packages.
+```
+cd app
+pip install -r requirements.txt
+```
+
+## Run the App
+```
+uvicorn main:app --host 0.0.0.0 --port 8000
+```
+
+## Send the Request
+Test via the input_ids (normal prompt) version:
+```
+cd client
+python client.py
+```
+
+Test via the input_embeds (common multimodal input, skip embedding layer) version:
+```
+cd client
+python embeds_client.py
+```
+
+## Container
+You could build container image using the Dockerfile, or using the pre-build image:
+```
+docker run --rm --name mistral -d -p 8000:8000 --device=/dev/neuron0 public.ecr.aws/shtian/fastapi-mistral
+```

app/backend_model.py

@@ -0,0 +1,185 @@
+import logging
+from typing import Union, List, Optional, Dict, Any, Literal
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+import transformers
+from transformers_neuronx import MistralForSampling, GQA, NeuronConfig, QuantizationConfig
+import time
+import math
+import concurrent.futures
+
+
+def padding_ceiling(n):
+    if n <= 0:
+        return 1
+    elif n & (n - 1) == 0:  # Check if n is already a power of 2
+        return n
+    else:
+        return 2 ** math.ceil(math.log2(n))
+
+
+class MyStreamer(transformers.generation.streamers.BaseStreamer):
+    def __init__(self) -> None:
+        self.reset()
+
+    def reset(self):
+        self.token_latencies = []
+        self.iter = 0
+        self.now = time.time()
+
+    def put(self, tokens):
+        now = time.time()
+        token_latency = now - self.now
+        self.now = now
+        self.iter += 1
+        self.token_latencies.append(token_latency)
+
+    def end(self):
+        print("\n\n")
+        print("First 5 token latencies:", self.token_latencies[:5])
+        print("All token latencies:", sum(self.token_latencies[:]))
+
+
+class MistralModel:
+    """
+    A class for generating text using the Mistral language model.
+    """
+
+    def __init__(self, model_name):
+        self.neuron_config = NeuronConfig(group_query_attention=GQA.SHARD_OVER_HEADS,
+                                          quant=QuantizationConfig(quant_dtype='s8', dequant_dtype='bf16'))
+        # self.model_name = 'mistralai/Mistral-7B-Instruct-v0.2'
+        self.model_name = model_name
+        self.amp: Literal['bf16', 'fp32'] = 'bf16'
+        self.batch_size = 1
+        self.tp_degree = 2
+        self.n_positions = 4096
+        self.context_length_estimate = [2289, 4096]
+        # self.context_length_estimate = 2289
+
+        self.model = self._load_model()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.prompt_template = "<s>[INST] {prompt} [/INST]"
+
+    def _load_model(self) -> MistralForSampling:
+        """
+        Load and initialize the Mistral model.
+
+        Returns:
+            MistralForSampling: The initialized Mistral model.
+        """
+        model = MistralForSampling.from_pretrained(
+            self.model_name,
+            amp=self.amp,
+            batch_size=self.batch_size,
+            tp_degree=self.tp_degree,
+            n_positions=self.n_positions,
+            neuron_config=self.neuron_config,
+            context_length_estimate=self.context_length_estimate,
+            # compiler_args=["--model-type=transformer", "--target=inf2", "--auto-cast=all", "--auto-cast-type=fp8_e4m3", "--optlevel=3", "--enable-saturate-infinity"]
+        )
+        model.to_neuron()
+        return model
+
+    def generate(self, inputs: Union[str, List[int]], parameters: Optional[Dict[str, Any]] = None) -> str:
+        """
+        Generate text using the Mistral model.
+
+        Args:
+            inputs (Union[str, List[int]]): The input prompt or a list of input embeddings.
+            parameters (Optional[Dict[str, Any]]): Optional parameters for text generation.
+
+        Returns:
+            str: The generated text.
+
+        Raises:
+            ValueError: If the input type is invalid.
+        """
+        try:
+            max_new_tokens = parameters.get("max_new_tokens", 256)
+            top_k = parameters.get("top_k", 100)
+            top_p = parameters.get("top_p", 0.1)
+            temperature = parameters.get("temperature", 0.1)
+            no_repeat_ngram_size = parameters.get("no_repeat_ngram_size", 3)
+            print(
+                f"parameters max_new_tokens: {max_new_tokens}, top_k: {top_k}, top_p: {top_p}, temperature: {temperature}, no_repeat_ngram_size: {no_repeat_ngram_size}")
+
+            if isinstance(inputs, str):
+                generated_text = self._generate_from_prompt(inputs, max_new_tokens, top_k, top_p, temperature,
+                                                            no_repeat_ngram_size)
+            elif isinstance(inputs, list):
+                generated_text = self._generate_from_embeddings(inputs, max_new_tokens, top_k, top_p, temperature,
+                                                                no_repeat_ngram_size)
+            else:
+                raise ValueError("Invalid input type. Must be str or List[int]")
+
+            return generated_text
+        except Exception as e:
+            logging.error(f"Error generating text: {e}")
+            raise
+
+    def _generate_from_prompt(self, prompt: str, max_new_tokens: int, top_k: float, top_p: float, temperature: float,
+                              no_repeat_ngram_size: int) -> str:
+        """
+        Generate text from a given prompt using the Mistral model.
+
+        Args:
+            prompt (str): The input prompt.
+            max_new_tokens (int): The maximum number of new tokens to generate.
+
+        Returns:
+            str: The generated text.
+        """
+        input_prompt = self.prompt_template.format(prompt=prompt)
+        encoded_input = self.tokenizer(input_prompt, return_tensors='pt')
+        input_ids = encoded_input.input_ids
+
+        with torch.inference_mode():
+            generated_sequence = self.model.sample(input_ids, sequence_length=min(self.n_positions,
+                                                                                  input_ids.shape[1] + max_new_tokens),
+                                                   start_ids=None, top_k=top_k, top_p=top_p, temperature=temperature,
+                                                   no_repeat_ngram_size=no_repeat_ngram_size)
+            decoded_output = [self.tokenizer.decode(tok) for tok in generated_sequence]
+
+        generated_text = decoded_output[0].split('[/INST]')[1].strip("</s>").strip()
+        return generated_text
+
+    def _generate_from_embeddings(self, input_embeddings: List[int], max_new_tokens: int, top_k: float, top_p: float,
+                                  temperature: float, no_repeat_ngram_size: int) -> str:
+        """
+        Generate text from a given list of input embeddings using the Mistral model.
+
+        Args:
+            input_embeddings (List[int]): A list of input embeddings.
+            max_new_tokens (int): The maximum number of new tokens to generate.
+
+        Returns:
+            str: The generated text.
+        """
+        s1 = time.time()
+        input_embeds_tensor = torch.tensor(input_embeddings)
+        input_embeds_length = input_embeds_tensor.shape[1]
+        padding_size = padding_ceiling(input_embeds_length)
+        if padding_size >= self.n_positions:
+            padding_size = input_embeds_length
+            padded_input_embeds = input_embeds_tensor
+        else:
+            padding_gap = padding_size - input_embeds_length
+            padded_input_embeds = F.pad(input_embeds_tensor, (0, 0, padding_gap, 0), value=self.tokenizer.pad_token_id)
+        print("ms1 - input_embeds time: ", time.time() - s1)
+
+        s2 = time.time()
+        with torch.inference_mode():
+            generated_sequence = self.model.sample(padded_input_embeds,
+                                                   sequence_length=min(self.n_positions, padding_size + max_new_tokens),
+                                                   start_ids=None, top_k=top_k, top_p=top_p, temperature=temperature,
+                                                   no_repeat_ngram_size=no_repeat_ngram_size, streamer=MyStreamer())
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                decoded_output = list(executor.map(self.tokenizer.decode, generated_sequence))
+            # decoded_output = [self.tokenizer.decode(tok) for tok in generated_sequence]
+        print("ms2 - decoded_output time: ", time.time() - s2)
+
+        generated_text = decoded_output[0].strip("</s>").strip()
+        return generated_text
+

app/llava/__init__.py

@@ -0,0 +1,9 @@
+from .model import LlavaLlamaForCausalLM, LlavaMistralForCausalLM
+try:
+    from .model import LlavaGemmaForCausalLM
+except:
+    pass
+try:
+    from .model import LlavaThothForCausalLM
+except:
+    pass

app/llava/configs/action_dataset_ablation/finetune_webvid.yaml

@@ -0,0 +1,11 @@
+datasets:
+
+  lk_image:
+    data_type: image
+
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    select_datasets: ['webvid10m', 'webvid2m']
+    # select_datasets: ['webvid10m', 'webvid2m', 'activitynet', 'vidal', 'hdvila']

app/llava/configs/action_dataset_ablation/finetune_webvid_act.yaml

@@ -0,0 +1,11 @@
+datasets:
+
+  lk_image:
+    data_type: image
+
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    select_datasets: ['webvid10m', 'webvid2m', 'activitynet']
+    # select_datasets: ['webvid10m', 'webvid2m', 'activitynet', 'vidal', 'hdvila']

app/llava/configs/action_dataset_ablation/finetune_webvid_hdvila.yaml

@@ -0,0 +1,11 @@
+datasets:
+
+  lk_image:
+    data_type: image
+
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    select_datasets: ['webvid10m', 'webvid2m', 'hdvila']
+    # select_datasets: ['webvid10m', 'webvid2m', 'activitynet', 'vidal', 'hdvila']

app/llava/configs/action_dataset_ablation/finetune_webvid_vidal.yaml

@@ -0,0 +1,11 @@
+datasets:
+
+  lk_image:
+    data_type: image
+
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    select_datasets: ['webvid10m', 'webvid2m', 'vidal']
+    # select_datasets: ['webvid10m', 'webvid2m', 'activitynet', 'vidal', 'hdvila']

app/llava/configs/adso_increasing_ablation/finetune_data_pure_gpt4v.yaml

@@ -0,0 +1,55 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+#  tt_vqa:
+#    data_type: frames
+#    sample_ratio: 1
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']
+    sample_method: sequential
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_130k.json
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']
+
+

app/llava/configs/adso_increasing_ablation/finetune_gpt4v_adso135k.yaml

@@ -0,0 +1,57 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  tt_vqa:
+    data_type: frames
+    sample_ratio: 2
+    fps: 2.0
+    conv_type: single  
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/20240208_meta_data_single_135k_caption_160k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']
+    sample_method: sequential
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_130k.json
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']

app/llava/configs/adso_increasing_ablation/finetune_gpt4v_adso185k.yaml

@@ -0,0 +1,57 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  tt_vqa:
+    data_type: frames
+    sample_ratio: 3
+    fps: 2.0
+    conv_type: single  
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/20240220_meta_data_single_190k_caption_160k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']
+    sample_method: sequential
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_130k.json
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']

app/llava/configs/adso_increasing_ablation/finetune_gpt4v_adso185k_baseline.yaml

@@ -0,0 +1,55 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  tt_vqa:
+    data_type: frames
+    sample_ratio: 3
+    fps: 2.0
+    conv_type: single  
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/20240220_meta_data_single_190k_caption_160k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    sample_ratio: 6
+  
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['detail']
+

app/llava/configs/adso_increasing_ablation/finetune_gpt4v_adso185k_no_qa.yaml

@@ -0,0 +1,57 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  tt_vqa:
+    data_type: frames
+    sample_ratio: 3
+    fps: 2.0
+    conv_type: single  
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/20240220_meta_data_single_190k_caption_no_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']
+    sample_method: sequential
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_130k.json
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']

app/llava/configs/adso_increasing_ablation/finetune_gpt4v_adso65k.yaml

@@ -0,0 +1,57 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  tt_vqa:
+    data_type: frames
+    sample_ratio: 2
+    fps: 2.0
+    conv_type: single  
+    train_data_path: /mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_single_60k_caption_170k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']
+    sample_method: sequential
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_130k.json
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']

app/llava/configs/finetune_debug.yaml

@@ -0,0 +1,8 @@
+datasets:
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_130k.json

app/llava/configs/finetune_gpt4v_adso65k.yaml

@@ -0,0 +1,56 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  tt_vqa:
+    data_type: frames
+    sample_ratio: 2
+    fps: 2.0
+    conv_type: single  
+    train_data_path: /mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_single_60k_caption_170k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_130k.json
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']

app/llava/configs/gpt4v_increasing_ablation/finetune_gpt4v_public500k.yaml

@@ -0,0 +1,57 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  # tt_vqa:
+  #   data_type: frames
+  #   sample_ratio: 2
+  #   fps: 2.0
+  #   conv_type: single  
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_single_60k_caption_170k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 10
+    conv_type: single
+    task_types: ['summary', 'detail']
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']

app/llava/configs/gpt4v_increasing_ablation/finetune_gpt4v_public500k_no_summary.yaml

@@ -0,0 +1,57 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  # tt_vqa:
+  #   data_type: frames
+  #   sample_ratio: 2
+  #   fps: 2.0
+  #   conv_type: single  
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_single_60k_caption_170k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+
+  gpt4v_public:
+    data_type: frames
+    fps: 1.0
+    sample_ratio: 4
+    conv_type: single
+    task_types: ['detail']
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']

app/llava/configs/gpt4v_increasing_ablation/finetune_gpt4v_public800k.yaml

@@ -0,0 +1,62 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  # tt_vqa:
+  #   data_type: frames
+  #   sample_ratio: 2
+  #   fps: 2.0
+  #   conv_type: single  
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_single_60k_caption_170k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption', 'qas']
+
+  # gpt4v_public:
+  #   data_type: frames
+  #   fps: 1.0
+  #   sample_ratio: 10
+  #   conv_type: single
+  #   task_types: ['summary', 'detail']
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json
+  
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    sample_ratio: 6
+  
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['summary', 'detail', 'qa_pairs']

app/llava/configs/gpt4v_increasing_ablation/finetune_videollava.yaml

@@ -0,0 +1,20 @@
+datasets:
+
+  # llava_pretrain:
+  #   data_type: image
+  #   sample_ratio: 1
+
+  # gpt4v_public:
+  #   data_type: frames
+  #   sample_ratio: 2
+  #   task_types: ['summary']
+  #   fps: 1.0
+  #   conv_type: single
+
+  lk_image:
+    data_type: image
+
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0

app/llava/configs/pretrain_data.yaml

@@ -0,0 +1,17 @@
+datasets:
+
+  llava_pretrain:
+    data_type: image
+    sample_ratio: 1
+
+  # internvid:
+  #   data_type: frames
+  #   sample_ratio: 10
+  
+  gpt4v_public:
+    data_type: frames
+    sample_ratio: 1
+    task_types: ['summary']
+    fps: 1.0
+    conv_type: single
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json

app/llava/configs/pretrain_data_large.yaml

@@ -0,0 +1,17 @@
+datasets:
+
+  llava_pretrain:
+    data_type: image
+    sample_ratio: 1
+
+  internvid:
+    data_type: frames
+    sample_ratio: 10
+  
+  gpt4v_public:
+    data_type: frames
+    sample_ratio: 1
+    task_types: ['summary']
+    fps: 1.0
+    conv_type: single
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json

app/llava/configs/pretrain_debug.yaml

@@ -0,0 +1,27 @@
+datasets:
+
+  llava_pretrain:
+    data_type: image
+    sample_ratio: 1
+
+  # gpt4v_public:
+  #   data_type: frames
+  #   sample_ratio: 2
+  #   task_types: ['summary']
+  #   fps: 1.0
+  #   conv_type: single
+
+  # lk_image:
+  #   data_type: image
+
+  # lk_video:
+  #   data_type: frames
+  #   conv_type: multi
+  #   fps: 1.0
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: multi
+    task_types: ['qa_pairs']

app/llava/configs/promptv1_2_increasing_ablation/finetune_gpt4_prompt_140k.yaml

@@ -0,0 +1,35 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  # gpt4v_internal:
+  #   data_type: frames
+  #   fps: 2.0
+  #   sample_ratio: 1
+  #   conv_type: single
+  #   task_types: ['summary', 'detail', 'qa_pairs']
+    
+  promptv1_2_internal:
+    data_type: frames
+    sample_ratio: 1
+    train_data_path: /mnt/bn/algo-masp-nas-2/kaili.zhao/data/masp_data/train/gpt4v_annotation/202400401week_gpt4v_all_videos_unique_ids.json
+    task_types: ['refine_caption']

app/llava/configs/release_version/finetune_250k_no_public.yaml

@@ -0,0 +1,50 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  tt_vqa:
+    data_type: frames
+    sample_ratio: 3
+    fps: 2.0
+    conv_type: single  
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/20231201_20240322_caption_250k.json
+
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption']
+
+
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['detail']

app/llava/configs/release_version/finetune_all_data.yaml

@@ -0,0 +1,63 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  tt_vqa:
+    data_type: frames
+    sample_ratio: 3
+    fps: 2.0
+    conv_type: single  
+    train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/20231201_20240322_caption_250k.json
+
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption']
+
+  # gpt4v_public:
+  #   data_type: frames
+  #   fps: 1.0
+  #   sample_ratio: 10
+  #   conv_type: single
+  #   task_types: ['summary', 'detail']
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json
+  
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    sample_ratio: 6
+  
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['detail']

app/llava/configs/release_version/finetune_gpt4v_caption.yaml

@@ -0,0 +1,62 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  # tt_vqa:
+  #   data_type: frames
+  #   sample_ratio: 2
+  #   fps: 2.0
+  #   conv_type: single  
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_single_60k_caption_170k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption']
+
+  # gpt4v_public:
+  #   data_type: frames
+  #   fps: 1.0
+  #   sample_ratio: 10
+  #   conv_type: single
+  #   task_types: ['summary', 'detail']
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json
+  
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    sample_ratio: 6
+  
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['detail']

app/llava/configs/release_version/finetune_gpt4v_caption_ocr.yaml

@@ -0,0 +1,67 @@
+datasets:
+
+#  m3it:
+#    data_type: images
+#    sample_ratio: 4
+#    tasks:
+#      - coco
+#      - coco-goi
+#      - coco-text
+#      - imagenet
+#      - coco-itm
+#      - iqa
+#      - mocheg
+#      - vsr
+#      - refcoco
+#      - science-qa
+#      - vqa-v2
+#      - gqa
+#      - st-vqa
+#      - text-vqa
+#      - okvqa
+#      - a-okvqa
+#
+  # tt_vqa:
+  #   data_type: frames
+  #   sample_ratio: 2
+  #   fps: 2.0
+  #   conv_type: single  
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_single_60k_caption_170k_QA.json
+
+  ShareGPT4V:
+    data_type: images
+    sample_ratio: 1
+
+
+  gpt4v_tt_vqa:
+    data_type: frames
+    fps: 0.5
+    sample_ratio: 6
+    conv_type: single
+    task_types: ['caption']
+
+  # gpt4v_public:
+  #   data_type: frames
+  #   fps: 1.0
+  #   sample_ratio: 10
+  #   conv_type: single
+  #   task_types: ['summary', 'detail']
+  #   train_data_path: /mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json
+  
+  lk_video:
+    data_type: frames
+    conv_type: multi
+    fps: 1.0
+    sample_ratio: 6
+  
+  gpt4v_internal:
+    data_type: frames
+    fps: 2.0
+    sample_ratio: 1
+    conv_type: single
+    task_types: ['detail']
+
+  synthetic_ocr:
+    data_type: video
+    sample_ratio: 1
+    fps: 0.5

app/llava/constants.py

@@ -0,0 +1,17 @@
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+# Model Constants
+IGNORE_INDEX = -100
+MM_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+IMAGE_PLACEHOLDER = "<image-placeholder>"
+DEFAULT_VIDEO_TOKEN = "<video>"
+DEFAULT_VIDEO_PATCH_TOKEN = "<vid_patch>"
+DEFAULT_VIDEO_START_TOKEN = "<vid_start>"
+DEFAULT_VIDEO_END_TOKEN = "<vid_end>"

app/llava/conversation.py

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

app/llava/datasets/__init__.py

@@ -0,0 +1,24 @@
+from .data_cfgs import *
+from .base_dataset import *
+from .prompts import *
+from .super_dataset import *
+from .cc_sbu_dataset import *
+from .llava_pretrain_dataset import *
+# from .llava_instruct_dataset import *
+# from .lrv_instruct_dataset import *
+from .internvid_dataset import *
+from .tt_vqa_dataset import *
+from .m3it_dataset import *
+from .sharegpt4v_dataset import *
+from .gpt4v_tt_vqa_dataset import *
+from .gpt4v_public_dataset import *
+from .gpt4v_internal_dataset import *
+# from .synthdog_dataset import *
+# from .ocr_vqa_dataset import *
+# from .sharegpt_dataset import *
+from .textcaps_dataset import * 
+from .synthetic_ocr_dataset import *
+from .lk_image_dataset import *
+from .lk_video_dataset import *
+
+from .promptv1_2_internal_dataset import *

app/llava/datasets/base_dataset.py

@@ -0,0 +1,234 @@
+import os
+import json
+import sys
+import copy
+import math
+import torch
+import decord
+import random
+import numpy as np
+from PIL import Image
+from decord import VideoReader
+from torch.utils.data import Dataset
+from llava.utils import master_print
+from typing import Dict, Optional, Sequence, List
+from llava.datasets.data_cfgs import data_configs
+from transformers import CLIPImageProcessor, SiglipImageProcessor
+
+from llava.mm_utils import get_frame_indices, process_anyres_image
+from torch.utils.data.dataloader import default_collate
+
+decord.bridge.set_bridge("torch")
+
+class TaskBaseDataset(Dataset):
+    """ Implementation of base task dataset """
+    def __init__(self, anno_path=None, data_args=None, name=None, **kwargs):
+
+        self.anno_path = anno_path
+        self.data_args = data_args
+        self.image_aspect_ratio = data_args.image_aspect_ratio
+        self.image_grid_pinpoints = data_args.image_grid_pinpoints
+        self.vis_processor = data_args.image_processor
+        self.type = None
+        self.name = name
+
+        master_print(f"Loading dataset {name}...")
+        if (anno_path is not None):
+            if not hasattr(self, 'annotation'):
+                self.annotation = json.load(open(anno_path, 'r'))
+            master_print(f"Finish loading dataset {name} {len(self.annotation)} samples...")
+
+    def __len__(self):
+        return len(self.annotation)
+
+    def collater(self, samples):
+        return default_collate(samples)
+
+    def text_preprocess(self, sources) -> List[List[Dict[str, str]]]:
+        pass
+
+    def vis_preprocess(self, vis_path) -> Image:
+        pass
+
+    @property
+    def data_type(self):
+        return self.type
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        item = self.annotation[i]
+
+        vis_path = item['vis_path'] if 'vis_path' in item else item['video_path']
+
+        ret = {
+            'images': self.vis_preprocess(vis_path),
+            'conversations': self.text_preprocess(item)
+        }
+        if 'id' in item:
+            ret['id'] = item['id']
+
+        return ret
+
+
+class ImageTaskDataset(TaskBaseDataset):
+    def __init__(self, anno_path=None, data_args=None, name=None):
+        super().__init__(anno_path=anno_path,
+                         data_args=data_args,
+                         name=name)
+        self.type = 'images'
+
+    @staticmethod
+    def expand2square(pil_img, background_color):
+        width, height = pil_img.size
+        if width == height:
+            return pil_img
+        elif width > height:
+            result = Image.new(pil_img.mode, (width, width), background_color)
+            result.paste(pil_img, (0, (width - height) // 2))
+            return result
+        else:
+            result = Image.new(pil_img.mode, (height, height), background_color)
+            result.paste(pil_img, ((height - width) // 2, 0))
+            return result
+
+    def preprocess_image(self, image):
+        if self.image_aspect_ratio == 'pad':
+            image = self.expand2square(image, tuple(int(x *255) for x in self.vis_processor.image_mean))
+            if isinstance(self.vis_processor, CLIPImageProcessor) or isinstance(self.vis_processor, SiglipImageProcessor):
+                image = self.vis_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+            else:
+                image = self.vis_processor.preprocess(image)
+        elif self.image_aspect_ratio == "anyres":
+            image = process_anyres_image(image, self.vis_processor, self.image_grid_pinpoints)
+        else:
+            if isinstance(self.vis_processor, CLIPImageProcessor) or isinstance(self.vis_processor, SiglipImageProcessor):
+                image = self.vis_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+            else:
+                image = self.vis_processor.preprocess(image)
+
+        return image
+
+    def vis_preprocess(self, vis_path):
+        image = Image.open(vis_path).convert('RGB')
+        image = self.preprocess_image(image)
+        if isinstance(image, list):
+            images = image
+        else:
+            images = [image]
+
+        return images
+
+
+class VideoTaskDataset(ImageTaskDataset):
+    def __init__(self, anno_path=None, data_args=None, name=None):
+        super().__init__(anno_path=anno_path,
+                         data_args=data_args,
+                         name=name)
+
+        # if not specify num_segments, use default
+        self.num_segments = self.data_args.num_segments
+        self.sample_strategy = self.data_args.sample_strategy
+        self.type = 'video'
+
+    def vis_preprocess(self, vis_path):
+        images = None
+        try:
+            video_reader = VideoReader(vis_path)
+            vlen = len(video_reader)
+            fps = video_reader.get_avg_fps()
+            duration = vlen / float(fps)
+
+            frame_indices = get_frame_indices(self.num_segments, vlen,
+                                              sample=self.sample_strategy, input_fps=fps, pad_last=False)
+            frames = video_reader.get_batch(frame_indices)
+            frames = frames.numpy().astype(np.uint8)
+            images = [Image.fromarray(frame).convert('RGB') for frame in frames]
+            images = [self.preprocess_image(image) for image in images]
+        except Exception as e:
+            print(e, vis_path)
+            sys.stdout.flush()
+            images = None
+
+        # print(f"images: {len(images)}, {images[0].shape}")
+
+        return images
+
+
+class FramesTaskDataset(ImageTaskDataset):
+    def __init__(self, anno_path=None, data_args=None, fps=0.5, name=None):
+        super().__init__(anno_path=anno_path,
+                         data_args=data_args,
+                         name=name)
+
+        # if not specify num_segments, use default
+        self.num_segments = self.data_args.num_segments
+        # print("self.num_segments:", self.num_segments)
+        self.type = 'video'
+        self.default_fps = 2.0
+        self.fps = fps
+
+    @staticmethod
+    def _downsample_frames(frames, interval, keep_first_last=True):
+        if keep_first_last:
+            first, last, mid = frames[0], frames[-1], frames[1:-1]
+            sampled_frames = mid[interval - 1::interval]
+            ret = [first] + sampled_frames + [last]
+
+        else:
+            # may output empty list, recommend keep first and last frame
+            ret = frames[interval - 1::interval]
+
+        return ret
+
+    @staticmethod
+    def _sample_frames(frames, num_segments):
+        frame_indices = list(range(len(frames)))
+        cand_indices  = copy.deepcopy(frame_indices)
+        intervals = np.linspace(start=0, stop=len(frame_indices), num=num_segments + 1).astype(int)
+        ranges = []
+
+        for idx, interv in enumerate(intervals[:-1]):
+            ranges.append((interv, intervals[idx + 1] - 1))
+
+        try:
+            frame_indices = [cand_indices[random.choice(range(x[0], x[1]))] for x in ranges]
+        except:
+            frame_indices = [cand_indices[x[0]] for x in ranges]
+
+        sampled_frames = [frames[indice] for indice in frame_indices]
+
+        return sampled_frames
+
+    def vis_preprocess(self, vis_path):
+        image_files = [(os.path.splitext(img)[0], img) for img in os.listdir(vis_path) if not img.startswith('cuttime')]
+        if image_files[0][1].endswith('jpeg'):
+            # gpt4v public data
+            image_files = [(int(x[0].split('_')[-1]), x[1]) for x in image_files]
+        else:
+            image_files = [(int(x[0]), x[1]) for x in image_files]
+
+        image_files = sorted(image_files, key=lambda img: img[0])
+
+        if self.fps < self.default_fps:
+            interval = math.floor(self.default_fps / self.fps)
+            image_files = self._downsample_frames(image_files, interval, keep_first_last=True)
+
+        if self.num_segments > 0 and len(image_files) > self.num_segments:
+            image_files = self._sample_frames(image_files, self.num_segments)
+
+        images = []
+        for image_file in image_files:
+            try:
+                images.append(Image.open(os.path.join(vis_path, image_file[1])).convert('RGB'))
+            except Exception as e:
+                continue
+        formatted_images = []
+        for image in images:
+            im = self.preprocess_image(image)
+            if isinstance(im, list):
+                formatted_images.extend(im)
+            else:
+                formatted_images.append(im)
+        return formatted_images
+
+
+

app/llava/datasets/builder.py

@@ -0,0 +1,5 @@
+from .registry import Registry
+
+__all__ = ['DATASETS']
+
+DATASETS = Registry('datasets')

app/llava/datasets/cc_sbu_dataset.py

@@ -0,0 +1,40 @@
+import os
+import random
+from llava.datasets.builder import DATASETS
+
+from typing import Dict, Optional, Sequence, List
+from llava.datasets.data_cfgs import data_configs
+from llava.datasets.base_dataset import ImageTaskDataset
+from llava.datasets.prompts import cc_sbu_prompt
+from llava.constants import DEFAULT_IMAGE_TOKEN
+
+
+class CCSBUDataset(ImageTaskDataset):
+    def __init__(self, anno_path, data_args=None, name='cc_sbu'):
+        super().__init__(anno_path=anno_path,
+                         data_args=data_args,
+                         name=name)
+
+    def text_preprocess(self, item) -> List[Dict[str, str]]:
+        caption = item['caption']
+
+        conversations = [
+            {
+                'from': 'human',
+                'value': DEFAULT_IMAGE_TOKEN + random.choice(cc_sbu_prompt)
+            },
+            {
+                'from': 'model',
+                'value': caption
+            }
+        ]
+
+        return conversations
+
+
+@DATASETS.register_obj
+def cc_sbu(data_args):
+    return CCSBUDataset(data_configs["cc_sbu"]['train_data_path'], data_args)
+
+
+

app/llava/datasets/data_cfgs.py

@@ -0,0 +1,157 @@
+data_configs = {
+    'llava_pretrain': {
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/blip_laion_cc_sbu_558k/meta_data.json'
+    },
+    'llava_instruct': {
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/llava_instruct_150k/meta_data.json'
+    },
+    'lrv_instruct': {
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/lrv_instructions/meta_data.json'
+    },
+    'coco_caption': {
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/data-tns-algo-masp/baiyi.by/data/coco_caption/train.json'
+    },
+    'cc_sbu': {
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/baiyi-arnold-nas/data/masp/vlm_data/cc_sbu/meta_data.json'
+    },
+    'laion': {
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/data-tns-algo-masp/baiyi.by/data/laion/train.json'
+    },
+    'webvid': {
+        'data_type': 'video',
+        'train_data_path': '/mnt/bn/baiyi-arnold-nas/data/masp/vlm_data/webvid_10M_video/train.json',
+        'val_data_path': '/mnt/bn/baiyi-arnold-nas/data/masp/vlm_data/webvid_10M_video/val.json'
+    },
+    'internvid': {
+        'data_type': 'frames',
+        'fps': 0.5,
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/InternVid/meta_data.json'
+    },
+    'video_chatgpt_instruct_single': {
+        'data_type': 'video',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/VideoChatGPT_Instruct_100K_single/train.json'
+    },
+    'video_chatgpt_instruct_multi': {
+        'data_type': 'video',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/VideoChatGPT_Instruct_100K_multi/train.json'
+    },
+    'video_chatgpt': {
+        'data_type': 'frames',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/video_chatgpt_instruct/meta_data.json'
+    },
+    'm3it': {
+        'data_type': 'images',
+        'default_tasks': [
+            'coco',
+            'textcap',
+            'image-paragraph-captioning',
+            'coco-goi',
+            'coco-itm',
+            'vqa-v2',
+            'shapes',
+            'docvqa',
+            'ocr-vqa',
+            'st-vqa',
+            'text-vqa',
+            'gqa',
+            'okvqa',
+            'a-okvqa',
+            'viquae',
+            'clevr',
+            'nlvr',
+            'vcr',
+            'visual-mrc',
+            'visual-dialog',
+            'multi30k'
+        ]
+    },
+    'tt_vqa': {
+        'data_type': 'frames',
+        'fps': 2,
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_single_60k_caption_170k_QA.json'
+        # 'train_data_path': '/mnt/bn/yukunfeng-nasdrive/xiangchen/dataset/masp/20240208_meta_data_single_135k_caption_160k_QA.json'
+        # 'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/ADSO_Anno_Data/batch_20231128/meta_data_final_single_non_empty.json'
+    },
+    'gpt4v_tt_vqa': {
+        'data_type': 'frames',
+        'fps': 0.5,
+        # 'train_data_path': '/mnt/bn/algo-masp-nas-2/baiyi.by/data/GPT4V_Negs/20231127_81k_single.json'
+        # 'train_data_path': '/mnt/bn/yukunfeng-nasdrive/xiangchen/dataset/masp/20231127_81k_25k_filtered_single_non_empty.json'
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/20231222_120k_multi_filtered.json',
+        'task_types': ['caption', 'qas'],
+        'conv_type': 'single'
+    },
+    'sharegpt4v': {
+        'data_type': 'images',
+        'coco_dir': '/mnt/bn/data-tns-algo-masp/data',
+        'llava_dir': '/mnt/bn/data-tns-algo-masp/baiyi.by/data/blip_laion_cc_sbu_558k',
+        'other_dir': '/mnt/bn/algo-masp-nas-2/xiangchen/dataset/sharegpt4v',
+    },
+    'gpt4v_public': {
+        'data_type': 'frames',
+        'fps': 1,
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_130k.json',
+        # 'train_data_path': '/mnt/bn/algo-masp-nas-2/xiangchen/data/shared_gpt4v_data/data_500k_filtered.json',
+        'task_types': ['summary', 'detail', 'qa_pairs'],
+        'conv_type': 'single',
+        'sample_method': 'uniform'
+    },
+
+    'gpt4v_internal': {
+        'data_type': 'frames',
+        'fps': 2,
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/gpt4v_internal_28k.json',
+        'task_types': ['summary','detail','qa_pairs'],
+        'conv_type': 'single'
+    },
+
+    'synthdog': { #500k
+        'data_type': 'images',
+    },
+
+    'ocr_vqa': { #200k
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/xiangchen/dataset/OCR-VQA/training_meta.json'
+    },
+
+    'sharegpt': { #50k
+        'data_type': 'text'
+    },
+
+    'text_caps':{ #100k
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/xiangchen/dataset/TextCaps/TextCaps_0.1_train.json'
+    },
+
+    'synthetic_ocr':{ # 50k
+        'data_type': 'frames',
+        'fps': 0.5,  # total 10 frames for each video
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/xiangchen/dataset/masp/synthetic_ocr/train_filtered.json'
+    },
+
+    'lk_image':{  # 600k
+        'data_type': 'images',
+        'train_data_path': '/mnt/bn/liangkeg/data/xiangchen/finetune_all_detail_vidal200k_videollava_images_im.json'
+    },
+
+    'lk_video':{  # 850k
+        'data_type': 'frames',
+        'fps': 1,
+        'train_data_path': '/mnt/bn/liangkeg/data/xiangchen/finetune_all_detail_vidal200k_videollava_images_vid.json',
+        'select_datasets': ['webvid10m', 'webvid2m', 'activitynet', 'vidal', 'hdvila'],
+    },
+
+    'promptv1_2_internal':{  # 210k
+        'data_type': 'frames',
+        'train_data_path': '/mnt/bn/algo-masp-nas-2/kaili.zhao/data/masp_data/train/gpt4v_annotation/202400401week_gpt4v_all_videos_unique_ids.json',
+        'task_types': ['caption']
+    }  
+}
+
+