Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,148 @@
+*.pt
+
+wandb/
+
+checkpoints/
+tests/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Pycharm project settings
+.idea
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+*.out
+src/wandb
+wandb
+
+# Pyre type checker
+.pyre/
+
+# Training
+batchscript*
+work_dirs
+data

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

LICENSE

@@ -0,0 +1,204 @@
+Copyright 2018-2023 OpenMMLab. All rights reserved.
+
+                                 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 2018-2023 OpenMMLab.
+
+   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,12 +1,196 @@
----
-title: Multimodal GPT
-emoji: 🐠
-colorFrom: indigo
-colorTo: gray
-sdk: gradio
-sdk_version: 5.21.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Multimodal-GPT
+app_file: app.py
+sdk: gradio
+sdk_version: 5.21.0
+---
+# 🤖 Multi-modal GPT
+
+Train a multi-modal chatbot with visual and language instructions!
+
+Based on the open-source multi-modal model [OpenFlamingo](https://github.com/mlfoundations/open_flamingo), we create various **visual instruction** data with open datasets, including VQA, Image Captioning, Visual Reasoning, Text OCR, and Visual Dialogue. Additionally, we also train the language model component of OpenFlamingo using only **language-only instruction** data.
+
+The **joint training** of visual and language instructions effectively improves the performance of the model! For more details please refer to our [technical report](https://arxiv.org/abs/2305.04790).
+
+Welcome to join us!
+
+</div>
+
+<div align="center">
+
+English | [简体中文](README_zh-CN.md)
+
+</div>
+
+<div align="center">
+  <a href="https://openmmlab.medium.com/" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/219255827-67c1a27f-f8c5-46a9-811d-5e57448c61d1.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://discord.com/channels/1037617289144569886/1046608014234370059" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/218347213-c080267f-cbb6-443e-8532-8e1ed9a58ea9.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://twitter.com/OpenMMLab" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/218346637-d30c8a0f-3eba-4699-8131-512fb06d46db.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://www.youtube.com/openmmlab" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/218346691-ceb2116a-465a-40af-8424-9f30d2348ca9.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://space.bilibili.com/1293512903" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/219026751-d7d14cce-a7c9-4e82-9942-8375fca65b99.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://www.zhihu.com/people/openmmlab" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/219026120-ba71e48b-6e94-4bd4-b4e9-b7d175b5e362.png" width="3%" alt="" /></a>
+</div>
+
+## Features
+
+- Support various vision and language instruction data
+- Parameter efficient fine-tuning with LoRA
+- Tuning vision and language at the same time, complement each other
+
+
+## Installation
+
+To install the package in an existing environment, run
+
+```bash
+git clone https://github.com/open-mmlab/Multimodal-GPT.git
+cd Multimodal-GPT
+pip install -r requirements.txt
+pip install -v -e .
+```
+
+or create a new conda environment
+
+```bash
+conda env create -f environment.yml
+```
+
+
+## Launch Demo Locally
+
+1. Download the pre-trained weights.
+
+    Use [this script](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py) for converting LLaMA weights to Hugging Face format.
+
+    Download the OpenFlamingo pre-trained model from [openflamingo/OpenFlamingo-9B](https://huggingface.co/openflamingo/OpenFlamingo-9B).
+
+    Download our LoRA Weight from [here](https://download.openmmlab.com/mmgpt/v0/mmgpt-lora-v0-release.pt).
+
+    Then place these models in `checkpoints` folders like this:
+
+    ```
+    checkpoints
+    ├── llama-7b_hf
+    │   ├── config.json
+    │   ├── pytorch_model-00001-of-00002.bin
+    │   ├── ......
+    │   └── tokenizer.model
+    ├── OpenFlamingo-9B
+    │   └──checkpoint.pt
+    ├──mmgpt-lora-v0-release.pt
+
+2. launch the gradio demo
+
+    ```bash
+    python app.py
+    ```
+
+## Examples
+
+### Recipe:
+![image4](https://user-images.githubusercontent.com/12907710/234554562-8f3be88f-d563-47ba-97d9-ade8d47c46b0.png)
+
+### Travel plan:
+![image3](https://user-images.githubusercontent.com/12907710/234523464-80c4e3f0-f99f-4498-96ef-dc43ef89c64b.png)
+
+### Movie:
+![image2](https://user-images.githubusercontent.com/12907710/234523468-e11905a6-491f-4b87-934f-90da7d14d1c3.png)
+
+### Famous person:
+![image](https://user-images.githubusercontent.com/12907710/234523475-fd91f979-a344-4228-813f-6b55a1bc250f.png)
+
+
+## Fine-tuning
+
+### Prepare datasets
+
+1. [A-OKVQA](https://allenai.org/project/a-okvqa/home)
+
+    Download annotation from [this link](https://prior-datasets.s3.us-east-2.amazonaws.com/aokvqa/aokvqa_v1p0.tar.gz) and unzip to `data/aokvqa/annotations`.
+
+    It also requires images from coco dataset which can be downloaded from [here](https://cocodataset.org/#home). 
+
+2. [COCO Caption](https://cs.stanford.edu/people/karpathy/deepimagesent/)
+
+    Download from [this link](https://cs.stanford.edu/people/karpathy/deepimagesent/coco.zip) and unzip to `data/coco`.
+
+    It also requires images from coco dataset which can be downloaded from [here](https://cocodataset.org/#home).
+
+3. [OCR VQA](https://ocr-vqa.github.io/)
+
+    Download from [this link](https://drive.google.com/drive/folders/1_GYPY5UkUy7HIcR0zq3ZCFgeZN7BAfm_?usp=sharing) and place in `data/OCR_VQA/`.
+
+4. [LlaVA](https://llava-vl.github.io/)
+
+    Download from [liuhaotian/LLaVA-Instruct-150K](https://huggingface.co/datasets/liuhaotian/LLaVA-Instruct-150K) and place in `data/llava/`.
+
+    It also requires images from coco dataset which can be downloaded from [here](https://cocodataset.org/#home).
+
+5. [Mini-GPT4](https://minigpt-4.github.io/)
+
+    Download from [Vision-CAIR/cc_sbu_align](https://huggingface.co/datasets/Vision-CAIR/cc_sbu_align) and place in `data/cc_sbu_align/`.
+
+6. [Dolly 15k](https://www.databricks.com/blog/2023/03/24/hello-dolly-democratizing-magic-chatgpt-open-models.html)
+
+    Download from [databricks/databricks-dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k) and place it in `data/dolly/databricks-dolly-15k.jsonl`.
+
+7. [Alpaca GPT4](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+
+    Download it from [this link](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM/raw/main/data/alpaca_gpt4_data.json) and place it in `data/alpaca_gpt4/alpaca_gpt4_data.json`.
+
+You can also customize the data path in the [configs/dataset_config.py](configs/dataset_config.py).
+
+8. [Baize](https://github.com/project-baize/baize-chatbot)
+
+    Download it from [this link](https://github.com/project-baize/baize-chatbot/blob/main/data/quora_chat_data.json) and place it in `data/baize/quora_chat_data.json`.
+
+
+## Start training
+
+```bash
+torchrun --nproc_per_node=8 mmgpt/train/instruction_finetune.py \
+  --lm_path checkpoints/llama-7b_hf \
+  --tokenizer_path checkpoints/llama-7b_hf \
+  --pretrained_path checkpoints/OpenFlamingo-9B/checkpoint.pt \
+  --run_name train-my-gpt4 \
+  --learning_rate 1e-5 \
+  --lr_scheduler cosine \
+  --batch_size 1 \ 
+  --tuning_config configs/lora_config.py \
+  --dataset_config configs/dataset_config.py \
+  --report_to_wandb
+```
+
+
+## Acknowledgements
+
+- [OpenFlamingo](https://github.com/mlfoundations/open_flamingo)
+- [LAVIS](https://github.com/salesforce/LAVIS)
+- [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca)
+- [MiniGPT-4](https://github.com/Vision-CAIR/MiniGPT-4)
+- [LLaVA](https://github.com/haotian-liu/LLaVA/tree/main)
+- [Instruction Tuning with GPT-4](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+
+If you find our project useful for your research and applications, please cite using this BibTeX:
+
+```bibtex
+@misc{gong2023multimodalgpt,
+      title={MultiModal-GPT: A Vision and Language Model for Dialogue with Humans}, 
+      author={Tao Gong and Chengqi Lyu and Shilong Zhang and Yudong Wang and Miao Zheng and Qian Zhao and Kuikun Liu and Wenwei Zhang and Ping Luo and Kai Chen},
+      year={2023},
+      eprint={2305.04790},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

README_zh-CN.md

@@ -0,0 +1,184 @@
+# 🤖 Multi-modal GPT
+
+使用视觉和语言指令训练一个多模态聊天机器人！
+
+基于开源多模态模型 [OpenFlamingo](https://github.com/mlfoundations/open_flamingo)，我们使用公开数据集创建了各种**视觉指令**数据，包括视觉问答、图像字幕、视觉推理、文本 OCR 和视觉对话。此外，我们还使用仅包含**语言指令**数据的语言模型组件进行了训练。
+
+视觉和语言指令的**联合训练**有效提高了模型的性能！更多细节请参阅我们的[技术报告](https://arxiv.org/abs/2305.04790)。
+
+欢迎加入我们！
+
+</div>
+
+<div align="center">
+
+[English](README.md) | 简体中文
+
+</div>
+
+<div align="center">
+  <a href="https://openmmlab.medium.com/" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/219255827-67c1a27f-f8c5-46a9-811d-5e57448c61d1.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://discord.com/channels/1037617289144569886/1046608014234370059" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/218347213-c080267f-cbb6-443e-8532-8e1ed9a58ea9.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://twitter.com/OpenMMLab" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/218346637-d30c8a0f-3eba-4699-8131-512fb06d46db.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://www.youtube.com/openmmlab" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/218346691-ceb2116a-465a-40af-8424-9f30d2348ca9.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://space.bilibili.com/1293512903" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/219026751-d7d14cce-a7c9-4e82-9942-8375fca65b99.png" width="3%" alt="" /></a>
+  <img src="https://user-images.githubusercontent.com/25839884/218346358-56cc8e2f-a2b8-487f-9088-32480cceabcf.png" width="3%" alt="" />
+  <a href="https://www.zhihu.com/people/openmmlab" style="text-decoration:none;">
+    <img src="https://user-images.githubusercontent.com/25839884/219026120-ba71e48b-6e94-4bd4-b4e9-b7d175b5e362.png" width="3%" alt="" /></a>
+</div>
+
+## 特性
+
+- 支持各种视觉和语言指令数据
+- 使用 LoRA 进行参数高效微调
+- 同时调整视觉和语言，相互补充
+
+## 安装
+
+在一个已有环境中安装依赖包，运行以下指令
+
+```bash
+git clone https://github.com/open-mmlab/Multimodal-GPT.git
+cd Multimodal-GPT
+pip install -r requirements.txt
+pip install -v -e .
+```
+
+或者创建一个新的 conda 环境
+
+```bash
+conda env create -f environment.yml
+```
+
+## Demo
+
+1. 下载预训练权重
+
+    使用[这个脚本](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py)把 LLaMA 权重转换成 HuggingFace 格式。
+
+    从 [openflamingo/OpenFlamingo-9B](https://huggingface.co/openflamingo/OpenFlamingo-9B) 下载 OpenFlamingo 预训练模型。
+
+    从[这个链接](https://download.openmmlab.com/mmgpt/v0/mmgpt-lora-v0-release.pt) 下载我们的 LoRA 权重。
+
+    然后把所有模型权重放到 `checkpoints` 文件夹下，目录结构如下：
+
+    ```
+    checkpoints
+    ├── llama-7b_hf
+    │   ├── config.json
+    │   ├── pytorch_model-00001-of-00002.bin
+    │   ├── ......
+    │   └── tokenizer.model
+    ├── OpenFlamingo-9B
+    │   └──checkpoint.pt
+    ├──mmgpt-lora-v0-release.pt
+
+2. 启动 gradio demo
+
+    ```bash
+    python app.py
+    ```
+
+## 示例
+
+### 菜单：
+![image4](https://user-images.githubusercontent.com/12907710/234554562-8f3be88f-d563-47ba-97d9-ade8d47c46b0.png)
+
+### 旅行计划：
+![image3](https://user-images.githubusercontent.com/12907710/234523464-80c4e3f0-f99f-4498-96ef-dc43ef89c64b.png)
+
+### 电影：
+![image2](https://user-images.githubusercontent.com/12907710/234523468-e11905a6-491f-4b87-934f-90da7d14d1c3.png)
+
+### 名人：
+![image](https://user-images.githubusercontent.com/12907710/234523475-fd91f979-a344-4228-813f-6b55a1bc250f.png)
+
+
+## 微调 Fine-tuning
+
+### 准备数据集
+
+1. [A-OKVQA](https://allenai.org/project/a-okvqa/home)
+
+    从[这个链接](https://prior-datasets.s3.us-east-2.amazonaws.com/aokvqa/aokvqa_v1p0.tar.gz)下载标注，解压到 `data/aokvqa/annotations` 路径下。
+
+    同时还需要 coco 数据集的图像，可以从[这里](https://cocodataset.org/#home)下载。
+
+2. [COCO Caption](https://cs.stanford.edu/people/karpathy/deepimagesent/)
+
+    从[这个链接](https://cs.stanford.edu/people/karpathy/deepimagesent/coco.zip)，解压到 `data/coco` 路径下。
+
+    同时还需要 coco 数据集的图像，可以从[这里](https://cocodataset.org/#home)下载。
+
+3. [OCR VQA](https://ocr-vqa.github.io/)
+
+    从 [这个链接](https://drive.google.com/drive/folders/1_GYPY5UkUy7HIcR0zq3ZCFgeZN7BAfm_?usp=sharing) 下载数据集，放到 `data/OCR_VQA/` 路径下。
+
+4. [LlaVA](https://llava-vl.github.io/)
+
+    从 [liuhaotian/LLaVA-Instruct-150K](https://huggingface.co/datasets/liuhaotian/LLaVA-Instruct-150K) 下载数据集，放到 `data/llava/` 路径下。
+
+    同时还需要 coco 数据集的图像，可以从[这里](https://cocodataset.org/#home)下载。
+
+5. [Mini-GPT4](https://minigpt-4.github.io/)
+
+    从 [Vision-CAIR/cc_sbu_align](https://huggingface.co/datasets/Vision-CAIR/cc_sbu_align) 下载数据集，放到 `data/cc_sbu_align/` 路径下。
+
+6. [Dolly 15k](https://www.databricks.com/blog/2023/03/24/hello-dolly-democratizing-magic-chatgpt-open-models.html)
+
+    从 [databricks/databricks-dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k) 下载数据集，放到 `data/dolly/databricks-dolly-15k.jsonl` 路径下。
+
+7. [Alpaca GPT4](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+
+    从[这个链接](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM/raw/main/data/alpaca_gpt4_data.json) 下载数据集，放到 `data/alpaca_gpt4/alpaca_gpt4_data.json` 路径下。
+
+你也可以在 [configs/dataset_config.py](configs/dataset_config.py) 文件中自定义数据集路径。
+
+
+## 开启训练
+
+```bash
+torchrun --nproc_per_node=8 mmgpt/train/instruction_finetune.py \
+  --lm_path checkpoints/llama-7b_hf \
+  --tokenizer_path checkpoints/llama-7b_hf \
+  --pretrained_path checkpoints/OpenFlamingo-9B/checkpoint.pt \
+  --run_name train-my-gpt4 \
+  --learning_rate 1e-5 \
+  --lr_scheduler cosine \
+  --batch_size 1 \ 
+  --tuning_config configs/lora_config.py \
+  --dataset_config configs/dataset_config.py \
+  --report_to_wandb
+```
+
+
+## 致谢
+
+- [OpenFlamingo](https://github.com/mlfoundations/open_flamingo)
+- [LAVIS](https://github.com/salesforce/LAVIS)
+- [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca)
+- [MiniGPT-4](https://github.com/Vision-CAIR/MiniGPT-4)
+- [LLaVA](https://github.com/haotian-liu/LLaVA/tree/main)
+- [Instruction Tuning with GPT-4](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+
+如果你觉得我们的项目对你的研究和应用有帮助，请用以下 BibTeX 进行引用
+
+```bibtex
+@misc{gong2023multimodalgpt,
+      title={MultiModal-GPT: A Vision and Language Model for Dialogue with Humans}, 
+      author={Tao Gong and Chengqi Lyu and Shilong Zhang and Yudong Wang and Miao Zheng and Qian Zhao and Kuikun Liu and Wenwei Zhang and Ping Luo and Kai Chen},
+      year={2023},
+      eprint={2305.04790},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

app.py

@@ -0,0 +1,379 @@
+import os
+
+import gradio as gr
+import torch
+from PIL import Image
+
+from mmgpt.models.builder import create_model_and_transforms
+
+TEMPLATE = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
+response_split = "### Response:"
+
+
+class Inferencer:
+
+    def __init__(self, finetune_path, llama_path, open_flamingo_path):
+        ckpt = torch.load(finetune_path, map_location="cpu", weights_only=False)
+        if "model_state_dict" in ckpt:
+            state_dict = ckpt["model_state_dict"]
+            # remove the "module." prefix
+            state_dict = {
+                k[7:]: v
+                for k, v in state_dict.items() if k.startswith("module.")
+            }
+        else:
+            state_dict = ckpt
+        tuning_config = ckpt.get("tuning_config")
+        if tuning_config is None:
+            print("tuning_config not found in checkpoint")
+        else:
+            print("tuning_config found in checkpoint: ", tuning_config)
+        model, image_processor, tokenizer = create_model_and_transforms(
+            model_name="open_flamingo",
+            clip_vision_encoder_path="ViT-L-14",
+            clip_vision_encoder_pretrained="openai",
+            lang_encoder_path=llama_path,
+            tokenizer_path=llama_path,
+            pretrained_model_path=open_flamingo_path,
+            tuning_config=tuning_config,
+        )
+        model.load_state_dict(state_dict, strict=False)
+        model.half()
+        
+        
+        device = torch.device("cpu")
+        model = model.to(device)
+
+        # model = model.to("cuda")
+        model.eval()
+        tokenizer.padding_side = "left"
+        tokenizer.add_eos_token = False
+        self.model = model
+        self.image_processor = image_processor
+        self.tokenizer = tokenizer
+
+    def __call__(self, prompt, imgpaths, max_new_token, num_beams, temperature,
+                 top_k, top_p, do_sample):
+        device = torch.device("cpu")
+        if len(imgpaths) > 1:
+            raise gr.Error(
+                "Current only support one image, please clear gallery and upload one image"
+            )
+        lang_x = self.tokenizer([prompt], return_tensors="pt")
+        if len(imgpaths) == 0 or imgpaths is None:
+            for layer in self.model.lang_encoder._get_decoder_layers():
+                layer.condition_only_lang_x(True)
+            output_ids = self.model.lang_encoder.generate(
+                input_ids=lang_x["input_ids"].to(device),
+                attention_mask=lang_x["attention_mask"].to(device),
+                max_new_tokens=max_new_token,
+                num_beams=num_beams,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                do_sample=do_sample,
+            )[0]
+            for layer in self.model.lang_encoder._get_decoder_layers():
+                layer.condition_only_lang_x(False)
+        else:
+            images = (Image.open(fp) for fp in imgpaths)
+            vision_x = [self.image_processor(im).unsqueeze(0) for im in images]
+            vision_x = torch.cat(vision_x, dim=0)
+            vision_x = vision_x.unsqueeze(1).unsqueeze(0).half()
+
+            output_ids = self.model.generate(
+                vision_x=vision_x.to(device),
+                lang_x=lang_x["input_ids"].to(device),
+                attention_mask=lang_x["attention_mask"].to(device),
+                max_new_tokens=max_new_token,
+                num_beams=num_beams,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                do_sample=do_sample,
+            )[0]
+        generated_text = self.tokenizer.decode(
+            output_ids, skip_special_tokens=True)
+        # print(generated_text)
+        result = generated_text.split(response_split)[-1].strip()
+        return result
+
+
+class PromptGenerator:
+
+    def __init__(
+        self,
+        prompt_template=TEMPLATE,
+        ai_prefix="Response",
+        user_prefix="Instruction",
+        sep: str = "\n\n### ",
+        buffer_size=0,
+    ):
+        self.all_history = list()
+        self.ai_prefix = ai_prefix
+        self.user_prefix = user_prefix
+        self.buffer_size = buffer_size
+        self.prompt_template = prompt_template
+        self.sep = sep
+
+    def add_message(self, role, message):
+        self.all_history.append([role, message])
+
+    def get_images(self):
+        img_list = list()
+        if self.buffer_size > 0:
+            all_history = self.all_history[-2 * (self.buffer_size + 1):]
+        elif self.buffer_size == 0:
+            all_history = self.all_history[-2:]
+        else:
+            all_history = self.all_history[:]
+        for his in all_history:
+            if type(his[-1]) == tuple:
+                img_list.append(his[-1][-1])
+        return img_list
+
+    def get_prompt(self):
+        format_dict = dict()
+        if "{user_prefix}" in self.prompt_template:
+            format_dict["user_prefix"] = self.user_prefix
+        if "{ai_prefix}" in self.prompt_template:
+            format_dict["ai_prefix"] = self.ai_prefix
+        prompt_template = self.prompt_template.format(**format_dict)
+        ret = prompt_template
+        if self.buffer_size > 0:
+            all_history = self.all_history[-2 * (self.buffer_size + 1):]
+        elif self.buffer_size == 0:
+            all_history = self.all_history[-2:]
+        else:
+            all_history = self.all_history[:]
+        context = []
+        have_image = False
+        for role, message in all_history[::-1]:
+            if message:
+                if type(message) is tuple and message[
+                        1] is not None and not have_image:
+                    message, _ = message
+                    context.append(self.sep + "Image:\n<image>" + self.sep +
+                                   role + ":\n" + message)
+                else:
+                    context.append(self.sep + role + ":\n" + message)
+            else:
+                context.append(self.sep + role + ":\n")
+
+        ret += "".join(context[::-1])
+        return ret
+
+
+def to_gradio_chatbot(prompt_generator):
+    ret = []
+    for i, (role, msg) in enumerate(prompt_generator.all_history):
+        if i % 2 == 0:
+            if type(msg) is tuple:
+                import base64
+                from io import BytesIO
+
+                msg, image = msg
+                if type(image) is str:
+                    from PIL import Image
+
+                    image = Image.open(image)
+                max_hw, min_hw = max(image.size), min(image.size)
+                aspect_ratio = max_hw / min_hw
+                max_len, min_len = 800, 400
+                shortest_edge = int(
+                    min(max_len / aspect_ratio, min_len, min_hw))
+                longest_edge = int(shortest_edge * aspect_ratio)
+                H, W = image.size
+                if H > W:
+                    H, W = longest_edge, shortest_edge
+                else:
+                    H, W = shortest_edge, longest_edge
+                image = image.resize((H, W))
+                # image = image.resize((224, 224))
+                buffered = BytesIO()
+                image.save(buffered, format="JPEG")
+                img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                msg = msg + img_str
+            ret.append([msg, None])
+        else:
+            ret[-1][-1] = msg
+    return ret
+
+
+def bot(
+    text,
+    image,
+    state,
+    prompt,
+    ai_prefix,
+    user_prefix,
+    seperator,
+    history_buffer,
+    max_new_token,
+    num_beams,
+    temperature,
+    top_k,
+    top_p,
+    do_sample,
+):
+    state.prompt_template = prompt
+    state.ai_prefix = ai_prefix
+    state.user_prefix = user_prefix
+    state.sep = seperator
+    state.buffer_size = history_buffer
+    if image:
+        state.add_message(user_prefix, (text, image))
+    else:
+        state.add_message(user_prefix, text)
+    state.add_message(ai_prefix, None)
+    inputs = state.get_prompt()
+    image_paths = state.get_images()[-1:]
+
+    inference_results = inferencer(inputs, image_paths, max_new_token,
+                                   num_beams, temperature, top_k, top_p,
+                                   do_sample)
+    state.all_history[-1][-1] = inference_results
+    memory_allocated = str(round(torch.cuda.memory_allocated() / 1024**3,
+                                 2)) + 'GB'
+    return state, to_gradio_chatbot(state), "", None, inputs, memory_allocated
+
+
+def clear(state):
+    state.all_history = []
+    return state, to_gradio_chatbot(state), "", None, ""
+
+
+title_markdown = ("""
+    # 🤖 Multi-modal GPT
+    [[Project]](https://github.com/open-mmlab/Multimodal-GPT.git)""")
+
+
+def build_conversation_demo():
+    with gr.Blocks(title="Multi-modal GPT") as demo:
+        gr.Markdown(title_markdown)
+
+        state = gr.State(PromptGenerator())
+        with gr.Row():
+            with gr.Column(scale=3):
+                memory_allocated = gr.Textbox(
+                    value=init_memory, label="Memory")
+                imagebox = gr.Image(type="filepath")
+                # TODO config parameters
+                with gr.Accordion(
+                        "Parameters",
+                        open=True,
+                ):
+                    max_new_token_bar = gr.Slider(
+                        0, 1024, 512, label="max_new_token", step=1)
+                    num_beams_bar = gr.Slider(
+                        0.0, 10, 3, label="num_beams", step=1)
+                    temperature_bar = gr.Slider(
+                        0.0, 1.0, 1.0, label="temperature", step=0.01)
+                    topk_bar = gr.Slider(0, 100, 20, label="top_k", step=1)
+                    topp_bar = gr.Slider(0, 1.0, 1.0, label="top_p", step=0.01)
+                    do_sample = gr.Checkbox(True, label="do_sample")
+                with gr.Accordion(
+                        "Prompt",
+                        open=False,
+                ):
+                    with gr.Row():
+                        ai_prefix = gr.Text("Response", label="AI Prefix")
+                        user_prefix = gr.Text(
+                            "Instruction", label="User Prefix")
+                        seperator = gr.Text("\n\n### ", label="Seperator")
+                    history_buffer = gr.Slider(
+                        -1, 10, -1, label="History buffer", step=1)
+                    prompt = gr.Text(TEMPLATE, label="Prompt")
+                    model_inputs = gr.Textbox(label="Actual inputs for Model")
+
+            with gr.Column(scale=6):
+                with gr.Row():
+                    with gr.Column():
+                         chatbot = gr.Chatbot(elem_id="chatbot", type="messages")
+                with gr.Row():
+                    with gr.Column(scale=8):
+                        textbox = gr.Textbox(
+                            show_label=False,
+                            placeholder="Enter text and press ENTER",
+                            container=False)
+                        submit_btn = gr.Button(value="Submit")
+                        clear_btn = gr.Button(value="🗑️  Clear history")
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        gr.Examples(
+            examples=[
+                [
+                    f"{cur_dir}/docs/images/demo_image.jpg",
+                    "What is in this image?"
+                ],
+            ],
+            inputs=[imagebox, textbox],
+        )
+        textbox.submit(
+            bot,
+            [
+                textbox,
+                imagebox,
+                state,
+                prompt,
+                ai_prefix,
+                user_prefix,
+                seperator,
+                history_buffer,
+                max_new_token_bar,
+                num_beams_bar,
+                temperature_bar,
+                topk_bar,
+                topp_bar,
+                do_sample,
+            ],
+            [
+                state, chatbot, textbox, imagebox, model_inputs,
+                memory_allocated
+            ],
+        )
+        submit_btn.click(
+            bot,
+            [
+                textbox,
+                imagebox,
+                state,
+                prompt,
+                ai_prefix,
+                user_prefix,
+                seperator,
+                history_buffer,
+                max_new_token_bar,
+                num_beams_bar,
+                temperature_bar,
+                topk_bar,
+                topp_bar,
+                do_sample,
+            ],
+            [
+                state, chatbot, textbox, imagebox, model_inputs,
+                memory_allocated
+            ],
+        )
+        clear_btn.click(clear, [state],
+                        [state, chatbot, textbox, imagebox, model_inputs])
+    return demo
+
+
+if __name__ == "__main__":
+    llama_path = "checkpoints/llama-7b_hf"
+    open_flamingo_path = "checkpoints/OpenFlamingo-9B/checkpoint.pt"
+    finetune_path = "checkpoints/mmgpt-lora-v0-release.pt"
+
+    inferencer = Inferencer(
+        llama_path=llama_path,
+        open_flamingo_path=open_flamingo_path,
+        finetune_path=finetune_path)
+    init_memory = str(round(torch.cuda.memory_allocated() / 1024**3, 2)) + 'GB'
+    demo = build_conversation_demo()
+    demo.queue(max_size=3)
+    IP = "0.0.0.0"
+    PORT = 8997
+    demo.launch(server_name=IP, server_port=PORT, share=True)
+
+

configs/dataset_config.py

@@ -0,0 +1,64 @@
+visual_datasets = [
+    dict(
+        type="llava",
+        vis_root="data/coco/train2017",
+        ann_paths=[
+            "data/llava/detail_23k.json",
+            "data/llava/complex_reasoning_77k.json",
+        ],
+    ),
+    dict(
+        type="llava_dial",
+        vis_root="data/coco/train2017",
+        ann_paths=[
+            "data/llava/conversation_58k.json",
+        ],
+    ),
+    dict(
+        type="aokvqa",
+        vis_root="data/coco/images",
+        ann_paths=[
+            "data/aokvqa/annotations/aokvqa_v1p0_train.json",
+        ],
+        sample=5000,
+    ),
+    dict(
+        type="minigpt4",
+        vis_root="data/cc_sbu_align/image",
+        ann_paths=[
+            "data/cc_sbu_align/filter_cap.json",
+        ],
+    ),
+    dict(
+        type="coco_caption",
+        vis_root="data/coco",
+        ann_paths=[
+            "data/coco/annotations/coco_karpathy_train_converted.json",
+            "data/coco/annotations/coco_karpathy_val.json",
+        ],
+        sample=512,
+    ),
+    dict(
+        type="ocr_vqa",
+        vis_root="data/OCR_VQA/image",
+        ann_paths=[
+            "data/OCR_VQA/downloaded_dataset.json",
+        ],
+        sample=512,
+    ),
+]
+
+language_datasets = [
+    dict(
+        type="dolly",
+        ann_path="data/dolly/databricks-dolly-15k.jsonl",
+    ),
+    dict(
+        type="alpaca_gpt4",
+        ann_path="data/alpaca_gpt4/alpaca_gpt4_data.json",
+    ),
+    dict(
+        type="baize",
+        ann_path="data/baize/quora_chat_data.json",
+    ),
+]

configs/lora_config.py

@@ -0,0 +1,9 @@
+tuning_config = dict(
+    lora=True,
+    lora_target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "to_q", "to_kv", "to_out", "ff.1", "ff.3"],
+    lora_r=16,
+    lora_alpha=16,
+    lora_dropout=0.0,
+    vis=True,
+    unfrozen=[],
+)

environment.yml

@@ -0,0 +1,10 @@
+name: mmgpt
+channels:
+  - defaults
+dependencies:
+  - python=3.9
+  - conda-forge::openjdk
+  - pip
+  - pip:
+    - -r requirements.txt
+    - -e .

mmgpt/__init__.py

@@ -0,0 +1,2 @@
+from .models.builder import create_model_and_transforms
+from .models.open_flamingo import Flamingo

mmgpt/datasets/__init__.py

@@ -0,0 +1,4 @@
+from .builder import build_dataset  # noqa: F401
+from .dial_dataset import DialDataset  # noqa: F401
+from .samplers import InfiniteSampler  # noqa: F401
+from .vqa_dataset import VQADataset  # noqa: F401

mmgpt/datasets/alpaca_gpt4_dataset.py

@@ -0,0 +1,26 @@
+import json
+
+from mmgpt.datasets.dolly_dataset import DollyDataset
+
+
+class AlpacaGPT4Dataset(DollyDataset):
+    """
+    ```json
+    [
+        {
+            "instruction": "Identify the odd one out.",
+            "input": "Twitter, Instagram, Telegram",
+            "output": "The odd one out is Telegram. Twitter and Instagram are social media platforms mainly for sharing information, images and videos while Telegram is a cloud-based instant messaging and voice-over-IP service."
+        },
+    ]
+    """
+
+    def load_annotation(self, ann_path):
+        self.annotation = json.load(open(ann_path, "r"))
+
+    def process_text(self, ann):
+        instruction = ann["instruction"]
+        input = ann["input"]
+        output = ann["output"]
+        instruction = self.prompter(instruction=instruction, input=input)
+        return dict(instruction=instruction, answer=output)

mmgpt/datasets/aokvqa_dataset.py

@@ -0,0 +1,51 @@
+import random
+
+from .vqa_dataset import VQADataset
+
+REASON_QUESTIONS = [
+    "Why?",
+    "Why is this?",
+    "And why?",
+    "What is the reason?",
+    "And can you tell me why?",
+    "Can you tell me why?",
+    "Can you tell me the reason?",
+]
+
+
+class AOKVQADataset(VQADataset):
+    def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, **kwargs):
+        super().__init__(tokenizer, vis_processor, vis_root, ann_paths, **kwargs)
+
+    def process_text(self, ann):
+        question = ann["question"]
+        question = question + " " + random.choice(REASON_QUESTIONS)
+
+        choices = ann["choices"]
+        true_answer = choices[ann["correct_choice_idx"]]
+        answer = "The answer is " + true_answer + ". Because " + " ".join(ann["rationales"])
+
+        is_option = random.random() < self.option_prob and len(choices) > 1
+        if is_option:
+            instruction = self.prompter(question, choices)
+        else:
+            instruction = self.prompter(question)
+
+        instruction = self.prompter(question)
+        return dict(instruction=instruction, answer=answer)
+
+
+def build_aokvqa_dataset(
+    tokenizer,
+    vis_processor,
+    vis_root="data/coco/images",
+    ann_paths=["data/aokvqa/annotations/aokvqa_v1p0_train.json"],
+    sample_image=False,
+):
+    return AOKVQADataset(
+        tokenizer=tokenizer,
+        vis_processor=vis_processor,
+        vis_root=vis_root,
+        ann_paths=ann_paths,
+        sample_image=sample_image,
+    )

mmgpt/datasets/baize_dataset.py

@@ -0,0 +1,86 @@
+import json
+
+from mmgpt.datasets.dolly_dataset import DollyDataset
+
+
+TEMPLATE = {
+    "description": "Template used by Alpaca-LoRA.",
+    "prompt_choice": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{question}\n\n### Input:\n{options}\n\n### Response:\n",
+    "prompt_qa": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{question}\n\n### Response:\n",
+    "prompt_dial": "\n\n### Instruction:\n{question}\n\n### Response:\n",
+    "response_split": "### Response:",
+}
+
+class LangDialPrompter:
+    def __call__(self, question, options=None):
+        if options:
+            options = ", ".join(options)
+            res = TEMPLATE["prompt_choice"].format(image="<image>", question=question, options=options)
+        else:
+            res = TEMPLATE["prompt_dial"].format(question=question)
+        return res
+
+    def get_response(self, output: str) -> str:
+        return output.split(TEMPLATE["response_split"])[-1].strip()
+
+class BaiZeDataset(DollyDataset):
+    """
+    ```json
+    [
+        {
+            "instruction": "Identify the odd one out.",
+            "input": "Twitter, Instagram, Telegram",
+            "output": "The odd one out is Telegram. Twitter and Instagram are social media platforms mainly for sharing information, images and videos while Telegram is a cloud-based instant messaging and voice-over-IP service."
+        },
+    ]
+    """
+    def __init__(self, *args, **kwargs):
+        super(BaiZeDataset, self).__init__(*args, **kwargs)
+        self.prompter = LangDialPrompter()
+
+    def load_annotation(self, ann_path):
+        self.annotation = json.load(open(ann_path, "r"))
+
+    def process_text(self, anns):
+        # TODO remove this
+        begin_string = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
+        convs = anns['input'].split("[|Human|] ")
+        conv_list = []
+        for conv_id, one_conv in enumerate(convs[1:-1]):
+            question, answer = one_conv.split("[|AI|] ")
+            question = question.replace("\n", "")
+            answer = answer.replace("\n", "")
+            instruction = self.prompter(question)
+            if conv_id == 0:
+                single_conv = dict(instruction=begin_string + instruction, answer=answer)
+            else:
+                single_conv = dict(instruction=instruction, answer=answer)
+            conv_list.append(single_conv)
+        return conv_list
+    
+    def __getitem__(self, index):
+        ann = self.annotation[index]
+        text_list = self.process_text(ann)
+        res_list = []
+        for text in text_list:
+            single_res = self.tokenize(text)
+            single_res["instruction"] = text["instruction"]
+            single_res["answer"] = text["answer"]
+            res_list.append(single_res)
+
+        input_ids = []
+        attention_mask = []
+        labels = []
+        instruction = []
+        answer = []
+        for res in res_list:
+            input_ids.extend(res["input_ids"])
+            attention_mask.extend(res["attention_mask"])
+            labels.extend(res["labels"])
+            instruction.append(res["instruction"])
+            answer.append(res["answer"])
+
+        res = dict(
+            input_ids=input_ids, attention_mask=attention_mask, labels=labels, instruction=instruction, answer=answer
+        )
+        return res

mmgpt/datasets/builder.py

@@ -0,0 +1,126 @@
+import numpy as np
+import torch
+
+from .alpaca_gpt4_dataset import AlpacaGPT4Dataset  # noqa: F401
+from .aokvqa_dataset import AOKVQADataset  # noqa: F401
+from .cc_sbu_align_dataset import CcSbuAlignDataset  # noqa: F401
+from .clevr_dataset import CLEVRDataset  # noqa: F401
+from .coco_caption_dataset import COCOCaptionDataset  # noqa: F401
+from .dial_dataset import DialDataset  # noqa: F401
+from .dolly_dataset import DollyDataset  # noqa: F401
+from .gqa_dataset import GQADataset  # noqa: F401
+from .llava_dataset import LlavaDataset  # noqa: F401
+from .nlvr_dataset import NLVRv1Dataset, NLVRv2Dataset  # noqa: F401
+from .ocr_vqa_dataset import OCRVQADataset  # noqa: F401
+from .snli_ve_datasets import SNLIVEDataset  # noqa: F401
+from .text_ocr_dataset import TextOCRDataset  # noqa: F401
+from .vqa_dataset import ConcatDataset, VQADataset  # noqa: F401
+from .baize_dataset import BaiZeDataset  # noqa: F401
+
+
+def build_dataset(dataset_config, **kwargs):
+    if isinstance(dataset_config, list):
+        datasets = [build_dataset(cfg, **kwargs) for cfg in dataset_config]
+        return ConcatDataset(datasets)
+    dataset_type = dataset_config.pop("type")
+    sample = dataset_config.pop("sample", -1)
+    if dataset_type == "llava":
+        dataset = LlavaDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "vqa":
+        dataset = VQADataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "minigpt4":
+        dataset = CcSbuAlignDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "llava_dial":
+        dataset = DialDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "coco_dial":
+        dataset = DialDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "aokvqa":
+        dataset = AOKVQADataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "okvqa":
+        dataset = VQADataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "text_ocr":
+        dataset = TextOCRDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "ocr_vqa":
+        dataset = OCRVQADataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "coco_caption":
+        dataset = COCOCaptionDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "gqa":
+        dataset = GQADataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "clevr":
+        dataset = CLEVRDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "nlvrv1":
+        dataset = NLVRv1Dataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "nlvrv2":
+        dataset = NLVRv2Dataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "snlive":
+        dataset = SNLIVEDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "dolly":
+        dataset = DollyDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "alpaca_gpt4":
+        dataset = AlpacaGPT4Dataset(
+            **dataset_config,
+            **kwargs,
+        )
+    elif dataset_type == "baize":
+        dataset = BaiZeDataset(
+            **dataset_config,
+            **kwargs,
+        )
+    else:
+        raise NotImplementedError
+
+    if sample > 0:
+        random_indices = np.random.choice(len(dataset), min(sample, len(dataset)), replace=False)
+        subsample_dataset = torch.utils.data.Subset(dataset, random_indices)
+        subsample_dataset.collater = dataset.collater
+        return subsample_dataset
+    else:
+        return dataset

mmgpt/datasets/cc_sbu_align_dataset.py

@@ -0,0 +1,107 @@
+import json
+import os
+import random
+
+from PIL import Image
+
+from .vqa_dataset import VQADataset, VQAPrompter
+
+QUESTIONS = [
+    "please describe the image",
+    "can you describe the image",
+    "Could you provide a description of the image?",
+    "What do you see in this image?",
+    "Share your thoughts on the content of the image.",
+    "Please narrate what's happening in the picture.",
+    "Can you give a brief explanation of the image?",
+    "Describe the main elements and details present in the image.",
+    "In your own words, what is depicted in the image?",
+    "Can you outline the key aspects of the image?",
+    "What are the most striking features in this image?",
+    "Please provide a summary of the image's content.",
+    "Describe the overall theme or concept captured in the image.",
+    "How would you explain the image's composition and focus?",
+    "What is the focal point or main subject of the image?",
+    "How do the different components of the image interact with each other?",
+    "What would be a fitting caption for this image?",
+    "Can you create a concise description that captures the essence of the image?",
+    "How would you briefly summarize the content of this image in a phrase or sentence?",
+    "Please provide a catchy and relevant caption for this picture.",
+    "If you were to give this image a title, what would it be?",
+    "Describe the image in one creative sentence.",
+    "Please suggest a memorable phrase that encapsulates the image's content.",
+    "What engaging phrase would best represent this image?",
+    "Can you create an expressive caption that highlights the main theme of the image?",
+    "How would you sum up the image's story for a caption?",
+    "Provide an eye-catching caption that conveys the image's core message.",
+    "If you were to give this image a headline, what would it say?",
+    "Can you craft a captivating caption that communicates the essence of the image?",
+    "How would you describe the image's content in a powerful caption?",
+    "Please provide an inventive title to summarize the scene depicted in the image.",
+    "Compose a concise and striking phrase that reflects the image's key elements.",
+    "If you were to create a caption for this image, what would it be?",
+    "Offer a compelling caption that highlights the central focus of the image.",
+    "Can you produce a unique caption that encapsulates the image's overall mood?",
+    "Please generate an attention-grabbing caption that would best illustrate the events captured in this image",
+    "How would you express the image's main idea in an impactful sentence?",
+    "Please create a vivid and concise title that conveys the essence of the picture.",
+    "Compose an imaginative caption that reflects the image's most striking features.",
+    "What memorable statement would best represent the scene illustrated in this image?",
+    "Draft an evocative caption that brings the image to life for the reader.",
+    "Can you suggest an insightful caption that highlights the underlying message of the image?",
+    "What engaging phrase would effectively convey the action or subject matter depicted in this picture?",
+    "How would you encapsulate the image's core theme in a concise and expressive manner?",
+    "Please provide a creative and impactful title that captures the spirit of the image.",
+    "Craft a captivating caption that showcases the image's most prominent attributes.",
+    "What intriguing statement would best sum up the scene presented in this image?",
+    "Develop a descriptive caption that paints a vivid picture for the viewer.",
+    "Can you give a detailed account of the image's contents?",
+    "What are the key elements and features visible in this image?",
+    "How would you narrate the events or actions depicted in the picture?",
+    "Please share your observations about the various components present in the image.",
+    "What is the overall theme or concept captured in this image? Can you describe it?",
+]
+
+
+class CcSbuAlignDataset(VQADataset):
+    def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, add_eos=True, ignore_instruction=True):
+        self.tokenizer = tokenizer
+        self.vis_root = vis_root
+
+        self.annotation = []
+        for ann_path in ann_paths:
+            self.annotation.extend(json.load(open(ann_path, "r"))["annotations"])
+
+        self.vis_processor = vis_processor
+        self.prompter = VQAPrompter()
+        self.add_eos = add_eos
+        self.ignore_instruction = ignore_instruction
+
+    def process_text(self, ann):
+        # random select a question
+        question = random.choice(QUESTIONS)
+        answer = ann["caption"]
+        instruction = self.prompter(question)
+        return dict(instruction=instruction, answer=answer)
+
+    def process_image(self, ann):
+        image_path = os.path.join(self.vis_root, ann["image_id"] + ".jpg")
+        image = Image.open(image_path).convert("RGB")
+
+        image = self.vis_processor(image)
+        return image
+
+
+def build_ccsbualign_dataset(
+    tokenizer,
+    vis_processor,
+    vis_root="data/cc_sbu_align/image/",
+    ann_paths=["data/cc_sbu_align/filter_cap.json"],
+    **kwargs,
+):
+    return CcSbuAlignDataset(
+        tokenizer=tokenizer,
+        vis_processor=vis_processor,
+        vis_root=vis_root,
+        ann_paths=ann_paths,
+    )

mmgpt/datasets/clevr_dataset.py

@@ -0,0 +1,74 @@
+import json
+import os
+import random
+from collections import defaultdict
+
+from PIL import Image
+
+from .vqa_dataset import VQADataset
+
+
+class CLEVRDataset(VQADataset):
+    """Visual Reasoning Dataset. It also contains Dialog.
+
+    Note: The image is a little bit simple. with several objects and simple background.
+    """
+
+    def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, **kwargs):
+        super().__init__(tokenizer, vis_processor, vis_root, ann_paths=[], **kwargs)
+
+        self.annotation = self.load_annotations(ann_paths)
+        if self.sample_image:
+            print("randomly sample one annotation for each image")
+            self.annotation = self.parse_annotation(self.annotation)
+        self._add_instance_ids()
+
+    @staticmethod
+    def load_annotations(ann_paths):
+        annotation = []
+        for ann_path in ann_paths:
+            ann = json.load(open(ann_path, "r"))
+            annotation.extend(ann["questions"])
+        return annotation
+
+    def parse_annotation(self, annotation):
+        image_list = defaultdict(list)
+        for ann in annotation:
+            image_list[ann["image_filename"]].append(ann)
+        annotation = []
+        for ann_list in image_list.values():
+            annotation.append(random.choice(ann_list))
+        return annotation
+
+    def process_text(self, ann):
+        question = ann["question"]
+        answer = ann["answer"]
+        instruction = self.prompter(question)
+        return dict(instruction=instruction, answer=answer)
+
+    def process_image(self, ann):
+        split = ann["split"]
+        image_path = os.path.join(self.vis_root, split, ann["image_filename"])
+        image = Image.open(image_path).convert("RGB")
+
+        image = self.vis_processor(image)
+        return image
+
+
+def build_clevr_dataset(
+    tokenizer,
+    vis_processor,
+    vis_root="data/clevr/CLEVR_v1.0/images",
+    ann_paths=[
+        "data/clevr/CLEVR_v1.0/questions/CLEVR_train_questions.json",
+        "data/clevr/CLEVR_v1.0/questions/CLEVR_val_questions.json",
+    ],
+    sample_image=False,
+):
+    return CLEVRDataset(
+        tokenizer=tokenizer,
+        vis_processor=vis_processor,
+        vis_root=vis_root,
+        ann_paths=ann_paths,
+        sample_image=sample_image,
+    )

mmgpt/datasets/coco_caption_dataset.py

@@ -0,0 +1,119 @@
+"""
+ Copyright (c) 2022, salesforce.com, inc.
+ All rights reserved.
+ SPDX-License-Identifier: BSD-3-Clause
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+"""
+
+import json
+import os
+import random
+
+import numpy as np
+from PIL import Image
+from transformers import LlamaTokenizer
+
+from .vqa_dataset import VQADataset
+
+QUESTIONS = [
+    "please describe the image",
+    "can you describe the image",
+    "Could you provide a description of the image?",
+    "What do you see in this image?",
+    "Share your thoughts on the content of the image.",
+    "Please narrate what's happening in the picture.",
+    "Can you give a brief explanation of the image?",
+    "Describe the main elements and details present in the image.",
+    "In your own words, what is depicted in the image?",
+    "Can you outline the key aspects of the image?",
+    "What are the most striking features in this image?",
+    "Please provide a summary of the image's content.",
+    "Describe the overall theme or concept captured in the image.",
+    "How would you explain the image's composition and focus?",
+    "What is the focal point or main subject of the image?",
+    "How do the different components of the image interact with each other?",
+    "What would be a fitting caption for this image?",
+    "Can you create a concise description that captures the essence of the image?",
+    "How would you briefly summarize the content of this image in a phrase or sentence?",
+    "Please provide a catchy and relevant caption for this picture.",
+    "If you were to give this image a title, what would it be?",
+    "Describe the image in one creative sentence.",
+    "Please suggest a memorable phrase that encapsulates the image's content.",
+    "What engaging phrase would best represent this image?",
+    "Can you create an expressive caption that highlights the main theme of the image?",
+    "How would you sum up the image's story for a caption?",
+    "Provide an eye-catching caption that conveys the image's core message.",
+    "If you were to give this image a headline, what would it say?",
+    "Can you craft a captivating caption that communicates the essence of the image?",
+    "How would you describe the image's content in a powerful caption?",
+    "Please provide an inventive title to summarize the scene depicted in the image.",
+    "Compose a concise and striking phrase that reflects the image's key elements.",
+    "If you were to create a caption for this image, what would it be?",
+    "Offer a compelling caption that highlights the central focus of the image.",
+    "Can you produce a unique caption that encapsulates the image's overall mood?",
+    "Please generate an attention-grabbing caption that would best illustrate the events captured in this image",
+    "How would you express the image's main idea in an impactful sentence?",
+    "Please create a vivid and concise title that conveys the essence of the picture.",
+    "Compose an imaginative caption that reflects the image's most striking features.",
+    "What memorable statement would best represent the scene illustrated in this image?",
+    "Draft an evocative caption that brings the image to life for the reader.",
+    "Can you suggest an insightful caption that highlights the underlying message of the image?",
+    "What engaging phrase would effectively convey the action or subject matter depicted in this picture?",
+    "How would you encapsulate the image's core theme in a concise and expressive manner?",
+    "Please provide a creative and impactful title that captures the spirit of the image.",
+    "Craft a captivating caption that showcases the image's most prominent attributes.",
+    "What intriguing statement would best sum up the scene presented in this image?",
+    "Develop a descriptive caption that paints a vivid picture for the viewer.",
+    "Can you give a detailed account of the image's contents?",
+    "What are the key elements and features visible in this image?",
+    "How would you narrate the events or actions depicted in the picture?",
+    "Please share your observations about the various components present in the image.",
+    "What is the overall theme or concept captured in this image? Can you describe it?",
+]
+
+
+class COCOCaptionDataset(VQADataset):
+    def __init__(
+        self, tokenizer, vis_processor=None, vis_root=None, ann_paths=[], add_eos=True, ignore_instruction=True
+    ):
+        """
+        vis_root (string): Root directory of images (e.g. coco/images/)
+        ann_root (string): directory to store the annotation file
+        """
+        self.tokenizer: LlamaTokenizer = tokenizer
+        self.vis_root = vis_root
+
+        self.annotation = []
+        for ann_path in ann_paths:
+            self.annotation.extend(json.load(open(ann_path, "r")))
+
+        self.vis_processor = vis_processor
+
+        instructions = []
+        for question in QUESTIONS:
+            # instruction = f"Below is a question about an image. Write a response to answer the question.\n\n### Image:\n<image>\n\n### Question:\n{question}\n\n### Answer:\n".format(
+            #    question
+            # )
+            instruction = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Image:\n{image}\n\n### Instruction:\n{question}\n\n### Response:\n".format(
+                image="<image>", question=question
+            )
+            instructions.append(instruction)
+        self.instructions = instructions
+        self.add_eos = add_eos
+        self.ignore_instruction = ignore_instruction
+
+    def process_image(self, ann):
+        image_path = os.path.join(self.vis_root, ann["image"])
+        image = Image.open(image_path).convert("RGB")
+
+        image = self.vis_processor(image)
+        return image
+
+    def process_text(self, ann):
+        all_captions = ann["caption"]
+        if not isinstance(all_captions, list):
+            all_captions = [all_captions]
+        caption = random.choice(all_captions)
+        instruction = random.choice(self.instructions)
+
+        return dict(instruction=instruction, answer=caption)

mmgpt/datasets/dial_dataset.py

@@ -0,0 +1,83 @@
+from .vqa_dataset import VQADataset
+
+TEMPLATE = {
+    "description": "Template used by Alpaca-LoRA.",
+    # "prompt_choice": "Below is a multiple choice question about an image, along with answer options. Please choose the correct answer from these options.\n\n### Image:\n{image}\n\n### Question:\n{question}\n\n### Options:\n{options}\n\n### Answer:\n",
+    # "prompt_qa": "Below is a question about an image. Write a response to answer the question.\n\n### Image:\n{image}\n\n### Question:\n{question}\n\n### Answer:\n",
+    "prompt_choice": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Image:\n{image}\n\n### Instruction:\n{question}\n\n### Input:\n{options}\n\n### Response:\n",
+    "prompt_qa": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Image:\n{image}\n\n### Instruction:\n{question}\n\n### Response:\n",
+    "prompt_dial": "\n\n### Instruction:\n{question}\n\n### Response:\n",
+    "response_split": "### Response:",
+}
+
+
+class DialPrompter:
+    def __call__(self, question, options=None):
+        if options:
+            options = ", ".join(options)
+            res = TEMPLATE["prompt_choice"].format(image="<image>", question=question, options=options)
+        else:
+            res = TEMPLATE["prompt_dial"].format(question=question)
+        return res
+
+    def get_response(self, output: str) -> str:
+        return output.split(TEMPLATE["response_split"])[-1].strip()
+
+
+class DialDataset(VQADataset):
+    def __init__(self, *args, **kwargs):
+        super(DialDataset, self).__init__(*args, **kwargs)
+        self.prompter = DialPrompter()
+
+    def _add_instance_ids(self, key="id"):
+        for idx, ann in enumerate(self.annotation):
+            ann[key] = str(idx)
+
+    def process_text(self, anns):
+        # TODO remove this
+        begin_string = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Image:\n{image}".format(
+            image="<image>"
+        )
+        num_convs = len(anns["conversations"]) // 2
+        conv_list = []
+        for conv_id in range(num_convs):
+            question = anns["conversations"][conv_id]["value"]
+            # remove '<image>' tag and '\n'
+            question = question.replace("<image>", "").replace("\n", "")
+            answer = anns["conversations"][conv_id + 1]["value"]
+            instruction = self.prompter(question)
+            if conv_id == 0:
+                single_conv = dict(instruction=begin_string + instruction, answer=answer)
+            else:
+                single_conv = dict(instruction=instruction, answer=answer)
+            conv_list.append(single_conv)
+        return conv_list
+
+    def __getitem__(self, index):
+        ann = self.annotation[index]
+        image = self.process_image(ann)
+        text_list = self.process_text(ann)
+        res_list = []
+        for text in text_list:
+            single_res = self.tokenize(text)
+            single_res["instruction"] = text["instruction"]
+            single_res["answer"] = text["answer"]
+            res_list.append(single_res)
+
+        input_ids = []
+        attention_mask = []
+        labels = []
+        instruction = []
+        answer = []
+        for res in res_list:
+            input_ids.extend(res["input_ids"])
+            attention_mask.extend(res["attention_mask"])
+            labels.extend(res["labels"])
+            instruction.extend(res["instruction"])
+            answer.extend(res["answer"])
+
+        res = dict(
+            input_ids=input_ids, attention_mask=attention_mask, labels=labels, instruction=instruction, answer=answer
+        )
+        res.update(image=image)
+        return res

mmgpt/datasets/dolly_dataset.py

@@ -0,0 +1,150 @@
+import copy
+import json
+
+import numpy as np
+from torch.utils.data import Dataset
+from transformers import LlamaTokenizer
+
+TEMPLATE = {
+    "description": "Template used by LLM.",
+    "prompt_no_input_format": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n",
+    "prompt_with_input_format": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
+    "response_split": "### Response:",
+}
+
+
+class LMPrompter:
+    def __call__(self, instruction, input=None):
+        if input is None or len(input) == 0:
+            return TEMPLATE["prompt_no_input_format"].format(instruction=instruction)
+        else:
+            return TEMPLATE["prompt_with_input_format"].format(instruction=instruction, input=input)
+
+    def get_response(self, output: str) -> str:
+        return output.split(TEMPLATE["response_split"])[-1].strip()
+
+
+class DollyDataset(Dataset):
+    """Each line of the annotation file is a json object with the following fields:
+
+    {
+        "instruction": "What is a dispersive prism?",
+        "context": "In optics, a dispersive prism is an optical prism that is used to disperse light, that is, to separate light into its spectral components (the colors of the rainbow). Different wavelengths (colors) of light will be deflected by the prism at different angles.[1] This is a result of the prism material's index of refraction varying with wavelength (dispersion). Generally, longer wavelengths (red) undergo a smaller deviation than shorter wavelengths (blue). The dispersion of white light into colors by a prism led Sir Isaac Newton to conclude that white light consisted of a mixture of different colors.",
+        "response": "A dispersive prism is an optical prism that disperses the light's different wavelengths at different angles. When white light is shined through a dispersive prism it will separate into the different colors of the rainbow.",
+        "category": "summarization"
+    }
+
+    """
+
+    def __init__(self, tokenizer, ann_path: str, add_eos=True, ignore_instruction=True, **kwargs):
+        """
+        ann_path (string): directory to store the annotation file
+        """
+        assert tokenizer.add_eos_token is False, "tokenizer should not add eos token by default"
+        self.tokenizer: LlamaTokenizer = tokenizer
+
+        self.annotation = []
+        self.prompter = LMPrompter()
+        self.add_eos = add_eos
+        self.ignore_instruction = ignore_instruction
+        self.load_annotation(ann_path)
+
+    def load_annotation(self, ann_path):
+        self.annotation = []
+        for line in open(ann_path, "r").readlines():
+            self.annotation.append(json.loads(line))
+
+    def __len__(self):
+        return len(self.annotation)
+
+    def process_text(self, ann):
+        instruction = ann["instruction"]
+        context = ann["context"]
+        response = ann["response"]
+        instruction = self.prompter(instruction=instruction, input=context)
+        return dict(instruction=instruction, answer=response)
+
+    def tokenize(self, text):
+        res = self.tokenizer(
+            text["instruction"] + text["answer"],
+            return_tensors=None,
+            padding="do_not_pad",
+            truncation=True,
+            max_length=512,
+        )
+
+        # manually add eos token
+        if res["input_ids"][-1] != self.tokenizer.eos_token_id and len(res["input_ids"]) < 512 and self.add_eos:
+            res["input_ids"].append(self.tokenizer.eos_token_id)
+            res["attention_mask"].append(1)
+        labels = copy.deepcopy(res["input_ids"])
+        # ignore instruction_token
+        if self.ignore_instruction:
+            instruction_token = self.tokenizer(
+                text["instruction"], return_tensors=None, padding="do_not_pad", truncation=True, max_length=512
+            )
+            labels = [-100] * len(instruction_token["input_ids"]) + labels[len(instruction_token["input_ids"]) :]
+
+        res.update(labels=labels)
+        return res
+
+    def __getitem__(self, index):
+        ann = self.annotation[index]
+        text = self.process_text(ann)
+        res = self.tokenize(text)
+        res.update(text)
+        return res
+
+    def collater(self, samples):
+        question_list, answer_list, input_id_list, attention_mask_list, labels_list = [], [], [], [], []
+
+        for sample in samples:
+            question_list.append(sample["instruction"])
+            answer_list.append(sample["answer"])
+            input_id_list.append(sample["input_ids"])
+            attention_mask_list.append(sample["attention_mask"])
+            labels_list.append(sample["labels"])
+
+        # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
+        # same length to return tensors.
+        max_label_length = max(len(l) for l in labels_list)
+        padding_side = self.tokenizer.padding_side
+        padded_labels = []
+        for l in labels_list:
+            remainder = [-100] * (max_label_length - len(l))
+            if isinstance(l, list):
+                l = l + remainder if padding_side == "right" else remainder + l
+            elif padding_side == "right":
+                l = np.concatenate([l, remainder]).astype(np.int64)
+            else:
+                l = np.concatenate([remainder, l]).astype(np.int64)
+            padded_labels.append(l)
+
+        padded_samples = self.tokenizer.pad(
+            {"input_ids": input_id_list, "attention_mask": attention_mask_list, "labels": padded_labels},
+            return_tensors="pt",
+            padding="longest",
+        )
+
+        labels = padded_samples["labels"]
+        labels[labels == self.tokenizer.pad_token_id] = -100
+        labels[:, 0] = -100
+        return {
+            "input_ids": padded_samples["input_ids"],
+            "attention_mask": padded_samples["attention_mask"],
+            "labels": labels,
+            "instruction": question_list,
+            "answer": answer_list,
+        }
+
+
+def build_dolly_dataset(
+    tokenizer,
+    ann_path="data/dolly/databricks-dolly-15k.jsonl",
+    **kwargs,
+):
+    return DollyDataset(
+        tokenizer=tokenizer,
+        ann_path=ann_path,
+        **kwargs,
+    )

mmgpt/datasets/gqa_dataset.py

@@ -0,0 +1,83 @@
+import json
+import os
+import random
+from collections import defaultdict
+
+from PIL import Image
+
+from .vqa_dataset import VQADataset
+
+
+class GQADataset(VQADataset):
+    """Visual Reasoning Dataset."""
+
+    def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, **kwargs):
+        super().__init__(tokenizer, vis_processor, vis_root, ann_paths=[], **kwargs)
+
+        self.annotation = self.load_annotations(ann_paths)
+        if self.sample_image:
+            print("randomly sample one annotation for each image")
+            self.annotation = self.parse_annotation(self.annotation)
+        self._add_instance_ids()
+        self.answer_prob = 1.0
+
+    @staticmethod
+    def load_annotations(ann_paths):
+        annotation = []
+        for ann_path in ann_paths:
+            ann = json.load(open(ann_path, "r"))
+            for k, v in ann.items():
+                v["question_id"] = k
+                annotation.append(v)
+        return annotation
+
+    def parse_annotation(self, annotation):
+        image_list = defaultdict(list)
+        for ann in annotation:
+            image_list[ann["imageId"]].append(ann)
+        annotation = []
+        for ann_list in image_list.values():
+            annotation.append(random.choice(ann_list))
+        return annotation
+
+    def process_text(self, ann):
+        question = ann["question"]
+
+        answer = ann["answer"]
+        full_answer = ann["fullAnswer"]
+
+        # TODO: check which one is better
+        # Random select answer or full_answer
+        if random.random() < self.answer_prob:
+            select_answer = full_answer
+        else:
+            select_answer = answer
+
+        instruction = self.prompter(question)
+        return dict(instruction=instruction, answer=select_answer)
+
+    def process_image(self, ann):
+        image_path = os.path.join(self.vis_root, ann["imageId"] + ".jpg")
+        image = Image.open(image_path).convert("RGB")
+
+        image = self.vis_processor(image)
+        return image
+
+
+def build_gqa_dataset(
+    tokenizer,
+    vis_processor,
+    vis_root="data/gqa/images",
+    ann_paths=[
+        "data/gqa/questions/train_all_questions/train_all_questions_0.json",
+        "data/gqa/questions/val_all_questions.json",
+    ],
+    sample_image=False,
+):
+    return GQADataset(
+        tokenizer=tokenizer,
+        vis_processor=vis_processor,
+        vis_root=vis_root,
+        ann_paths=ann_paths,
+        sample_image=sample_image,
+    )

mmgpt/datasets/llava_dataset.py

@@ -0,0 +1,18 @@
+from .vqa_dataset import VQADataset
+
+
+class LlavaDataset(VQADataset):
+    def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, **kwargs):
+        super().__init__(tokenizer, vis_processor, vis_root, ann_paths, **kwargs)
+
+    def _add_instance_ids(self, key="id"):
+        for idx, ann in enumerate(self.annotation):
+            ann[key] = str(idx)
+
+    def process_text(self, ann):
+        question = ann["conversations"][0]["value"]
+        # remove '<image>' tag and '\n'
+        question = question.replace("<image>", "").replace("\n", "")
+        answer = ann["conversations"][1]["value"]
+        instruction = self.prompter(question)
+        return dict(instruction=instruction, answer=answer)

mmgpt/datasets/nlvr_dataset.py

@@ -0,0 +1,212 @@
+import copy
+import json
+import os
+import random
+from collections import defaultdict
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from .vqa_dataset import VQADataset
+
+QUESTIONS = [
+    "Is this true?",
+    "Is this right?",
+    "Can you confirm this information?" "Do you agree with this statement?",
+    "Does this align with your understanding?",
+    "How do you interpret this information?",
+    "Does this align with your understanding?",
+    "Can you confirm this?",
+    "Is this statement correct?",
+    "Could you verify this information?",
+    "Do you agree with this?",
+    "Is this accurate?",
+    "Can you validate this claim?",
+    "Are these details valid?",
+    "Is this factually correct?",
+    "Is the following information correct?",
+    "Could you please verify this fact?",
+    "Do you agree with this assertion?",
+    "Are these details accurate?",
+    "Does this claim hold true?",
+]
+
+
+class NLVRv1Dataset(VQADataset):
+    """Visual Reasoning Dataset."""
+
+    def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, **kwargs):
+        super().__init__(tokenizer, vis_processor, vis_root, ann_paths=[], **kwargs)
+
+        self.annotation = self.load_annotations(ann_paths)
+        if self.sample_image:
+            print("randomly sample one annotation for each image")
+            self.annotation = self.parse_annotation(self.annotation)
+        self._add_instance_ids()
+
+    @staticmethod
+    def load_annotations(ann_paths):
+        annotation = []
+        for ann_path in ann_paths:
+            if "train.json" in ann_path:
+                split = "train"
+            elif "dev.json" in ann_path:
+                split = "dev"
+            elif "test.json" in ann_path:
+                split = "test"
+            else:
+                raise ValueError(f"Unknown split for {ann_path}")
+
+            with open(ann_path, "r") as f:
+                for line in f.readlines():
+                    line = line.strip()
+                    if len(line) != 0:
+                        ann = json.loads(line)
+                        ann["split"] = split
+                        annotation.append(ann)
+
+        return annotation
+
+    def parse_annotation(self, annotation):
+        image_list = defaultdict(list)
+        for ann in annotation:
+            img_key = f"{ann['split']}-{ann['identifier']}"
+            image_list[img_key].append(ann)
+        annotation = []
+        for ann_list in image_list.values():
+            annotation.append(random.choice(ann_list))
+        return annotation
+
+    def process_text(self, ann):
+        question = ann["sentence"] + " " + random.choice(QUESTIONS)
+        true_answer = ann["label"]
+
+        if random.random() < self.option_prob:
+            instruction = self.prompter(question, ["true", "false"])
+        else:
+            instruction = self.prompter(question)
+
+        return dict(instruction=instruction, answer=true_answer)
+
+    def process_image(self, ann):
+        # each question have 6 images, we can random select one of them.
+        # TODO: check whether using all 6 images?
+        random_id = random.randint(0, 5)
+        image_name = f"{ann['split']}-{ann['identifier']}-{random_id}.png"
+        image_path = os.path.join(self.vis_root, ann["split"], "images", ann["directory"], image_name)
+        image = Image.open(image_path).convert("RGB")
+
+        image = self.vis_processor(image)
+        return image
+
+
+class NLVRv2Dataset(VQADataset):
+    """Visual Reasoning Dataset."""
+
+    def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, **kwargs):
+        super().__init__(tokenizer, vis_processor, vis_root, ann_paths, **kwargs)
+        self.flip_prob = 0.5
+
+    def parse_annotation(self, annotation):
+        image_list = defaultdict(list)
+        for ann in annotation:
+            image_list[ann["images"][0]].append(ann)
+        # image_name_list = list(image_list.keys())
+        annotation = []
+        for ann_list in image_list.values():
+            annotation.append(random.choice(ann_list))
+        return annotation
+
+    def process_text(self, ann):
+        question = ann["sentence"] + " " + random.choice(QUESTIONS)
+        true_answer = ann["label"]
+
+        if random.random() < self.option_prob:
+            instruction = self.prompter(question, ["true", "false"])
+        else:
+            instruction = self.prompter(question)
+
+        return dict(instruction=instruction, answer=true_answer)
+
+    def process_image(self, ann):
+        image_0_path = os.path.join(self.vis_root, ann["images"][0])
+        image_1_path = os.path.join(self.vis_root, ann["images"][1])
+
+        image_0 = Image.open(image_0_path).convert("RGB")
+        image_1 = Image.open(image_1_path).convert("RGB")
+        image_0 = self.vis_processor(image_0)
+        image_1 = self.vis_processor(image_1)
+        return image_0, image_1
+
+    @staticmethod
+    def _flip(samples):
+        sentence = samples["sentence"]
+        image0, image1 = samples["image0"], samples["image1"]
+
+        if "left" not in sentence and "right" not in sentence:
+            if random.random() < 0.5:
+                image0, image1 = image1, image0
+        else:
+            if random.random() < 0.5:
+                sentence = sentence.replace("left", "[TEMP_TOKEN]")
+                sentence = sentence.replace("right", "left")
+                sentence = sentence.replace("[TEMP_TOKEN]", "right")
+
+                image0, image1 = image1, image0
+
+        samples["sentence"] = sentence
+        samples["image0"] = image0
+        samples["image1"] = image1
+
+        return samples
+
+    def __getitem__(self, index):
+        ann = copy.deepcopy(self.annotation[index])
+        image_0, image_1 = self.process_image(ann)
+        if random.random() < self.flip_prob:
+            samples = self._flip({"sentence": ann["sentence"], "image0": image_0, "image1": image_1})
+            image_0, image_1 = samples["image0"], samples["image1"]
+            ann["sentence"] = samples["sentence"]
+        # concat
+        # TODO: https://github.com/salesforce/LAVIS/blob/main/lavis/models/blip_models/blip_nlvr.py
+        # model logic need update if using nlvr2
+        image = torch.cat([image_0, image_1], dim=2)
+        image = F.interpolate(image[None, ...], size=(image_0.shape[1], image_0.shape[2]))[0]
+        text = self.process_text(ann)
+        res = self.tokenize(text)
+        res.update(image=image)
+        res.update(text)
+        return res
+
+
+def build_nlvrv1_dataset(
+    tokenizer,
+    vis_processor,
+    vis_root="data/nlvr",
+    ann_paths=["data/nlvr//train/train.json"],
+    sample_image=False,
+):
+    return NLVRv1Dataset(
+        tokenizer=tokenizer,
+        vis_processor=vis_processor,
+        vis_root=vis_root,
+        ann_paths=ann_paths,
+        sample_image=sample_image,
+    )
+
+
+def build_nlvrv2_dataset(
+    tokenizer,
+    vis_processor,
+    vis_root="data/nlvr2",
+    ann_paths=["data/nlvr2/annotations/nlvr_train.json"],
+    sample_image=False,
+):
+    return NLVRv2Dataset(
+        tokenizer=tokenizer,
+        vis_processor=vis_processor,
+        vis_root=vis_root,
+        ann_paths=ann_paths,
+        sample_image=sample_image,
+    )

mmgpt/datasets/ocr_vqa_dataset.py

@@ -0,0 +1,23 @@
+import os
+import random
+
+from PIL import Image
+
+from .vqa_dataset import VQADataset
+
+
+class OCRVQADataset(VQADataset):
+    def process_image(self, ann):
+        image_path = os.path.join(self.vis_root, ann["filename"])
+        image = Image.open(image_path).convert("RGB")
+
+        image = self.vis_processor(image)
+        return image
+
+    def process_text(self, ann):
+        index = random.choice(list(range(len(ann["questions"]))))
+        question = ann["questions"][index]
+        answer = ann["answers"][index]
+
+        instruction = self.prompter(question)
+        return dict(instruction=instruction, answer=answer)

mmgpt/datasets/samplers/__init__.py

@@ -0,0 +1 @@
+from .infinite_sampler import InfiniteSampler

mmgpt/datasets/samplers/infinite_sampler.py

@@ -0,0 +1,30 @@
+import itertools
+
+import torch
+from torch.utils.data.sampler import Sampler
+
+from mmgpt.train.distributed import world_info_from_env
+
+
+class InfiniteSampler(Sampler):
+    def __init__(self, dataset: int, shuffle: bool = True, seed: int = 0):
+        self._size = len(dataset)
+        self._shuffle = shuffle
+        self._seed = int(seed)
+        _, rank, world_size = world_info_from_env()
+
+        self._rank = rank
+        self._world_size = world_size
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)
+
+    def _infinite_indices(self):
+        g = torch.Generator()
+        g.manual_seed(self._seed)
+        while True:
+            if self._shuffle:
+                yield from torch.randperm(self._size, generator=g).tolist()
+            else:
+                yield from torch.arange(self._size).tolist()

mmgpt/datasets/snli_ve_datasets.py

@@ -0,0 +1,82 @@
+import json
+import os
+import random
+from collections import defaultdict
+
+from PIL import Image
+
+from .vqa_dataset import VQADataset
+
+QUESTIONS = [
+    "What do you think of the above sentence?",
+    "Can you confirm this statement?",
+    "How do you interpret the given information?",
+    "What is your opinion on this matter?",
+    "Could you provide your perspective on this statement?",
+    "How would you respond to the provided claim?",
+    "What are your thoughts regarding the mentioned subject?",
+    "Can you elaborate on this idea in English?",
+    "Do you have any insights or feedback on this topic?",
+    "What's your take on the given statement?",
+    "What is your perspective on the given statement?",
+    "How would you interpret this remark?",
+    "Could you please provide your opinion on this?",
+    "Can you share your understanding of the above point?",
+    "Would you mind elaborating on this topic?",
+    "What are your views about the given statement?",
+    "How do you feel about the presented information?",
+    "Could you provide your perspective on this?",
+    "What is your opinion regarding this statement?",
+    "Can you share your thoughts about the mentioned claim?",
+    "How would you interpret the above comment?",
+    "Would you mind sharing your insights on this issue?",
+]
+
+
+class SNLIVEDataset(VQADataset):
+    """Visual Reasoning Dataset."""
+
+    def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, **kwargs):
+        super().__init__(tokenizer, vis_processor, vis_root, ann_paths=[], **kwargs)
+
+        self.annotation = self.load_annotations(ann_paths)
+        if self.sample_image:
+            print("randomly sample one annotation for each image")
+            self.annotation = self.parse_annotation(self.annotation)
+        self._add_instance_ids()
+
+    @staticmethod
+    def load_annotations(ann_paths):
+        annotation = []
+        for ann_path in ann_paths:
+            with open(ann_path, "r") as f:
+                for line in f.readlines():
+                    line = line.strip()
+                    if len(line) != 0:
+                        ann = json.loads(line)
+                        annotation.append(ann)
+        return annotation
+
+    def parse_annotation(self, annotation):
+        image_list = defaultdict(list)
+        for ann in annotation:
+            image_list[ann["Flickr30K_ID"]].append(ann)
+        annotation = []
+        for ann_list in image_list.values():
+            annotation.append(random.choice(ann_list))
+        return annotation
+
+    def process_text(self, ann):
+        question = ann["sentence2"] + " " + random.choice(QUESTIONS)
+        answer = ann["gold_label"]
+        if random.random() < self.option_prob:
+            instruction = self.prompter(question, ["entailment", "neutral", "contradiction"])
+        else:
+            instruction = self.prompter(question)
+        return dict(instruction=instruction, answer=answer)
+
+    def process_image(self, ann):
+        image_path = os.path.join(self.vis_root, ann["Flickr30K_ID"] + ".jpg")
+        image = Image.open(image_path).convert("RGB")
+        image = self.vis_processor(image)
+        return image

mmgpt/datasets/text_ocr_dataset.py

@@ -0,0 +1,64 @@
+import json
+import os
+import random
+
+import numpy as np
+from PIL import Image
+from transformers import LlamaTokenizer
+
+from .vqa_dataset import VQADataset, VQAPrompter
+
+
+class TextOCRDataset(VQADataset):
+    def __init__(
+        self, tokenizer, vis_processor=None, vis_root=None, ann_paths=[], add_eos=True, ignore_instruction=True
+    ):
+        """
+        vis_root (string): Root directory of images (e.g. coco/images/)
+        ann_root (string): directory to store the annotation file
+        """
+        assert tokenizer.add_eos_token is False, "tokenizer should not add eos token by default"
+        self.tokenizer: LlamaTokenizer = tokenizer
+        self.vis_root = vis_root
+
+        self.annotation = []
+        for ann_path in ann_paths:
+            self.annotation.extend(json.load(open(ann_path, "r"))["data"])
+
+        self.vis_processor = vis_processor
+
+        self._add_instance_ids()
+        self.option_prob = 0.5
+        self.prompter = VQAPrompter()
+        self.add_eos = add_eos
+        self.ignore_instruction = ignore_instruction
+
+    def process_image(self, ann):
+        image_path = os.path.join(self.vis_root, ann["image_id"] + ".jpg")
+        image = Image.open(image_path).convert("RGB")
+
+        image = self.vis_processor(image)
+        return image
+
+    def process_text(self, ann):
+        question = ann["question"]
+
+        answer_weight = {}
+        for answer in ann["answers"]:
+            if answer in answer_weight.keys():
+                answer_weight[answer] += 1 / len(ann["answers"])
+            else:
+                answer_weight[answer] = 1 / len(ann["answers"])
+
+        answers = list(answer_weight.keys())
+        weights = list(answer_weight.values())
+
+        # create instruction
+        true_answer = answers[np.argmax(weights)]
+        is_option = random.random() < self.option_prob and len(answers) > 1
+        if is_option:
+            instruction = self.prompter(question, answers)
+        else:
+            instruction = self.prompter(question)
+
+        return dict(instruction=instruction, answer=true_answer)

mmgpt/datasets/vqa_dataset.py

@@ -0,0 +1,227 @@
+"""
+ Copyright (c) 2022, salesforce.com, inc.
+ All rights reserved.
+ SPDX-License-Identifier: BSD-3-Clause
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+"""
+
+import copy
+import json
+import os
+import random
+from collections import defaultdict
+from typing import Iterable
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import ConcatDataset, Dataset
+from torch.utils.data.dataloader import default_collate
+from transformers import LlamaTokenizer
+
+TEMPLATE = {
+    "description": "Template used by Alpaca-LoRA.",
+    # "prompt_choice": "Below is a multiple choice question about an image, along with answer options. Please choose the correct answer from these options.\n\n### Image:\n{image}\n\n### Question:\n{question}\n\n### Input:\n{options}\n\n### Answer:\n",
+    # "prompt_qa": "Below is a question about an image. Write a response to answer the question.\n\n### Image:\n{image}\n\n### Question:\n{question}\n\n### Answer:\n",
+    "prompt_choice": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Image:\n{image}\n\n### Instruction:\n{question}\n\n### Input:\n{options}\n\n### Response:\n",
+    "prompt_qa": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Image:\n{image}\n\n### Instruction:\n{question}\n\n### Response:\n",
+    "response_split": "### Response:",
+}
+
+
+class VQAPrompter:
+    def __call__(self, question, options=None):
+        if options:
+            options = ", ".join(options)
+            res = TEMPLATE["prompt_choice"].format(image="<image>", question=question, options=options)
+        else:
+            res = TEMPLATE["prompt_qa"].format(image="<image>", question=question)
+        return res
+
+    def get_response(self, output: str) -> str:
+        return output.split(TEMPLATE["response_split"])[-1].strip()
+
+
+class VQADataset(Dataset):
+    def __init__(
+        self,
+        tokenizer,
+        vis_processor=None,
+        vis_root=None,
+        ann_paths=[],
+        add_eos=True,
+        ignore_instruction=True,
+        sample_image=False,
+    ):
+        """
+        vis_root (string): Root directory of images (e.g. coco/images/)
+        ann_root (string): directory to store the annotation file
+        """
+        assert tokenizer.add_eos_token is False, "tokenizer should not add eos token by default"
+        self.tokenizer: LlamaTokenizer = tokenizer
+        self.vis_root = vis_root
+
+        self.annotation = []
+        for ann_path in ann_paths:
+            self.annotation.extend(json.load(open(ann_path, "r")))
+
+        self.sample_image = sample_image
+        if self.sample_image:
+            print("randomly sample one annotation for each image")
+            self.annotation = self.parse_annotation(self.annotation)
+
+        self.vis_processor = vis_processor
+
+        self._add_instance_ids()
+        self.option_prob = 0.5
+        self.prompter = VQAPrompter()
+        self.add_eos = add_eos
+        self.ignore_instruction = ignore_instruction
+
+    def parse_annotation(self, annotation):
+        image_list = defaultdict(list)
+        for ann in annotation:
+            image_list[ann["image"]].append(ann)
+        # image_name_list = list(image_list.keys())
+        annotation = []
+        for ann_list in image_list.values():
+            annotation.append(random.choice(ann_list))
+        return annotation
+
+    def __len__(self):
+        return len(self.annotation)
+
+    def _add_instance_ids(self, key="instance_id"):
+        for idx, ann in enumerate(self.annotation):
+            ann[key] = str(idx)
+
+    def process_image(self, ann):
+        image_path = os.path.join(self.vis_root, ann["image"])
+        image = Image.open(image_path).convert("RGB")
+
+        image = self.vis_processor(image)
+        return image
+
+    def process_text(self, ann):
+        question = ann["question"]
+
+        answer_weight = {}
+        for answer in ann["answer"]:
+            if answer in answer_weight.keys():
+                answer_weight[answer] += 1 / len(ann["answer"])
+            else:
+                answer_weight[answer] = 1 / len(ann["answer"])
+
+        answers = list(answer_weight.keys())
+        weights = list(answer_weight.values())
+
+        # create instruction
+        true_answer = answers[np.argmax(weights)]
+        is_option = random.random() < self.option_prob and len(answers) > 1
+        if is_option:
+            instruction = self.prompter(question, answers)
+        else:
+            instruction = self.prompter(question)
+
+        return dict(instruction=instruction, answer=true_answer)
+
+    def tokenize(self, text):
+        res = self.tokenizer(
+            text["instruction"] + text["answer"],
+            return_tensors=None,
+            padding="do_not_pad",
+            truncation=True,
+            max_length=512,
+        )
+
+        # manually add eos token
+        if res["input_ids"][-1] != self.tokenizer.eos_token_id and len(res["input_ids"]) < 512 and self.add_eos:
+            res["input_ids"].append(self.tokenizer.eos_token_id)
+            res["attention_mask"].append(1)
+        labels = copy.deepcopy(res["input_ids"])
+        # ignore instruction_token
+        if self.ignore_instruction:
+            instruction_token = self.tokenizer(
+                text["instruction"], return_tensors=None, padding="do_not_pad", truncation=True, max_length=512
+            )
+            labels = [-100] * len(instruction_token["input_ids"]) + labels[len(instruction_token["input_ids"]) :]
+
+        res.update(labels=labels)
+        return res
+
+    def __getitem__(self, index):
+        ann = self.annotation[index]
+        image = self.process_image(ann)
+        text = self.process_text(ann)
+        res = self.tokenize(text)
+        res.update(image=image)
+        res.update(text)
+        return res
+
+    def collater(self, samples):
+        image_list, question_list, answer_list, input_id_list, attention_mask_list, labels_list = [], [], [], [], [], []
+
+        for sample in samples:
+            image_list.append(sample["image"])
+            question_list.append(sample["instruction"])
+            answer_list.append(sample["answer"])
+            input_id_list.append(sample["input_ids"])
+            attention_mask_list.append(sample["attention_mask"])
+            labels_list.append(sample["labels"])
+
+        # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
+        # same length to return tensors.
+        max_label_length = max(len(l) for l in labels_list)
+        padding_side = self.tokenizer.padding_side
+        padded_labels = []
+        for l in labels_list:
+            remainder = [-100] * (max_label_length - len(l))
+            if isinstance(l, list):
+                l = l + remainder if padding_side == "right" else remainder + l
+            elif padding_side == "right":
+                l = np.concatenate([l, remainder]).astype(np.int64)
+            else:
+                l = np.concatenate([remainder, l]).astype(np.int64)
+            padded_labels.append(l)
+
+        padded_samples = self.tokenizer.pad(
+            {"input_ids": input_id_list, "attention_mask": attention_mask_list, "labels": padded_labels},
+            return_tensors="pt",
+            padding="longest",
+        )
+
+        labels = padded_samples["labels"]
+        media_token_id = self.tokenizer("<image>", add_special_tokens=False)["input_ids"][-1]
+        labels[labels == self.tokenizer.pad_token_id] = -100
+        labels[:, 0] = -100
+        labels[labels == media_token_id] = -100
+        return {
+            "image": torch.stack(image_list, dim=0),
+            "input_ids": padded_samples["input_ids"],
+            "attention_mask": padded_samples["attention_mask"],
+            "labels": labels,
+            "instruction": question_list,
+            "answer": answer_list,
+        }
+
+
+class ConcatDataset(ConcatDataset):
+    def __init__(self, datasets: Iterable[Dataset]) -> None:
+        super().__init__(datasets)
+
+    def collater(self, samples):
+        # TODO For now only supports datasets with same underlying collater implementations
+
+        all_keys = set()
+        for s in samples:
+            all_keys.update(s)
+
+        shared_keys = all_keys
+        for s in samples:
+            shared_keys = shared_keys & set(s.keys())
+
+        samples_shared_keys = []
+        for s in samples:
+            samples_shared_keys.append({k: s[k] for k in s.keys() if k in shared_keys})
+
+        return self.datasets[0].collater(samples_shared_keys)

mmgpt/models/builder.py

@@ -0,0 +1,74 @@
+from .open_flamingo import create_model_and_transforms as create_open_flamingo_model_and_transforms
+import torch.nn as nn
+from transformers import LlamaTokenizer, LlamaForCausalLM
+
+def create_model_and_transforms(
+    model_name: str,
+    clip_vision_encoder_path: str,
+    clip_vision_encoder_pretrained: str,
+    lang_encoder_path: str,
+    tokenizer_path: str,
+    tuning_config,
+    pretrained_model_path,
+    **kwargs,
+):
+    if model_name == "open_flamingo":
+        return create_open_flamingo_model_and_transforms(
+            clip_vision_encoder_path=clip_vision_encoder_path,
+            clip_vision_encoder_pretrained=clip_vision_encoder_pretrained,
+            lang_encoder_path=lang_encoder_path,
+            tokenizer_path=tokenizer_path,
+            tuning_config=tuning_config,
+            pretrained_model_path=pretrained_model_path,
+            **kwargs,
+        )
+    # TODO: support BLIP2
+    else:
+        raise ValueError(f"Unknown model name: {model_name}")
+
+# only for debugging
+def create_toy_model_and_transforms(
+    model_name: str,
+    clip_vision_encoder_path: str,
+    clip_vision_encoder_pretrained: str,
+    lang_encoder_path: str,
+    tokenizer_path: str,
+    tuning_config,
+    pretrained_model_path,
+    **kwargs,
+):
+    print("init toy vision encoder")
+    import torchvision
+
+    image_processor = torchvision.transforms.Compose(
+        [
+            torchvision.transforms.Resize((224, 224)),
+            torchvision.transforms.ToTensor(),
+        ]
+    )
+    print("init tokenizer")
+    text_tokenizer = LlamaTokenizer.from_pretrained(tokenizer_path)
+    # add Flamingo special tokens to the tokenizer
+    text_tokenizer.add_special_tokens({"additional_special_tokens": ["<|endofchunk|>", "<image>"]})
+    if text_tokenizer.pad_token is None:
+        # Issue: GPT models don't have a pad token, which we use to
+        # modify labels for the loss.
+        text_tokenizer.add_special_tokens({"pad_token": "<PAD>"})
+
+    class ToyModel(nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+            self.input_embeddings = nn.Embedding(38000, 512)
+            self.layer = nn.Linear(512, 512)
+            self.config = {"hidden_size": 512}
+
+        def forward(self, lang_x, **kwargs):
+            x = self.input_embeddings(lang_x)
+            x = self.layer(x)
+            loss = x.sum()
+
+            return (loss,)
+
+    model = ToyModel()
+
+    return model, image_processor, text_tokenizer

mmgpt/models/open_flamingo/__init__.py

@@ -0,0 +1,3 @@
+from .builder import create_model_and_transforms
+from .flamingo import Flamingo
+from .flamingo_lm import FlamingoLMMixin

mmgpt/models/open_flamingo/builder.py

@@ -0,0 +1,142 @@
+"""Modified from https://github.com/mlfoundations/open_flamingo"""
+import open_clip
+import torch
+import torch.nn as nn
+from bigmodelvis import Visualization
+from peft import LoraConfig, get_peft_model
+from transformers import LlamaForCausalLM, LlamaTokenizer
+
+from .flamingo import Flamingo
+from .flamingo_lm import FlamingoLMMixin
+from .utils import extend_instance
+
+
+def create_model_and_transforms(
+    clip_vision_encoder_path: str,
+    clip_vision_encoder_pretrained: str,
+    lang_encoder_path: str,
+    tokenizer_path: str,
+    decoder_layers_attr_name: str = None,
+    pretrained_model_path: str = None,
+    tuning_config=None,
+    **flamingo_kwargs,
+):
+    """
+    Initialize a Flamingo model from a pretrained vision encoder and language encoder.
+    Appends special tokens to the tokenizer and freezes backbones.
+
+    Args:
+        clip_vision_encoder_path (str): path to pretrained clip model (e.g. "ViT-B-32")
+        clip_vision_encoder_pretrained (str): name of pretraining dataset for clip model (e.g. "laion2b_s32b_b79k")
+        lang_encoder_path (str): path to pretrained language encoder
+        tokenizer_path (str): path to pretrained tokenizer
+        decoder_layers_attr_name (str, optional): name of the decoder layers attribute. Defaults to None.
+    Returns:
+        Flamingo: Flamingo model from pretrained vision and language encoders
+        Image processor: Pipeline to preprocess input images
+        Tokenizer: A tokenizer for the language model
+    """
+    print("init clip vision encoder")
+    vision_encoder, _, image_processor = open_clip.create_model_and_transforms(
+        clip_vision_encoder_path, pretrained=clip_vision_encoder_pretrained
+    )
+    # set the vision encoder to output the visual features
+    vision_encoder.visual.output_tokens = True
+    print("init tokenizer")
+    text_tokenizer = LlamaTokenizer.from_pretrained(tokenizer_path)
+    # add Flamingo special tokens to the tokenizer
+    text_tokenizer.add_special_tokens({"additional_special_tokens": ["<|endofchunk|>", "<image>"]})
+    if text_tokenizer.pad_token is None:
+        # Issue: GPT models don't have a pad token, which we use to
+        # modify labels for the loss.
+        text_tokenizer.add_special_tokens({"pad_token": "<PAD>"})
+    text_tokenizer.bos_token_id = 1
+    text_tokenizer.eos_token_id = 2
+
+    print("init llama")
+    lang_encoder = LlamaForCausalLM.from_pretrained(lang_encoder_path)
+    extend_instance(lang_encoder, FlamingoLMMixin)
+
+    if decoder_layers_attr_name is None:
+        decoder_layers_attr_name = _infer_decoder_layers_attr_name(lang_encoder)
+    lang_encoder.set_decoder_layers_attr_name(decoder_layers_attr_name)
+    lang_encoder.resize_token_embeddings(len(text_tokenizer))
+
+    model = Flamingo(
+        vision_encoder,
+        lang_encoder,
+        text_tokenizer.encode("<|endofchunk|>")[-1],
+        text_tokenizer.encode("<image>")[-1],
+        vis_dim=open_clip.get_model_config(clip_vision_encoder_path)["vision_cfg"]["width"],
+        cross_attn_every_n_layers=4,
+        **flamingo_kwargs,
+    )
+
+    if pretrained_model_path is not None:
+        print(f"loading pretrained model from {pretrained_model_path}")
+        model.load_state_dict(torch.load(pretrained_model_path), strict=False)
+
+    # Freeze all parameters
+    model.requires_grad_(False)
+    assert sum(p.numel() for p in model.parameters() if p.requires_grad) == 0
+
+    if tuning_config is not None:
+        model = prepare_model_for_tuning(model, tuning_config)
+    else:
+        raise ValueError("tuning_config must be provided")
+
+    print(
+        f"Flamingo model initialized with {sum(p.numel() for p in model.parameters() if p.requires_grad)} trainable parameters"
+    )
+
+    return model, image_processor, text_tokenizer
+
+
+def _infer_decoder_layers_attr_name(model):
+    for k in __KNOWN_DECODER_LAYERS_ATTR_NAMES:
+        if k.lower() in model.__class__.__name__.lower():
+            return __KNOWN_DECODER_LAYERS_ATTR_NAMES[k]
+
+    raise ValueError(
+        f"We require the attribute name for the nn.ModuleList in the decoder storing the transformer block layers. Please supply this string manually."
+    )
+
+
+__KNOWN_DECODER_LAYERS_ATTR_NAMES = {
+    "opt": "model.decoder.layers",
+    "gptneo": "transformer.h",
+    "gptj": "transformer.h",
+    "gpt-j": "transformer.h",
+    "pythia": "gpt_neox.layers",
+    "llama": "model.layers",
+}
+
+
+def prepare_model_for_tuning(model: nn.Module, config):
+    if config.lora:
+        lora_config = LoraConfig(
+            r=config.lora_r,
+            lora_alpha=config.lora_alpha,
+            target_modules=config.lora_target_modules,
+            lora_dropout=config.lora_dropout,
+            bias="none",  # won't use bias currently
+            modules_to_save=[],  # TODO: might be helpful if save partial model
+            task_type="CAUSAL_LM",
+        )
+        model.lang_encoder = get_peft_model(model.lang_encoder, peft_config=lora_config)
+
+    # manually unfreeze modules, we use a `substring` fashion mathcing
+    for name, param in model.named_parameters():
+        if any(substr in name for substr in config.unfrozen):
+            param.requires_grad = True
+
+    if config.vis and is_rank0():
+        Visualization(model).structure_graph()
+    return model
+
+
+# temporary workaround, should use a common utils in the future
+def is_rank0():
+    if not torch.distributed.is_initialized():
+        return True
+    return torch.distributed.get_rank() == 0

mmgpt/models/open_flamingo/flamingo.py

@@ -0,0 +1,208 @@
+"""Modified from https://github.com/mlfoundations/open_flamingo"""
+import torch
+from einops import rearrange
+from torch import nn
+
+from .helpers import PerceiverResampler
+
+
+class Flamingo(nn.Module):
+    def __init__(
+        self,
+        vision_encoder: nn.Module,
+        lang_encoder: nn.Module,
+        eoc_token_id: int,
+        media_token_id: int,
+        vis_dim: int,
+        cross_attn_every_n_layers: int = 1,
+        use_media_placement_augmentation: bool = False,
+    ):
+        """
+        Args:
+            vision_encoder (nn.Module): HF CLIPModel
+            lang_encoder (nn.Module): HF causal language model
+            eoc_token_id (int): Token id for <|endofchunk|>
+            media_token_id (int): Token id for <image>
+            vis_dim (int): Dimension of the visual features.
+                Visual features are projected to match this shape along the last dimension.
+            cross_attn_every_n_layers (int, optional): How often to apply cross attention after transformer layer. Defaults to 1.
+            use_media_placement_augmentation (bool, optional): Whether to randomly assign images to the preceding or following text in training. Defaults to False.
+        """
+        super().__init__()
+        self.eoc_token_id = eoc_token_id
+        self.media_token_id = media_token_id
+        self.use_media_placement_augmentation = use_media_placement_augmentation
+        self.vis_dim = vis_dim
+        self.vision_encoder = vision_encoder
+        self.perceiver = PerceiverResampler(dim=self.vis_dim)
+        self.lang_encoder = lang_encoder
+        self.lang_encoder.init_flamingo(
+            media_token_id=media_token_id,
+            vis_hidden_size=self.vis_dim,
+            cross_attn_every_n_layers=cross_attn_every_n_layers,
+            use_media_placement_augmentation=self.use_media_placement_augmentation,
+        )
+
+    def forward(
+        self,
+        vision_x: torch.Tensor,
+        lang_x: torch.Tensor,
+        attention_mask: torch.Tensor = None,
+        labels: torch.Tensor = None,
+        use_cached_vision_x: bool = False,
+        clear_conditioned_layers: bool = True,
+        past_key_values=None,
+        use_cache: bool = False,
+    ):
+        """
+        Forward pass of Flamingo.
+
+        Args:
+            vision_x (torch.Tensor): Vision input
+                shape (B, T_img, F, C, H, W) with F=1
+            lang_x (torch.Tensor): Language input ids
+                shape (B, T_txt)
+            attention_mask (torch.Tensor, optional): Attention mask. Defaults to None.
+            labels (torch.Tensor, optional): Labels. Defaults to None.
+            clear_conditioned_layers: if True, clear the conditioned layers
+                once the foward pass is completed. Set this to false if the
+                same set of images will be reused in another subsequent
+                forward pass.
+            past_key_values: pre-computed values to pass to language model.
+                See past_key_values documentation in Hugging Face
+                CausalLM models.
+            use_cache: whether to use cached key values. See use_cache
+                documentation in Hugging Face CausalLM models.
+        """
+        if vision_x is None and use_cached_vision_x is False:
+            for layer in self.lang_encoder._get_decoder_layers():
+                layer.condition_only_lang_x(True)
+            output = self.lang_encoder(
+                input_ids=lang_x,
+                attention_mask=attention_mask,
+                labels=labels,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            for layer in self.lang_encoder._get_decoder_layers():
+                layer.condition_only_lang_x(False)
+            return output
+        assert (
+            vision_x is not None
+        ) or use_cached_vision_x, "Must provide either vision_x or use_cached_vision_x to True."
+
+        if use_cached_vision_x:
+            # Case: use cached; vision_x should be cached and other
+            # vision-related inputs should not be provided.
+            assert vision_x is None, "Expect vision_x to be None when use_cached_vision_x is True."
+            assert self.lang_encoder.is_conditioned()
+
+        else:
+            # Case: do not use caching (i.e. this is a standard forward pass);
+            self._encode_vision_x(vision_x=vision_x)
+
+        output = self.lang_encoder(
+            input_ids=lang_x,
+            attention_mask=attention_mask,
+            labels=labels,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+        )
+
+        if clear_conditioned_layers:
+            self.lang_encoder.clear_conditioned_layers()
+
+        return output
+
+    def generate(
+        self,
+        vision_x: torch.Tensor,
+        lang_x: torch.Tensor,
+        attention_mask: torch.Tensor = None,
+        num_beams=1,
+        max_new_tokens=None,
+        temperature=1.0,
+        top_k=0,
+        top_p=1.0,
+        no_repeat_ngram_size=0,
+        prefix_allowed_tokens_fn=None,
+        length_penalty=1.0,
+        num_return_sequences=1,
+        do_sample=False,
+        early_stopping=False,
+    ):
+        """
+        Generate text conditioned on vision and language inputs.
+
+        Args:
+            vision_x (torch.Tensor): Vision input
+                shape (B, T_img, F, C, H, W)
+                images in the same chunk are collated along T_img, and frames are collated along F
+                currently only F=1 is supported (single-frame videos)
+            lang_x (torch.Tensor): Language input
+                shape (B, T_txt)
+            max_length (int, optional): Maximum length of the output. Defaults to None.
+            attention_mask (torch.Tensor, optional): Attention mask. Defaults to None.
+            num_beams (int, optional): Number of beams. Defaults to 1.
+            max_new_tokens (int, optional): Maximum new tokens. Defaults to None.
+            temperature (float, optional): Temperature. Defaults to 1.0.
+            top_k (int, optional): Top k. Defaults to 0.
+            top_p (float, optional): Top p. Defaults to 1.0.
+            no_repeat_ngram_size (int, optional): No repeat ngram size. Defaults to 0.
+            length_penalty (float, optional): Length penalty. Defaults to 1.0.
+            num_return_sequences (int, optional): Number of return sequences. Defaults to 1.
+            do_sample (bool, optional): Do sample. Defaults to False.
+            early_stopping (bool, optional): Early stopping. Defaults to False.
+        Returns:
+            torch.Tensor: lang_x with generated tokens appended to it
+        """
+        if num_beams > 1:
+            vision_x = vision_x.repeat_interleave(num_beams, dim=0)
+
+        self._encode_vision_x(vision_x=vision_x)
+
+        output = self.lang_encoder.generate(
+            lang_x,
+            attention_mask=attention_mask,
+            # eos_token_id=self.eoc_token_id,
+            num_beams=num_beams,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            length_penalty=length_penalty,
+            num_return_sequences=num_return_sequences,
+            do_sample=do_sample,
+            early_stopping=early_stopping,
+        )
+
+        self.lang_encoder.clear_conditioned_layers()
+        return output
+
+    def _encode_vision_x(self, vision_x: torch.Tensor):
+        """
+        Compute media tokens from vision input by passing it through vision encoder and conditioning language model.
+        Args:
+            vision_x (torch.Tensor): Vision input
+                shape (B, T_img, F, C, H, W)
+                Images in the same chunk are collated along T_img, and frames are collated along F
+                Currently only F=1 is supported (single-frame videos)
+
+        rearrange code based on https://github.com/dhansmair/flamingo-mini
+        """
+
+        assert vision_x.ndim == 6, "vision_x should be of shape (b, T_img, F, C, H, W)"
+        b, T, F = vision_x.shape[:3]
+        assert F == 1, "Only single frame supported"
+
+        vision_x = rearrange(vision_x, "b T F c h w -> (b T F) c h w")
+        with torch.no_grad():
+            vision_x = self.vision_encoder.visual(vision_x)[1]
+        vision_x = rearrange(vision_x, "(b T F) v d -> b T F v d", b=b, T=T, F=F)
+
+        vision_x = self.perceiver(vision_x)  # reshapes to (b, T, n, d)
+
+        for layer in self.lang_encoder._get_decoder_layers():
+            layer.condition_vis_x(vision_x)

mmgpt/models/open_flamingo/flamingo_lm.py

@@ -0,0 +1,131 @@
+"""Modified from https://github.com/mlfoundations/open_flamingo"""
+import random
+
+import torch.nn as nn
+
+from .helpers import GatedCrossAttentionBlock
+from .utils import getattr_recursive, setattr_recursive
+
+
+class FlamingoLayer(nn.Module):
+    def __init__(self, gated_cross_attn_layer, decoder_layer):
+        super().__init__()
+        self.gated_cross_attn_layer = gated_cross_attn_layer
+        self.decoder_layer = decoder_layer
+        self.vis_x = None
+        self.media_locations = None
+        self.only_lang_x = False
+
+    def is_conditioned(self) -> bool:
+        """Check whether the layer is conditioned."""
+        return self.vis_x is not None
+
+    # Used this great idea from this implementation of Flamingo (https://github.com/dhansmair/flamingo-mini/)
+    def condition_vis_x(self, vis_x):
+        self.vis_x = vis_x
+
+    def condition_only_lang_x(self, only_lang_x=False):
+        self.only_lang_x = only_lang_x
+
+    def condition_media_locations(self, media_locations):
+        self.media_locations = media_locations
+
+    def condition_attend_previous(self, attend_previous):
+        self.attend_previous = attend_previous
+
+    def forward(
+        self,
+        lang_x,
+        attention_mask=None,
+        **decoder_layer_kwargs,
+    ):
+        if self.gated_cross_attn_layer is None or self.only_lang_x:
+            return self.decoder_layer(lang_x, attention_mask=attention_mask, **decoder_layer_kwargs)
+
+        if self.vis_x is None:
+            raise ValueError("vis_x must be conditioned before forward pass")
+
+        if self.media_locations is None:
+            raise ValueError("media_locations must be conditioned before forward pass")
+
+        lang_x = self.gated_cross_attn_layer(
+            lang_x,
+            self.vis_x,
+            media_locations=self.media_locations,
+            attend_previous=self.attend_previous,
+        )
+        lang_x = self.decoder_layer(lang_x, attention_mask=attention_mask, **decoder_layer_kwargs)
+        return lang_x
+
+
+class FlamingoLMMixin(nn.Module):
+    """
+    Mixin to add cross-attention layers to a language model.
+    """
+
+    def set_decoder_layers_attr_name(self, decoder_layers_attr_name):
+        self.decoder_layers_attr_name = decoder_layers_attr_name
+
+    def _get_decoder_layers(self):
+        return getattr_recursive(self, self.decoder_layers_attr_name)
+
+    def _set_decoder_layers(self, value):
+        setattr_recursive(self, self.decoder_layers_attr_name, value)
+
+    def init_flamingo(
+        self,
+        media_token_id,
+        vis_hidden_size,
+        cross_attn_every_n_layers,
+        use_media_placement_augmentation,
+    ):
+        """
+        Initialize Flamingo by adding a new gated cross attn to the decoder. Store the media token id for computing the media locations.
+        """
+
+        self.gated_cross_attn_layers = nn.ModuleList(
+            [
+                GatedCrossAttentionBlock(dim=self.config.hidden_size, dim_visual=vis_hidden_size)
+                if (layer_idx + 1) % cross_attn_every_n_layers == 0
+                else None
+                for layer_idx, _ in enumerate(self._get_decoder_layers())
+            ]
+        )
+        self._set_decoder_layers(
+            nn.ModuleList(
+                [
+                    FlamingoLayer(gated_cross_attn_layer, decoder_layer)
+                    for gated_cross_attn_layer, decoder_layer in zip(
+                        self.gated_cross_attn_layers, self._get_decoder_layers()
+                    )
+                ]
+            )
+        )
+        self.media_token_id = media_token_id
+        self.use_media_placement_augmentation = use_media_placement_augmentation
+        self.initialized_flamingo = True
+
+    def forward(self, *input, **kwargs):
+        """Condition the Flamingo layers on the media locations before forward()"""
+        if not self.initialized_flamingo:
+            raise ValueError("Flamingo layers are not initialized. Please call `init_flamingo` first.")
+
+        input_ids = kwargs["input_ids"] if "input_ids" in kwargs else input[0]
+        media_locations = input_ids == self.media_token_id
+        attend_previous = (random.random() < 0.5) if self.use_media_placement_augmentation else False
+
+        for layer in self.get_decoder().layers:
+            layer.condition_media_locations(media_locations)
+            layer.condition_attend_previous(attend_previous)
+
+        return super().forward(*input, **kwargs)  # Call the other parent's forward method
+
+    def is_conditioned(self) -> bool:
+        """Check whether all decoder layers are already conditioned."""
+        return all(l.is_conditioned() for l in self._get_decoder_layers())
+
+    def clear_conditioned_layers(self):
+        for layer in self._get_decoder_layers():
+            layer.condition_vis_x(None)
+            layer.condition_media_locations(None)
+            layer.condition_attend_previous(None)

mmgpt/models/open_flamingo/helpers.py

@@ -0,0 +1,263 @@
+"""
+Taken from https://github.com/lucidrains/flamingo-pytorch
+"""
+
+import torch
+from einops import rearrange, repeat
+from einops_exts import rearrange_many
+from torch import einsum, nn
+
+
+def exists(val):
+    return val is not None
+
+
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim, bias=False),
+        nn.GELU(),
+        nn.Linear(inner_dim, dim, bias=False),
+    )
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm_media = nn.LayerNorm(dim)
+        self.norm_latents = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, T, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, T, n2, D)
+        """
+        x = self.norm_media(x)
+        latents = self.norm_latents(latents)
+
+        h = self.heads
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+        q, k, v = rearrange_many((q, k, v), "b t n (h d) -> b h t n d", h=h)
+        q = q * self.scale
+
+        # attention
+        sim = einsum("... i d, ... j d  -> ... i j", q, k)
+        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
+        attn = sim.softmax(dim=-1)
+
+        out = einsum("... i j, ... j d -> ... i d", attn, v)
+        out = rearrange(out, "b h t n d -> b t n (h d)", h=h)
+        return self.to_out(out)
+
+
+class PerceiverResampler(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth=6,
+        dim_head=64,
+        heads=8,
+        num_latents=64,
+        max_num_media=None,
+        max_num_frames=None,
+        ff_mult=4,
+    ):
+        super().__init__()
+        self.latents = nn.Parameter(torch.randn(num_latents, dim))
+        self.frame_embs = nn.Parameter(torch.randn(max_num_frames, dim)) if exists(max_num_frames) else None
+        self.media_time_embs = nn.Parameter(torch.randn(max_num_media, 1, dim)) if exists(max_num_media) else None
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, T, F, v, D)
+        Returns:
+            shape (b, T, n, D) where n is self.num_latents
+        """
+        b, T, F, v = x.shape[:4]
+
+        # frame and media time embeddings
+        if exists(self.frame_embs):
+            frame_embs = repeat(self.frame_embs[:F], "F d -> b T F v d", b=b, T=T, v=v)
+            x = x + frame_embs
+        x = rearrange(x, "b T F v d -> b T (F v) d")  # flatten the frame and spatial dimensions
+        if exists(self.media_time_embs):
+            x = x + self.media_time_embs[:T]
+
+        # blocks
+        latents = repeat(self.latents, "n d -> b T n d", b=b, T=T)
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+        return self.norm(latents)
+
+
+# gated cross attention
+
+
+class MaskedCrossAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_visual,
+        dim_head=64,
+        heads=8,
+        only_attend_immediate_media=True,
+    ):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim_visual, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+        # whether for text to only attend to immediate preceding image, or all previous images
+        self.only_attend_immediate_media = only_attend_immediate_media
+
+    def forward(self, x, media, media_locations=None, attend_previous=True):
+        """
+        Args:
+            x (torch.Tensor): text features
+                shape (B, T_txt, D_txt)
+            media (torch.Tensor): image features
+                shape (B, T_img, n, D_img) where n is the dim of the latents
+            media_locations: boolean mask identifying the media tokens in x
+                shape (B, T_txt)
+            attend_previous: bool
+                If false, ignores immediately preceding image and starts attending when following image
+        """
+        _, T_img, n = media.shape[:3]
+        h = self.heads
+
+        x = self.norm(x)
+
+        q = self.to_q(x)
+        media = rearrange(media, "b t n d -> b (t n) d")
+
+        k, v = self.to_kv(media).chunk(2, dim=-1)
+        q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=h)
+
+        q = q * self.scale
+
+        sim = einsum("... i d, ... j d -> ... i j", q, k)
+
+        if exists(media_locations):
+            # at each boolean of True, increment the time counter (relative to media time)
+            text_time = media_locations.cumsum(dim=-1)
+            media_time = torch.arange(T_img, device=x.device) + 1
+
+            if not attend_previous:
+                text_time[~media_locations] += 1
+                # make sure max is still the number of images in the sequence
+                text_time[
+                    text_time
+                    > repeat(
+                        torch.count_nonzero(media_locations, dim=1),
+                        "b -> b i",
+                        i=text_time.shape[1],
+                    )
+                ] = 0
+
+            # text time must equal media time if only attending to most immediate image
+            # otherwise, as long as text time is greater than media time (if attending to all previous images / media)
+            mask_op = torch.eq if self.only_attend_immediate_media else torch.ge
+
+            text_to_media_mask = mask_op(
+                rearrange(text_time, "b i -> b 1 i 1"),
+                repeat(media_time, "j -> 1 1 1 (j n)", n=n),
+            )
+            sim = sim.masked_fill(~text_to_media_mask, -torch.finfo(sim.dtype).max)
+
+        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
+        attn = sim.softmax(dim=-1)
+
+        if exists(media_locations) and self.only_attend_immediate_media:
+            # any text without a preceding media needs to have attention zeroed out
+            text_without_media_mask = text_time == 0
+            text_without_media_mask = rearrange(text_without_media_mask, "b i -> b 1 i 1")
+            attn = attn.masked_fill(text_without_media_mask, 0.0)
+
+        out = einsum("... i j, ... j d -> ... i d", attn, v)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)
+
+
+class GatedCrossAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_visual,
+        dim_head=64,
+        heads=8,
+        ff_mult=4,
+        only_attend_immediate_media=True,
+    ):
+        super().__init__()
+        self.attn = MaskedCrossAttention(
+            dim=dim,
+            dim_visual=dim_visual,
+            dim_head=dim_head,
+            heads=heads,
+            only_attend_immediate_media=only_attend_immediate_media,
+        )
+        self.attn_gate = nn.Parameter(torch.tensor([0.0]))
+
+        self.ff = FeedForward(dim, mult=ff_mult)
+        self.ff_gate = nn.Parameter(torch.tensor([0.0]))
+
+    def forward(
+        self,
+        x,
+        media,
+        media_locations=None,
+        attend_previous=True,
+    ):
+        x = (
+            self.attn(
+                x,
+                media,
+                media_locations=media_locations,
+                attend_previous=attend_previous,
+            )
+            * self.attn_gate.tanh()
+            + x
+        )
+        x = self.ff(x) * self.ff_gate.tanh() + x
+
+        return x

mmgpt/models/open_flamingo/utils.py

@@ -0,0 +1,31 @@
+def extend_instance(obj, mixin):
+    """Apply mixins to a class instance after creation"""
+    base_cls = obj.__class__
+    base_cls_name = obj.__class__.__name__
+    obj.__class__ = type(
+        base_cls_name, (mixin, base_cls), {}
+    )  # mixin needs to go first for our forward() logic to work
+
+
+def getattr_recursive(obj, att):
+    """
+    Return nested attribute of obj
+    Example: getattr_recursive(obj, 'a.b.c') is equivalent to obj.a.b.c
+    """
+    if att == "":
+        return obj
+    i = att.find(".")
+    if i < 0:
+        return getattr(obj, att)
+    else:
+        return getattr_recursive(getattr(obj, att[:i]), att[i + 1 :])
+
+
+def setattr_recursive(obj, att, val):
+    """
+    Set nested attribute of obj
+    Example: setattr_recursive(obj, 'a.b.c', val) is equivalent to obj.a.b.c = val
+    """
+    if "." in att:
+        obj = getattr_recursive(obj, ".".join(att.split(".")[:-1]))
+    setattr(obj, att.split(".")[-1], val)

mmgpt/train/__init__.py

@@ -0,0 +1 @@
+

mmgpt/train/distributed.py

@@ -0,0 +1,131 @@
+"""Modified from https://github.com/mlfoundations/open_flamingo"""
+import os
+
+import torch
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+def is_global_master(args):
+    return args.rank == 0
+
+
+def is_local_master(args):
+    return args.local_rank == 0
+
+
+def is_master(args, local=False):
+    return is_local_master(args) if local else is_global_master(args)
+
+
+def is_using_horovod():
+    # NOTE w/ horovod run, OMPI vars should be set, but w/ SLURM PMI vars will be set
+    # Differentiating between horovod and DDP use via SLURM may not be possible, so horovod arg still required...
+    ompi_vars = ["OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"]
+    pmi_vars = ["PMI_RANK", "PMI_SIZE"]
+    if all([var in os.environ for var in ompi_vars]) or all([var in os.environ for var in pmi_vars]):
+        return True
+    else:
+        return False
+
+
+def is_using_distributed():
+    if "WORLD_SIZE" in os.environ:
+        return int(os.environ["WORLD_SIZE"]) > 1
+    if "SLURM_NTASKS" in os.environ:
+        return int(os.environ["SLURM_NTASKS"]) > 1
+    return False
+
+
+def world_info_from_env():
+    local_rank = 0
+    for v in (
+        "LOCAL_RANK",
+        "MPI_LOCALRANKID",
+        "SLURM_LOCALID",
+        "OMPI_COMM_WORLD_LOCAL_RANK",
+    ):
+        if v in os.environ:
+            local_rank = int(os.environ[v])
+            break
+    global_rank = 0
+    for v in ("RANK", "PMI_RANK", "SLURM_PROCID", "OMPI_COMM_WORLD_RANK"):
+        if v in os.environ:
+            global_rank = int(os.environ[v])
+            break
+    world_size = 1
+    for v in ("WORLD_SIZE", "PMI_SIZE", "SLURM_NTASKS", "OMPI_COMM_WORLD_SIZE"):
+        if v in os.environ:
+            world_size = int(os.environ[v])
+            break
+
+    return local_rank, global_rank, world_size
+
+
+def init_distributed_device(args):
+    # Distributed training = training on more than one GPU.
+    # Works in both single and multi-node scenarios.
+    args.distributed = False
+    args.world_size = 1
+    args.rank = 0  # global rank
+    args.local_rank = 0
+    if args.horovod:
+        assert hvd is not None, "Horovod is not installed"
+        hvd.init()
+        args.local_rank = int(hvd.local_rank())
+        args.rank = hvd.rank()
+        args.world_size = hvd.size()
+        args.distributed = True
+        os.environ["LOCAL_RANK"] = str(args.local_rank)
+        os.environ["RANK"] = str(args.rank)
+        os.environ["WORLD_SIZE"] = str(args.world_size)
+    elif is_using_distributed():
+        if "SLURM_PROCID" in os.environ:
+            # DDP via SLURM
+            args.local_rank, args.rank, args.world_size = world_info_from_env()
+            # SLURM var -> torch.distributed vars in case needed
+            os.environ["LOCAL_RANK"] = str(args.local_rank)
+            os.environ["RANK"] = str(args.rank)
+            os.environ["WORLD_SIZE"] = str(args.world_size)
+            torch.distributed.init_process_group(
+                backend=args.dist_backend,
+                init_method=args.dist_url,
+                world_size=args.world_size,
+                rank=args.rank,
+            )
+        else:
+            # DDP via torchrun, torch.distributed.launch
+            args.local_rank, _, _ = world_info_from_env()
+            torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url)
+            args.world_size = torch.distributed.get_world_size()
+            args.rank = torch.distributed.get_rank()
+        args.distributed = True
+    else:
+        # needed to run on single gpu
+        torch.distributed.init_process_group(
+            backend=args.dist_backend,
+            init_method=args.dist_url,
+            world_size=1,
+            rank=0,
+        )
+
+    if torch.cuda.is_available():
+        if args.distributed and not args.no_set_device_rank:
+            device = "cuda:%d" % args.local_rank
+        else:
+            device = "cuda:0"
+        torch.cuda.set_device(device)
+    else:
+        device = "cpu"
+    args.device = device
+    device = torch.device(device)
+    return device
+
+
+def is_rank0():
+    if not torch.distributed.is_initialized():
+        return True
+    return torch.distributed.get_rank() == 0

mmgpt/train/instruction_finetune.py

@@ -0,0 +1,460 @@
+"""Modified from https://github.com/mlfoundations/open_flamingo"""
+
+import argparse
+import copy
+import glob
+import os
+import random
+import time
+
+import numpy as np
+import torch
+import wandb
+from mmengine import Config
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, DistributedSampler
+from tqdm import tqdm
+from transformers import (
+    get_constant_schedule_with_warmup,
+    get_cosine_schedule_with_warmup,
+    get_linear_schedule_with_warmup,
+)
+
+from mmgpt import create_model_and_transforms
+from mmgpt.models.builder import create_toy_model_and_transforms
+from mmgpt.datasets import InfiniteSampler, build_dataset
+from mmgpt.train.distributed import init_distributed_device, world_info_from_env
+from mmgpt.train.train_utils import AverageMeter, get_autocast, get_cast_dtype, get_checkpoint
+
+
+def random_seed(seed=42, rank=0):
+    torch.manual_seed(seed + rank)
+    np.random.seed(seed + rank)
+    random.seed(seed + rank)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--vision_encoder_path", default="ViT-L-14", type=str)
+    parser.add_argument("--vision_encoder_pretrained", default="openai", type=str)
+    parser.add_argument("--lm_path", default="checkpoints/llama-7b_hf", type=str)
+    parser.add_argument(
+        "--tokenizer_path",
+        default="checkpoints/llama-7b_hf",
+        type=str,
+        help="path to tokenizer",
+    )
+    parser.add_argument(
+        "--pretrained_path",
+        default="checkpoints/OpenFlamingo-9B/checkpoint.pt",
+        type=str,
+        help="path to pretrained model",
+    )
+    parser.add_argument(
+        "--run_name",
+        type=str,
+        default="train-my-gpt4",
+        help="used to name saving directory and wandb run",
+    )
+    parser.add_argument("--use_media_placement_augmentation", action="store_true")
+    parser.add_argument("--offline", action="store_true")
+    parser.add_argument("--num_epochs", type=int, default=1)
+    parser.add_argument("--logging_steps", type=int, default=100, help="log loss every n steps")
+    # Sum of gradient optimization batch size
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        help="path to checkpoint to resume from, this should contain model, optimizer, and lr_scheduler states",
+        default=None,
+    )
+    parser.add_argument(
+        "--delete_previous_checkpoint",
+        action="store_true",
+        help="delete previous checkpoint when saving new checkpoint",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--learning_rate", default=1e-5, type=float)
+    parser.add_argument(
+        "--lr_scheduler",
+        default="constant",
+        type=str,
+        help="constant, linear, or cosine",
+    )
+    parser.add_argument("--warmup_steps", default=100, type=int)
+    parser.add_argument("--weight_decay", default=0.1, type=float)
+    parser.add_argument(
+        "--precision",
+        choices=["amp", "amp_bf16", "amp_bfloat16", "bf16", "fp16", "fp32"],
+        default="amp",
+        help="Floating point precision.",
+    )
+    # data args
+    parser.add_argument("--workers", type=int, default=0)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--dataset_config", type=str, default=None, help="path to dataset config file")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=16)
+    # Finetune config
+    parser.add_argument("--tuning_config", type=str, default=None, help="path to tuning config file")
+    # distributed training args
+    parser.add_argument(
+        "--dist-url",
+        default="env://",
+        type=str,
+        help="url used to set up distributed training",
+    )
+    parser.add_argument("--dist-backend", default="nccl", type=str, help="distributed backend")
+    parser.add_argument(
+        "--horovod",
+        default=False,
+        action="store_true",
+        help="Use horovod for distributed training.",
+    )
+    parser.add_argument(
+        "--no-set-device-rank",
+        default=False,
+        action="store_true",
+        help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc).",
+    )
+    # wandb args
+    parser.add_argument("--report_to_wandb", default=False, action="store_true")
+    parser.add_argument(
+        "--wandb_project",
+        type=str,
+    )
+    parser.add_argument(
+        "--wandb_entity",
+        type=str,
+    )
+    parser.add_argument(
+        "--save_checkpoints_to_wandb",
+        default=False,
+        action="store_true",
+        help="save checkpoints to wandb",
+    )
+
+    args = parser.parse_args()
+
+    if args.save_checkpoints_to_wandb and not args.report_to_wandb:
+        raise ValueError("save_checkpoints_to_wandb requires report_to_wandb")
+
+    if args.offline:
+        os.environ["WANDB_MODE"] = "offline"
+        os.environ["TRANSFORMERS_OFFLINE"] = "1"
+
+    args.local_rank, args.rank, args.world_size = world_info_from_env()
+
+    if args.rank == 0:
+        if not os.path.exists(args.run_name):
+            os.makedirs(args.run_name)
+
+    device_id = init_distributed_device(args)
+
+    random_seed(args.seed)
+
+    if args.tuning_config is not None:
+        tuning_config = Config.fromfile(args.tuning_config)
+    else:
+        raise ValueError("tuning_config must be specified")
+
+    model, image_processor, tokenizer = create_model_and_transforms(
+        model_name="open_flamingo",
+        clip_vision_encoder_path=args.vision_encoder_path,
+        clip_vision_encoder_pretrained=args.vision_encoder_pretrained,
+        lang_encoder_path=args.lm_path,
+        tokenizer_path=args.tokenizer_path if args.tokenizer_path else args.lm_path,
+        use_media_placement_augmentation=args.use_media_placement_augmentation,
+        pretrained_model_path=args.pretrained_path,
+        tuning_config=tuning_config.tuning_config,
+    )
+
+    if args.dataset_config is not None:
+        dataset_config = Config.fromfile(args.dataset_config)
+    else:
+        raise ValueError("dataset_config must be specified")
+
+    dataset = build_dataset(
+        dataset_config=dataset_config.visual_datasets,
+        vis_processor=image_processor,
+        tokenizer=tokenizer,
+    )
+    train_dataloader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=args.workers,
+        sampler=DistributedSampler(dataset, shuffle=True, drop_last=True),
+        collate_fn=dataset.collater,
+    )
+
+    # build language dataset and dataloader for multi-modality training
+    if dataset_config.get('language_datasets') is not None and len(dataset_config.language_datasets) > 0:
+        lang_dataset = build_dataset(
+            dataset_config=dataset_config.language_datasets,
+            tokenizer=tokenizer,
+        )
+        lang_dataloader = DataLoader(
+            lang_dataset,
+            batch_size=args.batch_size,
+            num_workers=args.workers,
+            sampler=InfiniteSampler(lang_dataset, shuffle=True),
+            collate_fn=lang_dataset.collater,
+        )
+        lang_dataloader = iter(lang_dataloader)
+    else:
+        lang_dataloader = None
+
+    random_seed(args.seed, args.rank)
+
+    print(f"Start running training on rank {args.rank}.")
+
+    if args.rank == 0 and args.report_to_wandb:
+        wandb.init(
+            project=args.wandb_project,
+            entity=args.wandb_entity,
+            name=args.run_name,
+            config=vars(args),
+        )
+
+    device_id = args.rank % torch.cuda.device_count()
+    model = model.to(device_id)
+
+    ddp_model = DDP(model, device_ids=[device_id], find_unused_parameters=True)
+
+    def get_grouped_params(model):
+        params_with_wd, params_without_wd = [], []
+
+        def apply_decay(x):
+            return (
+                "gated_cross_attn_layer" in x
+                and "ff_gate" not in x
+                and "attn_gate" not in x
+                and "norm" not in x
+                and "bias" not in x
+            )
+
+        for n, p in model.named_parameters():
+            # if p.requires_grad:
+            if apply_decay(n):
+                params_with_wd.append(p)
+            else:
+                params_without_wd.append(p)
+
+        return [
+            {"params": params_with_wd, "weight_decay": args.weight_decay},
+            {"params": params_without_wd, "weight_decay": 0.0},
+        ]
+
+    optimizer = torch.optim.AdamW(get_grouped_params(ddp_model), lr=args.learning_rate)
+
+    total_training_steps = len(train_dataloader) * args.num_epochs
+
+    if args.rank == 0:
+        print(f"Total training steps: {total_training_steps}")
+
+    if args.lr_scheduler == "linear":
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=args.warmup_steps,
+            num_training_steps=total_training_steps // args.gradient_accumulation_steps,
+        )
+    elif args.lr_scheduler == "cosine":
+        lr_scheduler = get_cosine_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=args.warmup_steps,
+            num_training_steps=total_training_steps // args.gradient_accumulation_steps,
+        )
+    else:
+        lr_scheduler = get_constant_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps)
+
+    # check if a checkpoint exists for this run
+    if os.path.exists(f"{args.run_name}") and args.resume_from_checkpoint is None:
+        checkpoint_list = glob.glob(f"{args.run_name}/checkpoint_*.pt")
+        if len(checkpoint_list) == 0:
+            print(f"Found no checkpoints for run {args.run_name}.")
+        else:
+            args.resume_from_checkpoint = sorted(checkpoint_list, key=lambda x: int(x.split("_")[-1].split(".")[0]))[-1]
+            print(f"Found checkpoint {args.resume_from_checkpoint} for run {args.run_name}.")
+
+    resume_from_epoch = 0
+    if args.resume_from_checkpoint is not None:
+        if args.rank == 0:
+            print(f"Loading checkpoint from {args.resume_from_checkpoint}")
+        checkpoint = torch.load(args.resume_from_checkpoint, map_location="cpu")
+        ddp_model.load_state_dict(checkpoint["model_state_dict"], False)
+        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
+        resume_from_epoch = checkpoint["epoch"] + 1
+
+    ddp_model.train()
+
+    for epoch in range(resume_from_epoch, args.num_epochs):
+        train_dataloader.sampler.set_epoch(epoch)
+
+        train_one_epoch(
+            args=args,
+            model=ddp_model,
+            epoch=epoch,
+            tokenizer=tokenizer,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            train_dataloader=train_dataloader,
+            language_dataloader=lang_dataloader,
+            device_id=device_id,
+            wandb=wandb,
+        )
+
+        if args.rank == 0:
+            if not os.path.exists(args.run_name):
+                os.makedirs(args.run_name)
+
+            checkpoint_dict = {
+                "epoch": epoch,
+                "model_state_dict": get_checkpoint(ddp_model),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "lr_scheduler_state_dict": lr_scheduler.state_dict(),
+                "tuning_config": tuning_config,
+            }
+
+            print(f"Saving checkpoint to {args.run_name}/checkpoint_{epoch}.pt")
+            torch.save(checkpoint_dict, f"{args.run_name}/checkpoint_{epoch}.pt")
+            if args.report_to_wandb and args.save_checkpoints_to_wandb:
+                wandb.save(f"{args.run_name}/checkpoint_{epoch}.pt")
+
+            if args.delete_previous_checkpoint:
+                if epoch > 0:
+                    os.remove(f"{args.run_name}/checkpoint_{epoch-1}.pt")
+    if args.rank == 0:
+        torch.save(
+            {"model_state_dict": get_checkpoint(ddp_model.module), "tuning_config": tuning_config},
+            f"{args.run_name}/final_weights.pt",
+        )
+        if args.report_to_wandb and args.save_checkpoints_to_wandb:
+            wandb.save(f"{args.run_name}/final_weights.pt")
+
+
+def train_one_epoch(
+    args,
+    model,
+    epoch,
+    train_dataloader,
+    language_dataloader,
+    tokenizer,
+    optimizer,
+    lr_scheduler,
+    device_id,
+    wandb,
+):
+    num_batches_per_epoch = len(train_dataloader)
+
+    total_training_steps = num_batches_per_epoch * args.num_epochs
+
+    autocast = get_autocast(args.precision)
+    cast_dtype = get_cast_dtype(args.precision)
+
+    model.train()
+
+    # setup logging
+    step_time_m = AverageMeter()  # time for one optimizer step (> 1 batch if using gradient accum)
+    data_time_m = (
+        AverageMeter()
+    )  # avg time to load one batch of both C4 AND laion (= 1 batch regardless of gradient accum)
+    end = time.time()
+
+    # loop through dataloader
+    for num_steps, batch in tqdm(
+        enumerate(train_dataloader),
+        disable=args.rank != 0,
+        total=total_training_steps,
+        initial=(epoch * num_batches_per_epoch),
+    ):
+        data_time_m.update(time.time() - end)
+
+        global_step = num_steps + epoch * num_batches_per_epoch
+
+        #### VISION FORWARD PASS ####
+        images = batch["image"].to(device_id, dtype=cast_dtype, non_blocking=True).unsqueeze(1).unsqueeze(1)
+        input_ids = batch["input_ids"].to(device_id, dtype=cast_dtype, non_blocking=True)
+        attention_mask = batch["attention_mask"].to(device_id, dtype=cast_dtype, non_blocking=True)
+        labels = batch["labels"].to(device_id, dtype=cast_dtype, non_blocking=True)
+
+        with autocast():
+            loss_batch = model(
+                vision_x=images,
+                lang_x=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+            )[0]
+        loss = loss_batch / args.gradient_accumulation_steps
+        loss_vision = loss  # for logging
+
+        #### BACKWARD PASS ####
+        loss.backward()
+
+        #### LANGUAGE FORWARD PASS ####
+        if language_dataloader is not None:
+            batch_lang = next(language_dataloader)
+            lang_input_ids = batch_lang["input_ids"].to(device_id, dtype=cast_dtype, non_blocking=True)
+            lang_attention_mask = batch_lang["attention_mask"].to(device_id, dtype=cast_dtype, non_blocking=True)
+            lang_labels = batch_lang["labels"].to(device_id, dtype=cast_dtype, non_blocking=True)
+
+            with autocast():
+                lang_loss_batch = model(
+                    vision_x=None,
+                    lang_x=lang_input_ids,
+                    attention_mask=lang_attention_mask,
+                    labels=lang_labels,
+                )[0]
+            lang_loss = lang_loss_batch / args.gradient_accumulation_steps
+            #### BACKWARD PASS ####
+            lang_loss.backward()
+
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
+        # step optimizer and log
+        if (((num_steps + 1) % args.gradient_accumulation_steps) == 0) or (num_steps == num_batches_per_epoch - 1):
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+            # step time and reset end outside of rank 0
+            step_time_m.update(time.time() - end)
+            end = time.time()
+
+            if args.rank == 0 and args.report_to_wandb:
+                # compute within rank 0
+                samples_per_second = (
+                    args.gradient_accumulation_steps * args.batch_size * args.world_size / step_time_m.val
+                )
+                samples_per_second_per_gpu = args.gradient_accumulation_steps * args.batch_size / step_time_m.val
+
+                wandb.log(
+                    {
+                        "data_time": data_time_m.avg,
+                        "step_time": step_time_m.avg,
+                        "samples_per_second": samples_per_second,
+                        "samples_per_second_per_gpu": samples_per_second_per_gpu,
+                        "lr": optimizer.param_groups[0]["lr"],
+                    },
+                    commit=False,
+                )
+                step_time_m.reset()
+                data_time_m.reset()
+
+                loss_log = {
+                    "loss": loss.item(),
+                    "loss_vision": loss_vision.item(),
+                    "global_step": global_step,
+                }
+                if language_dataloader is not None:
+                    loss_log["loss_lang"] = lang_loss.item()
+
+                wandb.log(loss_log, commit=True)
+
+        # Log loss to console
+        if ((num_steps + 1) % args.logging_steps == 0) and args.rank == 0:
+            print(
+                f"Step {num_steps+1}/{num_batches_per_epoch} of epoch {epoch+1}/{args.num_epochs} complete. Loss: {loss.item():.3f}"
+            )
+
+
+if __name__ == "__main__":
+    main()

mmgpt/train/train_utils.py

@@ -0,0 +1,251 @@
+"""Modified from https://github.com/mlfoundations/open_flamingo"""
+import time
+from contextlib import suppress
+
+import torch
+from tqdm import tqdm
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == "bf16":
+        cast_dtype = torch.bfloat16
+    elif precision == "fp16":
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_autocast(precision):
+    if precision == "amp":
+        return torch.cuda.amp.autocast
+    elif precision == "amp_bfloat16" or precision == "amp_bf16":
+        # amp_bfloat16 is more stable than amp float16 for clip training
+        return lambda: torch.cuda.amp.autocast(dtype=torch.bfloat16)
+    else:
+        return suppress
+
+
+def train_one_epoch(
+    args,
+    model,
+    epoch,
+    laion_loader,
+    mmc4_loader,
+    tokenizer,
+    optimizer,
+    lr_scheduler,
+    device_id,
+    wandb,
+):
+    num_batches_per_epoch_laion = laion_loader.num_batches
+    num_batches_per_epoch_mmc4 = mmc4_loader.num_batches
+
+    assert (
+        num_batches_per_epoch_laion == num_batches_per_epoch_mmc4
+    ), "Number of batches in laion and mmc4 datasets must be the same"
+    num_batches_per_epoch = num_batches_per_epoch_mmc4
+    total_training_steps = num_batches_per_epoch * args.num_epochs
+
+    autocast = get_autocast(args.precision)
+    cast_dtype = get_cast_dtype(args.precision)
+
+    media_token_id = tokenizer("<image>", add_special_tokens=False)["input_ids"][-1]
+    endofchunk_token_id = tokenizer("<|endofchunk|>", add_special_tokens=False)["input_ids"][-1]
+
+    model.train()
+
+    # setup logging
+    step_time_m = AverageMeter()  # time for one optimizer step (> 1 batch if using gradient accum)
+    data_time_m = (
+        AverageMeter()
+    )  # avg time to load one batch of both C4 AND laion (= 1 batch regardless of gradient accum)
+    end = time.time()
+
+    # loop through dataloader
+    for num_steps, (batch_laion, batch_mmc4) in tqdm(
+        enumerate(zip(laion_loader, mmc4_loader)),
+        disable=args.rank != 0,
+        total=total_training_steps,
+        initial=(epoch * num_batches_per_epoch),
+    ):
+        data_time_m.update(time.time() - end)
+
+        global_step = num_steps + epoch * num_batches_per_epoch
+
+        #### LAION FORWARD PASS ####
+        images = batch_laion[0].to(device_id, dtype=cast_dtype, non_blocking=True).unsqueeze(1).unsqueeze(1)
+
+        input_ids = batch_laion[1][0].to(device_id, dtype=cast_dtype, non_blocking=True)
+        attention_mask = batch_laion[1][1].to(device_id, dtype=cast_dtype, non_blocking=True)
+
+        labels = input_ids.clone()
+        labels[labels == tokenizer.pad_token_id] = -100
+        labels[:, 0] = -100
+        labels[labels == media_token_id] = -100
+        labels.to(device_id)
+
+        with autocast():
+            loss_laion = model(
+                vision_x=images,
+                lang_x=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+            )[0]
+        divided_loss_laion = loss_laion / args.gradient_accumulation_steps
+
+        #### C4 FORWARD PASS ####
+        images = batch_mmc4[0].to(device_id, dtype=cast_dtype, non_blocking=True).unsqueeze(2)
+        input_ids = torch.stack([x[0] for x in batch_mmc4[1]]).squeeze(1)
+        attention_mask = torch.stack([x[1] for x in batch_mmc4[1]]).squeeze(1)
+
+        # NOTE: irena: expected shape of clip_text_input_ids / attention_mask is (N, I, max_seq_len)
+        labels = input_ids.clone()
+        labels[labels == tokenizer.pad_token_id] = -100
+        labels[:, 0] = -100
+
+        for i in range(labels.shape[0]):
+            # remove loss for any token before the first <image> token
+            label_idx = 0
+            while label_idx < labels.shape[1] and labels[i][label_idx] != media_token_id:
+                labels[i][label_idx] = -100
+                label_idx += 1
+
+            # get index of all endofchunk tokens in the sequence
+            endofchunk_idxs = torch.where(labels[i] == endofchunk_token_id)[0]
+            for endofchunk_idx in endofchunk_idxs:
+                token_idx = endofchunk_idx + 1
+                while token_idx < labels.shape[1] and labels[i][token_idx] != media_token_id:
+                    labels[i][token_idx] = -100
+                    token_idx += 1
+
+        labels[labels == media_token_id] = -100
+        labels.to(device_id)
+
+        with autocast():
+            loss_mmc4 = model(
+                vision_x=images,
+                lang_x=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+            )[0]
+
+            # if loss is nan, skip this batch
+            if torch.isnan(loss_mmc4):
+                print("loss is nan, skipping this batch")
+                print("input_ids: ", tokenizer.batch_decode(input_ids))
+                print("labels: ", labels)
+                print("images: ", images)
+                optimizer.zero_grad()
+                continue
+
+        divided_loss_mmc4 = loss_mmc4 / args.gradient_accumulation_steps
+
+        #### BACKWARD PASS ####
+        loss = divided_loss_laion * args.loss_multiplier_laion + divided_loss_mmc4 * args.loss_multiplier_mmc4
+        loss.backward()
+
+        #### MASK GRADIENTS FOR EMBEDDINGS ####
+        # Note (anas): Do not apply weight decay to embeddings as it will break this function.
+        def mask_embedding(m):
+            if isinstance(m, torch.nn.Embedding) and m.weight.requires_grad:
+                zero_mask = torch.zeros_like(m.weight.grad)
+                zero_mask[media_token_id] = torch.ones_like(zero_mask[media_token_id])
+                zero_mask[endofchunk_token_id] = torch.ones_like(zero_mask[endofchunk_token_id])
+                m.weight.grad = m.weight.grad * zero_mask
+
+        model.apply(mask_embedding)
+
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
+        # step optimizer and log
+        if (((num_steps + 1) % args.gradient_accumulation_steps) == 0) or (num_steps == num_batches_per_epoch - 1):
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+            # step time and reset end outside of rank 0
+            step_time_m.update(time.time() - end)
+            end = time.time()
+
+            if args.rank == 0 and args.report_to_wandb:
+                # compute within rank 0
+                laion_samples_per_second = (
+                    args.gradient_accumulation_steps * args.batch_size_laion * args.world_size / step_time_m.val
+                )
+                laion_samples_per_second_per_gpu = (
+                    args.gradient_accumulation_steps * args.batch_size_laion / step_time_m.val
+                )
+
+                c4_samples_per_second = (
+                    args.gradient_accumulation_steps * args.batch_size_mmc4 * args.world_size / step_time_m.val
+                )
+                c4_samples_per_second_per_gpu = (
+                    args.gradient_accumulation_steps * args.batch_size_mmc4 / step_time_m.val
+                )
+
+                wandb.log(
+                    {
+                        "data_time": data_time_m.avg,
+                        "step_time": step_time_m.avg,
+                        "laion_samples_per_second": laion_samples_per_second,
+                        "laion_samples_per_second_per_gpu": laion_samples_per_second_per_gpu,
+                        "c4_samples_per_second": c4_samples_per_second,
+                        "c4_samples_per_second_per_gpu": c4_samples_per_second_per_gpu,
+                        "lr": optimizer.param_groups[0]["lr"],
+                    },
+                    commit=False,
+                )
+                step_time_m.reset()
+                data_time_m.reset()
+
+                wandb.log(
+                    {
+                        "loss_laion": divided_loss_laion.item(),
+                        "global_step": global_step,
+                    },
+                    commit=False,
+                )
+                wandb.log(
+                    {"loss_mmc4": divided_loss_mmc4.item(), "global_step": global_step},
+                    commit=True,
+                )
+
+        # Log loss to console
+        if ((num_steps + 1) % args.logging_steps == 0) and args.rank == 0:
+            print(
+                f"Step {num_steps+1}/{num_batches_per_epoch} of epoch {epoch+1}/{args.num_epochs} complete. Loss LAION: {loss_laion.item():.3f} // Loss MMC4: {loss_mmc4.item():.3f}"
+            )
+
+
+def get_checkpoint(model: torch.nn.Module):
+    state_dict = model.state_dict()
+    parameters = {k: v for k, v in model.named_parameters()}
+    # remove duplicate parameters
+    duplicate_keys = set(state_dict.keys()) - set(parameters.keys())
+    for k in duplicate_keys:
+        del state_dict[k]
+    # remove non-grad parameters
+    for name, p in parameters.items():
+        if not p.requires_grad:
+            del state_dict[name]
+
+    return state_dict
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count

requirements.txt

@@ -0,0 +1,20 @@
+einops
+einops-exts
+transformers
+peft
+bigmodelvis
+torch
+torchvision
+pillow
+more-itertools
+datasets
+braceexpand
+webdataset
+wandb
+nltk
+scipy
+inflection
+sentencepiece
+open_clip_torch
+mmengine
+gradio

setup.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+from setuptools import find_packages, setup
+
+if __name__ == "__main__":
+    with Path(Path(__file__).parent, "README.md").open(encoding="utf-8") as file:
+        long_description = file.read()
+
+    # TODO: This is a hack to get around the fact that we can't read the requirements.txt file, we should fix this.
+    # def _read_reqs(relpath):
+    #     fullpath = os.path.join(Path(__file__).parent, relpath)
+    #     with open(fullpath) as f:
+    #         return [
+    #             s.strip()
+    #             for s in f.readlines()
+    #             if (s.strip() and not s.startswith("#"))
+    #         ]
+
+    REQUIREMENTS = [
+        "einops",
+        "einops-exts",
+        "transformers",
+        "torch",
+        "torchvision",
+        "pillow",
+        "more-itertools",
+        "datasets",
+        "braceexpand",
+        "webdataset",
+        "wandb",
+        "nltk",
+        "scipy",
+        "inflection",
+        "sentencepiece",
+        "open_clip_torch",
+    ]
+
+    setup(
+        name="mmgpt",
+        packages=find_packages(),
+        include_package_data=True,
+        version="0.0.1",
+        license="Apache 2.0",
+        description="An open-source framework for multi-modality instruction fine-tuning",
+        long_description=long_description,
+        long_description_content_type="text/markdown",
+        data_files=[(".", ["README.md"])],
+        keywords=["machine learning"],
+        install_requires=REQUIREMENTS,
+    )