Upload 32 files

.gitattributes

@@ -39,3 +39,5 @@ data/go1.4-basic.obo filter=lfs diff=lfs merge=lfs -text
 data/swissprot_exp/train_exp_prompt_bp_new.csv filter=lfs diff=lfs merge=lfs -text
 data/swissprot_exp/train_exp_prompt_cc_new.csv filter=lfs diff=lfs merge=lfs -text
 data/swissprot_exp/train_exp_prompt_mf_new.csv filter=lfs diff=lfs merge=lfs -text
+dataset_card/imgs/coco_caption.png filter=lfs diff=lfs merge=lfs -text
+docs/_static/merlion.png filter=lfs diff=lfs merge=lfs -text

dataset_card/coco_caption.md

@@ -0,0 +1,41 @@
+![Samples from the COCO Caption dataset (Image credit: "https://arxiv.org/pdf/1504.00325.pdf").](imgs/coco_caption.png)(Samples from the COCO Caption dataset. Image credit: "https://arxiv.org/pdf/1504.00325.pdf")
+
+# Microsoft COCO Dataset (Captioning)
+
+## Description
+[Microsoft COCO Captions dataset](https://github.com/tylin/coco-caption) contains over one and a half million captions describing over 330,000 images. For the training and validation images, five independent human generated captions are be provided for each image.
+
+## Task
+
+(from https://paperswithcode.com/task/image-captioning)
+
+**Image captioning** is the task of describing the content of an image in words. This task lies at the intersection of computer vision and natural language processing. Most image captioning systems use an encoder-decoder framework, where an input image is encoded into an intermediate representation of the information in the image, and then decoded into a descriptive text sequence.
+
+## Metrics
+Models are typically evaluated according to a [BLEU](https://aclanthology.org/P02-1040/) or [CIDER](https://www.cv-foundation.org/openaccess/content_cvpr_2015/papers/Vedantam_CIDEr_Consensus-Based_Image_2015_CVPR_paper.pdf) metric.
+
+## Leaderboard
+
+(Ranked by BLEU-4)
+
+| Rank |  Model  | BLEU-4 | CIDEr | METEOR | SPICE |                                                                    Resources                                                                     |
+| ---- | :-----: | :----: | :---: | :----: | :---: | :----------------------------------------------------------------------------------------------------------------------------------------------: |
+| 1    |   OFA   |  44.9  | 154.9 |  32.5  | 26.6  |                                [paper](https://arxiv.org/abs/2202.03052), [code](https://github.com/OFA-Sys/OFA)                                 |
+| 2    |  LEMON  |  42.6  | 145.5 |  31.4  | 25.5  |                                                                    [paper]()                                                                     |
+| 3    |  CoCa  |   40.9   |  143.6  | 33.9 | 24.7 | [paper](https://arxiv.org/pdf/2205.01917.pdf) |
+| 4    | SimVLM  |  40.6  | 143.3 |  33.7  | 25.4  |                                                [paper](https://openreview.net/pdf?id=GUrhfTuf_3)                                                 |
+| 5    |  VinVL  |  41.0  | 140.9 |  31.1  | 25.2  |                           [paper](https://arxiv.org/pdf/2101.00529v2.pdf), [code](https://github.com/microsoft/Oscar)                            |
+| 6    |  OSCAR  |  40.7  | 140.0 |  30.6  | 24.5  |                           [paper](https://arxiv.org/pdf/2004.06165v5.pdf), [code](https://github.com/microsoft/Oscar)                            |
+| 7    |  BLIP   |  40.4  | 136.7 |  31.4  | 24.3  | [paper](https://arxiv.org/pdf/2201.12086.pdf), [code](https://github.com/salesforce/BLIP), [demo](https://huggingface.co/spaces/Salesforce/BLIP) |
+| 8    |   M^2   |  39.1  | 131.2 |  29.2  | 22.6  |                 [paper](https://arxiv.org/pdf/1912.08226v2.pdf), [code](https://github.com/aimagelab/meshed-memory-transformer)                  |
+| 9    |  BUTD   |  36.5  | 113.5 |  27.0  | 20.3  |               [paper](https://arxiv.org/abs/1707.07998?context=cs), [code](https://github.com/peteanderson80/bottom-up-attention)                |
+| 10    | ClipCap |  32.2  | 108.4 |  27.1  | 20.1  |                     [paper](https://arxiv.org/pdf/2111.09734v1.pdf), [code](https://github.com/rmokady/clip_prefix_caption)                      |
+
+## Auto-Downloading
+
+```
+cd lavis/datasets/download_scripts && python download_coco.py
+```
+
+## References
+"Microsoft COCO Captions: Data Collection and Evaluation Server", Xinlei Chen, Hao Fang, Tsung-Yi Lin, Ramakrishna Vedantam, Saurabh Gupta, Piotr Dollar, C. Lawrence Zitnick

docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/benchmark.rst

@@ -0,0 +1,348 @@
+Benchmark
+############
+
+We provide scripts for evaluating and training models on task datasets. The following benchmark results are included for reference.
+
+
+ALBEF
+*******
+.. list-table::
+   :widths: 30 80 20
+
+   * - **Pretraining**
+     - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/pretrain.sh>`__
+   * -
+     - Visual Genome (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_vg.py>`__)
+     -
+   * -
+     - SBU (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_sbu.py>`__)
+     -
+   * -
+     - CC3M (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/DownloadConceptualCaptions/download_data_cc3m.py>`__)
+     -
+   * -
+     - CC12M (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/DownloadConceptualCaptions/download_data_cc12m.py>`__)
+     -
+
+.. list-table::
+   :widths: 30 40 20 20 20 30 30
+   :header-rows: 1
+
+   * -
+     - **Retrieval**
+     - **R1**
+     - **R5**
+     - **R10**
+     - **Training**
+     - **Evaluation**
+   * - TR
+     - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 77.6
+     - 94.1
+     - 97.2
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_coco_retrieval_albef.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/eval/eval_coco_retrieval.sh>`__
+   * - IR
+     - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 61.0
+     - 84.5
+     - 90.7
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_coco_retrieval_albef.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/eval/eval_coco_retrieval.sh>`__
+   * - TR
+     - Flickr30k (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_flickr.py>`__)
+     - 77.6
+     - 94.1
+     - 97.2
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_flickr30k_retrieval_albef.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/eval/eval_flickr30k_retrieval.sh>`__
+   * - IR
+     - Flickr30k (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_flickr.py>`__)
+     - 61.0
+     - 84.5
+     - 90.7
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_flickr30k_retrieval_albef.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/eval/eval_flickr30k_retrieval.sh>`__
+
+
+.. list-table::
+   :widths: 20 20 20 20 20
+   :header-rows: 1
+
+   * - **VQA**
+     - **test-dev**
+     - **test-std/test**
+     - **Training**
+     - **Evaluation**
+   * - VQAv2 (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 76.35
+     - 76.54
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_vqa_albef.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/eval/test_albef_vqa.sh>`__
+   * - OKVQA (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - NA
+     - 54.7 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_okvqa_albef.sh>`__
+     - NA
+   * - AOKVQA (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 54.5
+     - NA
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_aokvqa_albef.sh>`__
+     - NA
+
+  
+.. list-table::
+   :widths: 20 20 20 20 20
+   :header-rows: 1
+
+   * - **Multimodal Classification**
+     - **val**
+     - **test**
+     - **Training**
+     - **Evaluation**
+   * - SNLI-VE (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 80.60
+     - 81.04
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_ve_albef.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/eval/eval_albef_ve.sh>`__
+   * - NLVR2 (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 82.47 
+     - 82.91 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_nlvr_albef.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/eval/eval_albef_nlvr.sh>`__
+  
+BLIP
+*******
+.. list-table::
+   :widths: 30 80 20
+
+   * - **Pretraining (14M)**
+     - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/pretrain.sh>`__
+   * -
+     - Visual Genome (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_vg.py>`__)
+     -
+   * -
+     - SBU (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_sbu.py>`__)
+     -
+   * -
+     - CC3M (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/DownloadConceptualCaptions/download_data_cc3m.py>`__)
+     -
+   * -
+     - CC12M (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/DownloadConceptualCaptions/download_data_cc12m.py>`__)
+     -
+
+.. list-table::
+   :widths: 30 40 20 20 20 30 30
+   :header-rows: 1
+
+   * - **Tasks**
+     - **Retrieval**
+     - **R1**
+     - **R5**
+     - **R10**
+     - **Training**
+     - **Evaluation**
+   * - TR
+     - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 82.0
+     - 95.8
+     - 98.1
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/train_retrieval_coco.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_ret_coco.sh>`__
+   * - IR
+     - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 64.5
+     - 86.0
+     - 91.7
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/train_retrieval_coco.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_ret_coco.sh>`__
+   * - TR
+     - Flickr30k (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_flickr.py>`__)
+     - 96.9
+     - 99.9
+     - 100.0
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/train_retrieval_flickr.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_ret_flickr.sh>`__
+   * - IR
+     - Flickr30k (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_flickr.py>`__)
+     - 87.5
+     - 97.6
+     - 98.9
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/train_retrieval_flickr.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_ret_flickr.sh>`__
+
+
+.. list-table::
+   :widths: 20 20 20 20 20
+   :header-rows: 1
+
+   * - **VQA**
+     - **test-dev**
+     - **test-std/test**
+     - **Training**
+     - **Evaluation**
+   * - VQAv2 (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 78.23
+     - 78.29
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/train/train_vqa_albef.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/albef/eval/test_albef_vqa.sh>`__
+   * - OKVQA (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - NA
+     - 55.4 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/train_okvqa.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_okvqa.sh>`__
+   * - AOKVQA (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 56.2
+     - 50.1 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/train_aokvqa.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_aokvqa.sh>`__
+
+
+.. list-table::
+   :widths: 20 20 20 20 20 20
+   :header-rows: 1
+
+   * - **Image Captioning**
+     - **BLEU@4**
+     - **CIDEr**
+     - **SPICE**
+     - **Training**
+     - **Evaluation**
+   * - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 39.9
+     - 133.5
+     - 23.7
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/train_caption_coco.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_coco_cap.sh>`__
+   * - NoCaps (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_nocaps.py>`__)
+     - 31.9
+     - 109.1
+     - 14.7
+     - NA
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_nocaps.sh>`__
+
+
+.. list-table::
+   :widths: 20 20 20 20 20
+   :header-rows: 1
+
+   * - **Multimodal Classification**
+     - **val**
+     - **test**
+     - **Training**
+     - **Evaluation**
+   * - NLVR2 (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 82.48
+     - 83.25
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/train/train_nlvr.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/blip/eval/eval_nlvr.sh>`__
+
+CLIP
+*******
+.. list-table::
+   :widths: 30 40 20 20 20 30
+   :header-rows: 1
+
+   * - **Tasks**
+     - **Retrieval (Zero-shot)**
+     - **R1**
+     - **R5**
+     - **R10**
+     - **Evaluation**
+   * - TR
+     - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 57.2
+     - 80.5
+     - 87.8
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/clip/eval/eval_clip_ret_coco.sh>`__
+   * - IR
+     - COCO (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_coco.py>`__)
+     - 36.5
+     - 60.8
+     - 71.0
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/clip/eval/eval_clip_ret_coco.sh>`__
+   * - TR
+     - Flickr30k (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_flickr.py>`__)
+     - 86.5
+     - 98.0
+     - 99.1
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/clip/eval/eval_clip_ret_flickr.sh>`__
+   * - IR
+     - Flickr30k (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_flickr.py>`__)
+     - 67.0
+     - 88.9
+     - 93.3
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/clip/eval/eval_clip_ret_flickr.sh>`__
+
+.. list-table::
+   :widths: 20 20 20
+   :header-rows: 1
+
+   * - **Multimodal Classification**
+     - **val**
+     - **Evaluation**
+   * - ImageNet 
+     - 76.5 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/clip/eval/eval_clip_zs_imnet.sh>`__
+
+
+ALPRO
+*******
+.. list-table::
+   :widths: 30 40 20 20 20 20 30
+   :header-rows: 1
+
+   * - **Tasks**
+     - **Retrieval**
+     - **R1**
+     - **R5**
+     - **R10**
+     - **Training**
+     - **Evaluation**
+   * - TR
+     - MSRVTT (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_msrvtt.py>`__)
+     - 33.2
+     - 60.5 
+     - 71.7 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/train/train_msrvtt_ret.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/eval/eval_msrvtt_ret.sh>`__
+   * - VR
+     - MSRVTT (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_msrvtt.py>`__)
+     - 33.8
+     - 61.4
+     - 72.7
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/train/train_msrvtt_ret.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/eval/eval_msrvtt_ret.sh>`__
+   * - TR
+     - DiDeMo (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_didemo.py>`__)
+     - 38.8 
+     - 66.4
+     - 76.8
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/train/train_didemo_ret.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/eval/eval_didemo_ret.sh>`__
+   * - VR
+     - DiDeMo (`download <https://github.com/salesforce/LAVIS/blob/main/lavis/datasets/download_scripts/download_didemo.py>`__)
+     - 36.6
+     - 67.5
+     - 77.9
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/train/train_didemo_ret.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/eval/eval_didemo_ret.sh>`__
+
+.. list-table::
+   :widths: 20 20 20 20
+   :header-rows: 1
+
+   * - **Video QA**
+     - **test**
+     - **Training**
+     - **Evaluation**
+   * - MSRVTT 
+     - 42.1 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/train/train_msrvtt_qa.sh>`__
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/eval/eval_msrvtt_qa.sh>`__
+   * - MSVD 
+     - 46.0 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/train/train_msvd_qa.sh>`__ 
+     - `script <https://github.com/salesforce/LAVIS/blob/main/run_scripts/alpro/eval/eval_msvd_qa.sh>`__

docs/build_docs.sh

@@ -0,0 +1,101 @@
+#!/bin/bash
+set -euo pipefail
+
+# Change to root directory of repo
+DIRNAME=$(cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+cd "${DIRNAME}/.."
+
+# # Set up virtual environment
+pip3 install setuptools wheel virtualenv
+if [ ! -d venv ]; then
+  rm -f venv
+  virtualenv venv
+fi
+source venv/bin/activate
+
+# # Get current git branch & stash unsaved changes
+GIT_BRANCH=$(git branch --show-current)
+if [ -z "${GIT_BRANCH}" ]; then
+    GIT_BRANCH="main"
+fi
+git stash
+
+# Set up exit handler to restore git state & delete temp branches
+# function exit_handler {
+#     git reset --hard
+#     git checkout "${GIT_BRANCH}" --
+#     git stash pop || true
+#     for version in $(git tag --list 'v[0-9]*'); do
+#         branch="${version}_local_docs_only"
+#         if git show-ref --verify --quiet "refs/heads/$branch"; then
+#             git branch -D "$branch"
+#         fi
+#     done
+# }
+# trap exit_handler EXIT
+
+# Clean up build directory and install Sphinx requirements
+pip3 install -r "${DIRNAME}/requirements.txt"
+sphinx-build -M clean "${DIRNAME}" "${DIRNAME}/_build"
+
+# Build API docs for current head
+export current_version="latest"
+pip3 install "."
+sphinx-build -b html "${DIRNAME}" "${DIRNAME}/_build/html/${current_version}" -W --keep-going
+rm -rf "${DIRNAME}/_build/html/${current_version}/.doctrees"
+#pip3 uninstall -y omnixai
+
+# Install all previous released versions
+# and use them to build the appropriate API docs.
+# Uninstall after we're done with each one.
+# versions=()
+# checkout_files=("${DIRNAME}/*.rst" "lavis" "tutorials" "setup.py")
+# for version in $(git tag --list 'v[0-9]*'); do
+#     versions+=("$version")
+#     git checkout -b "${version}_local_docs_only"
+#     for f in $(git diff --name-only --diff-filter=A "tags/${version}" "${DIRNAME}/*.rst"); do
+#         git rm "$f"
+#     done
+#     git checkout "tags/${version}" -- "${checkout_files[@]}"
+#     export current_version=${version}
+#     pip3 install ".[all]"
+#     sphinx-build -b html "${DIRNAME}" "${DIRNAME}/_build/html/${current_version}" -W --keep-going
+#     rm -rf "${DIRNAME}/_build/html/${current_version}/.doctrees"
+#     #pip3 uninstall -y omnixai
+#     git reset --hard
+#     git checkout "${GIT_BRANCH}" --
+# done
+
+# Determine the latest stable version if there is one
+# if (( ${#versions[@]} > 0 )); then
+#   stable_hash=$(git rev-list --tags --max-count=1)
+#   stable_version=$(git describe --tags "$stable_hash")
+#   export stable_version
+# else
+export stable_version="latest"
+# fi
+
+# Create dummy HTML's for the stable version in the base directory
+while read -r filename; do
+    filename=$(echo "$filename" | sed "s/\.\///")
+    n_sub=$(echo "$filename" | (grep -o "/" || true) | wc -l)
+    prefix=""
+    for (( i=0; i<n_sub; i++ )); do
+        prefix+="../"
+    done
+    url="${prefix}${stable_version}/$filename"
+    mkdir -p "${DIRNAME}/_build/html/$(dirname "$filename")"
+    cat > "${DIRNAME}/_build/html/$filename" <<EOF
+<!DOCTYPE html>
+<html>
+   <head>
+      <title>LAVIS Documentation</title>
+      <meta http-equiv = "refresh" content="0; url='$url'" />
+   </head>
+   <body>
+      <p>Please wait while you're redirected to our <a href="$url">documentation</a>.</p>
+   </body>
+</html>
+EOF
+done < <(cd "${DIRNAME}/_build/html/$stable_version" && find . -name "*.html")
+echo "Finished writing to _build/html."

docs/conf.py

@@ -0,0 +1,56 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+# import os
+# import sys
+# sys.path.insert(0, os.path.abspath('.'))
+
+
+# -- Project information -----------------------------------------------------
+
+project = "LAVIS"
+copyright = "2022, salesforce.com inc."
+author = (
+    "Dongxu Li, Junnan Li, Hung Le, Guangsen Wang, Silvio Savarese, Steven C.H. Hoi"
+)
+
+
+# -- General configuration ---------------------------------------------------
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = ["nbsphinx"]
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ["_templates"]
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = []
+
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+# html_theme = "alabaster"
+html_theme = "sphinx_rtd_theme"
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ["_static"]
+
+# pygments_style = "sphinx"

docs/getting_started.rst

@@ -0,0 +1,233 @@
+Dataset Zoo
+##################
+LAVIS inherently supports a wide variety of common language-vision datasets by providing automatic download scripts to help download and organize these datasets; 
+and implements PyTorch datasets for these datasets. To view supported datasets, use the following code:
+
+.. code-block:: python
+
+    from lavis.datasets.builders import dataset_zoo
+    dataset_names = dataset_zoo.get_names()
+    print(dataset_names)
+    # ['aok_vqa', 'coco_caption', 'coco_retrieval', 'coco_vqa', 'conceptual_caption_12m',
+    #  'conceptual_caption_3m', 'didemo_retrieval', 'flickr30k', 'imagenet', 'laion2B_multi',
+    #  'msrvtt_caption', 'msrvtt_qa', 'msrvtt_retrieval', 'msvd_caption', 'msvd_qa', 'nlvr',
+    #  'nocaps', 'ok_vqa', 'sbu_caption', 'snli_ve', 'vatex_caption', 'vg_caption', 'vg_vqa']
+    print(len(dataset_names))
+    # 23
+
+
+Auto-Downloading and Loading Datasets
+######################################
+We now take COCO caption dataset as an example to demonstrate how to download and prepare the dataset.
+
+In ``lavis/datasets/download_scripts/``, we provide tools to download most common public language-vision datasets supported by LAVIS.
+The COCO caption dataset uses images from COCO dataset. Therefore, we first download COCO images via:
+
+.. code-block:: bash
+    
+    cd lavis/datasets/download_scripts/ && python download_coco.py
+
+This will automatically download and extract COCO images to the default LAVIS cache location.
+The default cache location is ``~/.cache/lavis``, defined in ``lavis/configs/default.yaml``.
+
+After downloading the images, we can use ``load_dataset()`` to obtain the dataset. On the first run, this will automatically download and cache annotation files.
+
+.. code-block:: python
+
+    from lavis.datasets.builders import load_dataset
+    coco_dataset = load_dataset("coco_caption")
+
+    print(coco_dataset.keys())
+    # dict_keys(['train', 'val', 'test'])
+
+    print(len(coco_dataset["train"]))
+    # 566747
+
+    print(coco_dataset["train"][0])
+    # {'image': <PIL.Image.Image image mode=RGB size=640x480>,
+    #  'text_input': 'A woman wearing a net on her head cutting a cake. ',
+    #  'image_id': 0}
+
+If you already host a local copy of the dataset, you can pass in the ``vis_path`` argument to change the default location to load images.
+
+.. code-block:: python
+
+    coco_dataset = load_dataset("coco_caption", vis_path=YOUR_LOCAL_PATH)
+
+
+Model Zoo
+####################################
+LAVIS supports a growing list of pre-trained models for different tasks,
+datatsets and of varying sizes. Let's get started by viewing the supported models.
+
+.. code-block:: python
+
+    from lavis.models import model_zoo
+    print(model_zoo)
+    # ==================================================
+    # Architectures                  Types
+    # ==================================================
+    # albef_classification           base, ve
+    # albef_nlvr                     base
+    # albef_pretrain                 base
+    # albef_retrieval                base, coco, flickr
+    # albef_vqa                      base, vqav2
+    # alpro_qa                       base, msrvtt, msvd
+    # alpro_retrieval                base, msrvtt, didemo
+    # blip_caption                   base, base_coco, large, large_coco
+    # blip_classification            base
+    # blip_feature_extractor         base
+    # blip_nlvr                      base
+    # blip_pretrain                  base
+    # blip_retrieval                 base, coco, flickr
+    # blip_vqa                       base, vqav2
+    # clip                           ViT-B-32, ViT-B-16, ViT-L-14, ViT-L-14-336, RN50
+
+    # show total number of support model variants
+    len(model_zoo)
+    # 33
+
+
+Inference with Pre-trained Models
+####################################
+
+Now let's see how to use models in LAVIS to perform inference on example data. We first
+load a sample image from local.
+
+.. code-block:: python
+
+    from PIL import Image
+
+    # setup device to use
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # load sample image
+    raw_image = Image.open("docs/_static/merlion.png").convert("RGB")
+
+This example image shows `Merlion park <https://en.wikipedia.org/wiki/Merlion>`_ (`image credit <https://theculturetrip.com/asia/singapore/articles/what-exactly-is-singapores-merlion-anyway/>`_), a landmark in Singapore.
+
+.. image:: _static/merlion.png
+
+Image Captioning
+*******************************
+We now use the BLIP model to generate a caption for the image. To make inference even easier, we also associate each
+pre-trained model with its preprocessors (transforms),  we use ``load_model_and_preprocess()`` with the following arguments:
+
+- ``name``: The name of the model to load. This could be a pre-trained model, task model, or feature extractor. See ``model_zoo`` for a full list of model names.
+- ``model_type``: Each architecture has variants trained on different datasets and at different scale. See Types column in ``model_zoo`` for a full list of model types.
+- ``is_eval``: if `True`, set the model to evaluation mode. This is desired for inference or feature extraction.
+- ``device``: device to load the model to.
+
+.. code-block:: python
+
+    from lavis.models import load_model_and_preprocess
+    # loads BLIP caption base model, with finetuned checkpoints on MSCOCO captioning dataset.
+    # this also loads the associated image processors
+    model, vis_processors, _ = load_model_and_preprocess(name="blip_caption", model_type="base_coco", is_eval=True, device=device)
+
+    # preprocess the image
+    # vis_processors stores image transforms for "train" and "eval" (validation / testing / inference)
+    image = vis_processors["eval"](raw_image).unsqueeze(0).to(device)
+
+    # generate caption
+    model.generate({"image": image})
+    # ['a large fountain spewing water into the air']
+
+
+You may also load models and their preprocessors separately via ``load_model()`` and ``load_processor()``.
+In BLIP, you can also generate diverse captions by turning nucleus sampling on.
+
+.. code-block:: python
+
+    from lavis.processors import load_processor
+    from lavis.models import load_model
+
+    # load image preprocesser used for BLIP
+    vis_processor = load_processor("blip_image_eval").build(image_size=384)
+    model = load_model(name="blip_caption", model_type="base_coco", is_eval=True, device=device)
+
+    image = vis_processor(image).unsqueeze(0).to(device)
+    model.generate({"image": raw_image}, use_nucleus_sampling=True)
+    # one generated random sample: ['some very pretty buildings and some water jets']
+
+
+Visual question answering (VQA)
+*******************************
+BLIP model is able to answer free-form questions about images in natural language.
+To access the VQA model, simply replace the ``name`` and ``model_type`` arguments 
+passed to ``load_model_and_preprocess()``.
+
+.. code-block:: python
+
+    from lavis.models import load_model_and_preprocess
+    model, vis_processors, txt_processors = load_model_and_preprocess(name="blip_vqa", model_type="vqav2", is_eval=True, device=device)
+
+    # ask a random question.
+    question = "Which city is this photo taken?"
+    
+    image = vis_processors["eval"](raw_image).unsqueeze(0).to(device)
+    question = txt_processors["eval"](question)
+
+    model.predict_answers(samples={"image": image, "text_input": question}, inference_method="generate")
+    # ['singapore']
+
+
+Unified Feature Extraction Interface
+####################################
+
+LAVIS provides a unified interface to extract multimodal features from each architecture.
+To extract features, we load the feature extractor variants of each model.
+The multimodal feature can be used for multimodal classification. The low-dimensional unimodal features can be used to compute cross-modal similarity.
+
+.. code-block:: python
+
+    from lavis.models import load_model_and_preprocess 
+    
+    model, vis_processors, txt_processors = load_model_and_preprocess(name="blip_feature_extractor", model_type="base", is_eval=True, device=device)
+    caption = "a large fountain spewing water into the air"
+
+    image = vis_processors["eval"](raw_image).unsqueeze(0).to(device)
+    text_input = txt_processors["eval"](caption)
+
+    sample = {"image": image, "text_input": [text_input]}
+
+    features_multimodal = model.extract_features(sample)
+    print(features_multimodal.keys())
+    # odict_keys(['image_embeds', 'multimodal_embeds'])
+    print(features_multimodal.multimodal_embeds.shape)
+    # torch.Size([1, 12, 768]), use features_multimodal[:, 0, :] for multimodal classification tasks
+
+    features_image = model.extract_features(sample, mode="image")
+    print(features_image.keys())
+    # odict_keys(['image_embeds', 'image_embeds_proj'])
+    print(features_image.image_embeds.shape)
+    # torch.Size([1, 197, 768])
+    print(features_image.image_embeds_proj.shape)
+    # torch.Size([1, 197, 256])
+
+    features_text = model.extract_features(sample, mode="text")
+    print(features_text.keys())
+    # odict_keys(['text_embeds', 'text_embeds_proj'])
+    print(features_text.text_embeds.shape)
+    # torch.Size([1, 12, 768])
+    print(features_text.text_embeds_proj.shape)
+    # torch.Size([1, 12, 256])
+    
+    similarity = features_image.image_embeds_proj[:, 0, :] @ features_text.text_embeds_proj[:, 0, :].t()
+    print(similarity)
+    # tensor([[0.2622]])
+
+Since LAVIS supports a unified feature extraction interface, minimal changes are necessary to use a different model as feature extractor. For example,
+to use ALBEF as the feature extractor, one only needs to change the following line:
+
+.. code-block:: python
+
+    model, vis_processors, txt_processors = load_model_and_preprocess(name="albef_feature_extractor", model_type="base", is_eval=True, device=device)
+
+Similarly, to use CLIP as feature extractor: 
+
+.. code-block:: python
+
+    model, vis_processors, txt_processors = load_model_and_preprocess(name="clip_feature_extractor", model_type="base", is_eval=True, device=device)
+    # model, vis_processors, txt_processors = load_model_and_preprocess(name="clip_feature_extractor", model_type="RN50", is_eval=True, device=device)
+    # model, vis_processors, txt_processors = load_model_and_preprocess(name="clip_feature_extractor", model_type="ViT-L-14", is_eval=True, device=device)

docs/index.rst

@@ -0,0 +1,46 @@
+.. LAVIS documentation master file, created by
+   sphinx-quickstart on Sun Jul 31 10:32:27 2022.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+Welcome to LAVIS's documentation!
+=================================
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Introduction
+
+   intro
+
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Getting Started
+
+   getting_started
+
+
+..    :maxdepth: 1
+..    :caption: Advanced Training
+
+..    advanced_training
+
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Advanced Usage
+
+   benchmark
+   tutorial
+
+
+.. Documentations
+.. ===================
+
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`

docs/intro.rst

@@ -0,0 +1,99 @@
+What is LAVIS?
+####################################
+
+LAVIS is a Python deep learning library for LAnguage-and-VISion research and applications.
+It features a unified design to access state-of-the-art foundation language-vision models (`ALBEF <https://arxiv.org/pdf/2107.07651.pdf>`_,
+`BLIP <https://arxiv.org/pdf/2201.12086.pdf>`_, `ALPRO <https://arxiv.org/pdf/2112.09583.pdf>`_, `CLIP <https://arxiv.org/pdf/2103.00020.pdf>`_), common tasks 
+(retrieval, captioning, visual question answering, multimodal classification etc.) and datasets (COCO, Flickr, Nocaps, Conceptual
+Commons, SBU, etc.).
+
+This library aims to provide engineers and researchers with a one-stop solution to rapidly develop models for their specific multimodal
+scenarios, and benchmark them across standard and customized datasets. 
+
+Key features of LAVIS include:
+
+- **Modular and Extensible Library Design**: facilitating to easily utilize and repurpose existing modules (datasets, models, preprocessors), also to add new modules.
+
+- **Easy Off-the-shelf Inference and Feature Extraction**: readily available pre-trained models let you take advantage of state-of-the-art multimodal understanding and generation capabilities on your own data.
+
+- **Reproducible Model Zoo**: provided training/pre-training recipies to easily replicate and extend state-of-the-art models.
+
+- **Dataset Zoo and Automatic Downloading Tools**: it can be a hassle to prepare the many language-vision datasets. LAVIS provides automatic downloaing scripts to help prepare a large variety of datasets and their annotations.
+
+Other features include:
+
+- **Distributed Training** using multiple GPUs on one machine or across multiple machines.
+
+- **Web Demo**: try supported models on your own pictures, questions etc.
+
+- **Leaderboard**: comparing state-of-the-art models across standard datasets. 
+
+- **Dataset Explorer**: help browse and understand language-vision datasets.
+
+Supported Tasks, Models and Datasets
+####################################
+
+The following table shows the supported models and language-vision tasks by LAVIS. Adapting existing models to more tasks is possible and next to come in future releases.
+
+======================================== =========================== ============================================= ============ 
+Tasks                                     Supported Models            Supported Datasets                            Modalities  
+======================================== =========================== ============================================= ============ 
+Image-text Pre-training                   ALBEF, BLIP                 COCO, VisualGenome, SBU, ConceptualCaptions  image, text  
+Image-text Retrieval                      ALBEF, BLIP, CLIP           COCO, Flickr30k                              image, text  
+Text-image Retrieval                      ALBEF, BLIP, CLIP           COCO, Flickr30k                              image, text  
+Visual Question Answering                 ALBEF, BLIP                 VQAv2, OKVQA, A-OKVQA                        image, text  
+Image Captioning                          BLIP                        COCO, NoCaps                                 image, text  
+Image Classification                      CLIP                        ImageNet                                     image        
+Natural Language Visual Reasoning (NLVR)  ALBEF, BLIP                 NLVR2                                        image, text  
+Visual Entailment (VE)                    ALBEF                       SNLI-VE                                      image, text  
+Visual Dialogue                           BLIP                        VisDial                                      image, text  
+Video-text Retrieval                      BLIP, ALPRO                 MSRVTT, DiDeMo                               video, text  
+Text-video Retrieval                      BLIP, ALPRO                 MSRVTT, DiDeMo                               video, text  
+Video Question Answering (VideoQA)        BLIP, ALPRO                 MSRVTT, MSVD                                 video, text  
+Video Dialogue                            VGD-GPT                     AVSD                                         video, text  
+Multimodal Feature Extraction             ALBEF, CLIP, BLIP, ALPRO    customized                                   image, text  
+======================================== =========================== ============================================= ============ 
+
+Library Design
+####################################
+
+.. image:: _static/architecture.png
+  :width: 550
+
+LAVIS has six key modules.
+
+- ``lavis.runners`` manages the overall training and evaluation lifecycle. It is also responsible for creating required components lazily as per demand, such as optimizers, learning rate schedulers and dataloaders. Currently ``RunnerBase`` implements epoch-based training and ``RunerIters`` implements iteration-based training.
+- ``lavis.tasks`` implements concrete training and evaluation logic per task. A task could be, for example, retrieval, captioning, pre-training. The rationale to have an abstraction of task is to accommodate task-specific training and evaluation. For example, evaluating a retrieval model is different from a classification model.
+- ``lavis.datasets`` is responsible for creating datasets, where ``lavis.datasets.builders`` loads dataset configurations, downloads annotations and returns a dataset object; ``lavis.datasets.datasets`` defines the supported datasets, each is a ``torch.utils.data.Dataset`` instance. We also provide `automatic dataset downloading tools` in ``datasets/download_scripts`` to help prepare common public datasets.
+- ``lavis.models`` holds definition for the supported models and shared model layers.
+- ``lavis.processors`` handles preprocessing of text and images/videos before feeding the model. For images and videos, a processor can be thought as transfroms in torchvision; for text input, this may include lowering case, truncation etc.
+- ``lavis.common`` module contains shared classes and methods used by multiple other modules. For example,
+
+   - ``lavis.common.config`` contains classes to store and manipulate configuration files used by LAVIS. In particular, we use a hierarchical configuration design, to allow highly customizable training and evaluation.
+   - ``lavis.common.registry``  serves as a centralized place to manage modules that share the same functionalities. It allows building datasets, models, tasks, and learning rate schedulers during runtime, by specifying their names as string in the configuration file.
+   - ``lavis.common.optims`` contains definitions of learning rate schedulers.
+   - ``lavis.common.dist_utils`` contains utilities for distributed training and evaluation.
+   - ``lavis.common.utils`` contains miscellaneous utilities, mostly IO-related helper functions.
+
+
+Installation
+############
+1. (Optional) Creating conda environment
+
+.. code-block:: bash
+
+   conda create -n lavis python=3.8
+   conda activate lavis
+
+2. Cloning and building from source
+
+.. code-block:: bash
+
+   git clone https://github.com/salesforce/LAVIS.git
+   cd LAVIS
+   pip install .
+
+If you would like to develop on LAVIS, you may find it easier to build with editable mode::
+
+   pip install -e .
+

docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.http://sphinx-doc.org/
+	exit /b 1
+)
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

docs/requirements.txt

@@ -0,0 +1,7 @@
+GitPython
+ipykernel
+nbsphinx==0.8.7
+pandoc
+sphinx
+sphinx_autodoc_typehints
+sphinx_rtd_theme

docs/tutorial.configs.rst

@@ -0,0 +1,172 @@
+.. _config:
+
+Training Models on Task Datasets (Commands and Configurations) 
+#################################################################
+
+LAVIS provides scripts to pre-train and finetune supported models on standard language-vision tasks, stored at ``lavis/run_scripts/``. 
+To replicate the experiments, just run these bash scripts. For example, to train BLIP model on the image-text retrieval task with MSCOCO dataset, we can run
+
+.. code-block::
+
+    bash run_scripts/blip/train/train_retrieval_coco.sh
+
+Inside the scripts, we can see 
+
+.. code-block:: bash
+
+    python -m torch.distributed.run --nproc_per_node=8 train.py --cfg-path lavis/projects/blip/train/retrieval_coco_ft.yaml
+
+where we start a pytorch distributed training on 8 GPUs (you may change according to your own hardware setup). The ``--cfg-path`` specifys a `runtime configuration file`, specifying
+the task, model, dataset and training recipes. 
+
+Available options and their descriptions are as below.
+
+.. LAVIS executes training and evaluation based on arguments specified in the configuration files. The default model and dataset configurations are defined in ``lavis/configs``. The task-specific configurations are defined in ``lavis/projects``. Task-specific configurations have higher priority over the default configurations.
+
+.. The following tables provide explanations for the arguments in the configuration files.
+
+.. list-table::
+   :widths: 30 40
+   :header-rows: 1
+
+   * - Model Configurations
+     - Functionalities
+   * - arch
+     - | name of the model from the model zoo
+       | default: task-dependent
+   * - model_type
+     - | the type of the model (e.g., base)
+       | default: task-dependent
+   * - load_pretrained
+     - | load pretrained weights
+       | default: True (for finetuning task) | False (for pretraining task) 
+   * - load_finetuned
+     - | load task-specific finetuned weights
+       | default: False (for finetuning task) | True (for evaluation) 
+   * - pretrained 
+     - | URL or local path which stores the pretrained model, defined in the default model configuration file
+       | default: task-dependent 
+   * - finetuned
+     - | URL or local path which stores the finetuned model, defined in the default model configuration file
+       | default: task-dependent
+
+.. list-table::
+   :widths: 30 50
+   :header-rows: 1
+
+   * - Dataset Configurations
+     - Functionalities
+   * - vis_processor
+     - | pre-processing of visual input
+       | default: task-dependent
+   * - text_processor
+     - | pre-processing of text input
+       | default: task-dependent
+   * - build_info
+     - | dataset information including the storage location, defined in the default dataset configuration file
+       | default: task-dependent
+
+.. list-table::
+   :widths: 30 50
+   :header-rows: 1
+
+   * - Runtime Configurations
+     - Functionalities
+   * - task
+     - | name of the task
+       | default: task-dependent
+   * - lr_sched
+     - | learning rate schedular
+       | default: linear_warmup_cosine_lr
+   * - init_lr
+     - | initial learning rate (after warmup)
+       | default: task-dependent
+   * - min_lr
+     - | final learning rate after decay
+       | default: task-dependent
+   * - warmup_lr
+     - | starting learning rate for warmup
+       | default: init_lr (no warmup)
+   * - lr_decay_rate
+     - | learning rate decay per epoch for step_lr_shedule
+       | default: 0.9
+   * - warmup_steps
+     - | number of steps for learning rate warmup
+       | default: 0
+   * - max_epoch
+     - | total number of training epochs
+       | default: task-dependent
+   * - weight_decay
+     - | weight decay coefficient for the optimizer
+       | default: 0.05
+   * - batch_size_train
+     - | batch size during training
+       | default: task-dependent
+   * - batch_size_eval
+     - | batch size during evaluation
+       | default: task-dependent
+   * - seed
+     - | pseudo random number generator seed
+       | default: 42
+   * - output_dir
+     - | directory to store logs, results and checkpoints
+       | default: task-dependent
+   * - resume_ckpt_path
+     - | path of the checkpoint to resume training from
+       | default: None
+   * - evaluate
+     - | only perform evaluation without training
+       | default: False
+   * - train_splits
+     - | dataset splits used for training
+       | default: ["train"]
+   * - valid_splits
+     - | dataset splits used for validation
+       | default: ["val"]
+   * - test
+     - | dataset splits used for test
+       | default: ["test"]
+   * - device
+     - | use cpu or gpu (cuda)
+       | default: cuda
+   * - world_size
+     - | number of processes participating in the job
+       | default: 1
+   * - dist_url
+     - | URL specifying how to initialize the process group
+       | default: "env://"
+   * - distributed
+     - | use distributed training
+       | default: True
+   * - amp
+     - | use automatic mixed precision training
+       | default: False
+
+.. list-table::
+   :widths: 40 50
+   :header-rows: 1
+
+   * - Text Generation Configurations
+     - Functionalities
+   * - max_len
+     - | maximum number of text tokens to generate
+       | default: 20 (for image captioning)
+   * - min_len
+     - | minimum number of text tokens to generate
+       | default: 5 (for image captioning)
+   * - num_beams
+     - | number of beams to perform beam search
+       | default: 3
+
+.. list-table::
+   :widths: 40 50
+   :header-rows: 1
+
+   * - Multimodal Retrieval Configurations
+     - Functionalities
+   * - negative_all_rank
+     - | collect negatives from all processes for the image-text matching loss
+       | default: True (for coco)
+   * - k_test
+     - | number of retrieval candidates ranked from contrastive similarity
+       | default: 256 (for coco)

docs/tutorial.datasets.rst

@@ -0,0 +1,424 @@
+Adding Datasets
+################################################
+
+This is a tutorial on adding a new dataset using ``lavis.datasets`` module. 
+
+The LAVIS library includes a standard dataset module, which allows customization to add new datasets. 
+The ``lavis.datasets`` module is designed such that any new dataset class can be easily added and adapted from our code base, including creating dataset configuration, and defining and associating new dataset classes.
+
+In this tutorial, we will replicate the steps to add a dataset class for the `Audio-Visual Scene-Aware Dialogue (AVSD) <https://arxiv.org/pdf/1901.09107.pdf>`_ benchmark for the video-grounded dialogue task.
+
+Dataset Configuration ``lavis.configs.datasets``
+**************************************************************
+
+First, we define the basic configurations for this dataset, including a new dataset class ``avsd_dialogue``, dataset card, and data types. 
+We can define any new dataset configuration in ``lavis.configs.datasets``. For instance, under this module, we can set up a configuration file ``avsd/defaults_dial.yaml`` as follows:  
+
+.. code-block:: yaml
+
+    datasets:
+      avsd_dialogue: # name of the dataset builder
+        dataset_card: dataset_card/avsd_dialogue.md # path to the dataset card 
+        data_type: features # [images|videos|features] we use features in this case for extracted video features 
+
+        build_info:
+          # Be careful not to append minus sign (-) before split to avoid itemizing
+          annotations:
+            train:
+              url: /export/home/data/avsd/train_set4DSTC7-AVSD.json
+              storage: avsd/annotations/train.json
+            val:
+              url: /export/home/data/avsd/valid_set4DSTC7-AVSD.json
+              storage: avsd/annotations/val.json 
+            test:
+              url: /export/home/data/avsd/test_set4DSTC7-AVSD.json
+              storage: avsd/annotations/test.json 
+          features:
+            storage: /export/home/data/avsd/features/ 
+
+
+Dataset Card
+===============
+One optional step to set up dataset configuration is defining a dataset card, which contains more details about the dataset such as description, tasks, and metrics. 
+For instance, we can define a dataset card for the AVSD benchmark in ``dataset_card/avsd_dialogue.md``.
+Depending on the dataset, we included in its corresponding dataset card the command for auto-downloading data (with python code defined in ``lavis.datasets.download_scripts``) that will automatically load the data and store it in a specific folder.
+Else, you should describe in the dataset card the external download instructions from the original data source to load the dataset properly. 
+
+One example of a dataset card for the AVSD benchmark is: 
+
+.. code-block:: md
+
+    ![Samples from the AVSD dataset (Image credit: "https://arxiv.org/pdf/1901.09107.pdf").](imgs/avsd_dialogue.png)(Samples from the AVSD dataset. Image credit: "https://arxiv.org/pdf/1901.09107.pdf")
+    
+    # Audio-Visual Scene-Aware Dialogues (AVSD) 
+    
+    ## Description
+    [Audio-Visual Scene-Aware Dialogues (AVSD)](https://github.com/hudaAlamri/DSTC7-Audio-Visual-Scene-Aware-Dialog-AVSD-Challenge) contains more than 10,000 dialogues, each of which is grounded on a unique video. In the test split, for each test sample, 6 reference dialogue responses are provided. 
+    
+    
+    ## Task
+    
+    (https://github.com/hudaAlamri/DSTC7-Audio-Visual-Scene-Aware-Dialog-AVSD-Challenge)
+    
+    In a **video-grounded dialogue task**, the system must generate responses to user input in the context of a given dialog.
+    This context consists of a dialog history (previous utterances by both user and system) in addition to video and audio information that comprise the scene. The quality of a system’s automatically generated sentences is evaluated using objective measures to determine whether or not the generated responses are natural and informative
+    
+    ## Metrics
+    Models are typically evaluated according to [BLEU](https://aclanthology.org/P02-1040/), [CIDER](https://www.cv-foundation.org/openaccess/content_cvpr_2015/papers/Vedantam_CIDEr_Consensus-Based_Image_2015_CVPR_paper.pdf), [METEOR](https://aclanthology.org/W05-0909/), and [ROUGE-L](https://aclanthology.org/W04-1013/) metrics. 
+    
+    ## Leaderboard
+    
+    ....
+    
+    
+    ## Auto-Downloading
+    
+    Please refer to [benchmark webite](https://github.com/hudaAlamri/DSTC7-Audio-Visual-Scene-Aware-Dialog-AVSD-Challenge) for instructions to download the dataset. 
+    
+    
+    ## References
+    "Audio Visual Scene-Aware Dialog", Huda Alamri, Vincent Cartillier, Abhishek Das, Jue Wang, Anoop Cherian, Irfan Essa, Dhruv Batra, Tim K. Marks, Chiori Hori, Peter Anderson, Stefan Lee, Devi Parikh
+
+Visual Data Type
+==============================
+We currently limit the visual data types to one of three options: ``images``, ``videos``, and ``features``. 
+"Images" and "videos" refer to the raw visual data, which is appropriate for models processing visual data in their original forms (e.g. ViT models). 
+"Features" are visual representations extracted from pretrained models (e.g. CNN models). 
+In this tutorial, the AVSD benchmark consists of video features extracted from 3D-CNN models. 
+
+Build Info
+==============================
+Build info refers to the specific locations where data is stored and cached. 
+
+For text annotations (e.g. captioning or dialogues), by default, we include three data splits, namely "train", "val", and "test", typically used in all machine learning projects. 
+For each split, we specify 2 parameters: ``url``  and ``storage``.
+``url`` can be either an online URL where the dataset can be loaded automatically (e.g. from *googleapis*), or a local directory where data is already downloaded beforehand. 
+``storage`` is the directory where the data will be cached over time, avoiding downloading data repeatedly.
+
+For visual data annotations, ensure the field name matches the data types defined earlier (e.g. one of "images", "videos" or features"). 
+As visual features are usually large and should be downloaded beforehand, we maintain only a ``storage`` parameter where visual data is cached. 
+
+Dataset ``lavis.datasets.datasets``
+**************************************************************
+
+Base Dataset ``lavis.datasets.datasets.base_dataset``
+=======================================================
+In this step, we want to define new dataset classes that inherit our base dataset class ``lavis.datasets.datasets.base_dataset``. This base dataset class already defines standard methods such as ``collater`` which uses the default collator from Pytorch. 
+
+.. code-block:: python
+
+    import json
+    from typing import Iterable
+    
+    from torch.utils.data import Dataset, ConcatDataset
+    from torch.utils.data.dataloader import default_collate
+        
+    class BaseDataset(Dataset):
+        def __init__(
+            self, vis_processor=None, text_processor=None, vis_root=None, ann_paths=[]
+        ):
+            """
+            vis_root (string): Root directory of images (e.g. coco/images/)
+            ann_root (string): directory to store the annotation file
+            """
+            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
+            self.text_processor = text_processor
+    
+            self._add_instance_ids()
+    
+        def __len__(self):
+            return len(self.annotation)
+    
+        def collater(self, samples):
+            return default_collate(samples)
+    
+        def set_processors(self, vis_processor, text_processor):
+            self.vis_processor = vis_processor
+            self.text_processor = text_processor
+    
+        def _add_instance_ids(self, key="instance_id"):
+            for idx, ann in enumerate(self.annotation):
+                ann[key] = str(idx)
+
+Any dataset subclass will inherit these methods and it is optional to define and overwrite these methods accordingly to the specifications of the dataset. 
+We encourage users not to modify the base dataset class as any modification will have cascading impacts on any other dataset classes that inherit this base dataset. 
+Instead, the users should independently create new dataset classes to cater to their specific requirements. 
+
+Dialogue Datasets ``lavis.datasets.datasets.dialogue_datasets``
+======================================================================
+
+For example, for the AVSD dataset, we want to define a new dataset subclass ``DialogueDataset`` for dialogue tasks. We can define this dataset class in ``lavis.datasets.datasets.dialogue_datasets`` as following: 
+
+.. code-block:: python
+
+    import os
+    from collections import OrderedDict
+        
+    from lavis.datasets.datasets.base_dataset import BaseDataset
+    
+    import json 
+    import copy 
+
+    class DialogueDataset(BaseDataset):
+        def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
+            """
+            vis_processor (string): visual processor 
+            text_processor (string): textual processor 
+            vis_root (string): Root directory of images (e.g. coco/images/)
+            ann_paths (string): Root directory of images (e.g. coco/images/)
+            """
+                
+            self.vis_root = vis_root
+    
+            self.annotation = []
+            for ann_path in ann_paths:
+                dialogs = json.load(open(ann_path, "r"))['dialogs']
+                for dialog in dialogs: 
+                    all_turns = dialog['dialog']
+                    dialogue_context = [] 
+                    for turn in all_turns: 
+                        dialog_instance = copy.deepcopy(dialog)
+                        question = turn['question']
+                        answer = turn['answer'] 
+                        
+                        dialog_instance['dialog'] = copy.deepcopy(dialogue_context) 
+                        dialog_instance['question'] = question
+                        dialog_instance['answer'] = answer 
+                        self.annotation.append(dialog_instance)
+                        dialogue_context.append(turn)
+                        
+            self.vis_processor = vis_processor
+            self.text_processor = text_processor
+    
+            self._add_instance_ids()
+    
+            self.img_ids = {}
+            n = 0
+            for ann in self.annotation:
+                img_id = ann["image_id"]
+                if img_id not in self.img_ids.keys():
+                    self.img_ids[img_id] = n
+                    n += 1
+
+Class inheritance allows us to define multiple subclasses. For instance, we want another dialogue dataset class that is defined only for the test split. We can define another dataset class ``DialogueEvalDataset`` as similarly defined above but the annotations are processed differently. 
+Typically, in dialogue tasks, during test time, only a single test sample is constructed per dialogue (rather than decomposing all dialogue turns as samples during training time).
+The dataset class can then be defined as: 
+
+.. code-block:: python
+
+    class DialogueEvalDataset(BaseDataset):
+        def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
+            # ...
+            # defined similarly as DialogueDataset above 
+            # except for the loading of dialogue annotation data            
+    
+            self.annotation = []
+            for ann_path in ann_paths:
+                dialogs = json.load(open(ann_path, "r"))['dialogs']
+                for dialog in dialogs: 
+                    all_turns = dialog['dialog']
+                    dialogue_context = all_turns[:-1]
+                    last_turn = all_turns[-1] 
+                    
+                    question = last_turn['question']
+                    answer = last_turn['answer'] 
+                        
+                    dialog['dialog'] = dialogue_context
+                    dialog['question'] = question
+                    dialog['answer'] = answer
+                                        
+                    self.annotation.append(dialog)
+
+
+Using class inheritance to define datasets also allows us to develop more fine-grain class implementations, each of which is specifically designated for a benchmark. 
+For instance, under the dialogue-based tasks, we can further define another dataset subclass that is specified for the AVSD dataset. 
+We can define a new class ``AVSDDialDataset`` that further specifies how to load individual samples and collate them accordingly to specific requirements: 
+
+.. code-block:: python
+
+    import os
+    from lavis.datasets.datasets.base_dataset import BaseDataset
+    from lavis.datasets.datasets.dialogue_datasets import DialogueDataset, DialogueEvalDataset
+    
+    import torch 
+        
+    class AVSDDialDataset(DialogueDataset):
+        def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
+
+            super().__init__(vis_processor, text_processor, vis_root, ann_paths)
+    
+        def __getitem__(self, index):
+    
+            ann = self.annotation[index]
+    
+            vname = ann["image_id"]
+    
+            video = self.vis_processor(self.vis_root, vname)
+            
+            dialogue = self.text_processor(ann)
+            
+            return {
+                "video_fts": video['video_fts'],
+                "video_token_type_ids": video['token_type_ids'], 
+                "input_ids": dialogue['input_ids'], 
+                "token_type_ids": dialogue['token_type_ids'],
+                "labels": dialogue['labels'], 
+                "image_id": ann["image_id"],
+                "instance_id": ann["instance_id"]
+            }
+        
+        def collater(self, samples):
+            
+            input_ids, token_type_ids, labels, video_fts, video_token_type_ids = [], [], [], [], []
+            
+            for i in samples:
+                input_ids.append(i['input_ids'])
+                token_type_ids.append(i['token_type_ids'])
+                labels.append(i['labels'])
+                video_fts.append(i['video_fts'])
+                video_token_type_ids.append(i['video_token_type_ids'])
+    
+            input_ids = self.text_processor.padding(input_ids)
+            
+            labels = self.text_processor.padding(labels, -1)
+            video_fts = self.vis_processor.padding(video_fts)
+            
+            token_type_ids = self.text_processor.padding(token_type_ids)
+            video_token_type_ids = self.text_processor.padding(video_token_type_ids)
+            token_type_ids = torch.cat([video_token_type_ids, token_type_ids], dim=1)
+            
+            attn_mask = self.text_processor.get_attention_mask(input_ids)
+            video_mask = self.vis_processor.get_attention_mask(video_fts)
+            attn_mask = torch.cat([video_mask, attn_mask], dim=1)
+            
+            video_labels = torch.ones((video_fts.size(0), video_fts.size(1))).long() * -1 # ignore token indice -1 by default 
+
+            labels = torch.cat([video_labels, labels], dim=1)
+            
+            samples = {}
+            samples['input_ids'] = input_ids
+            samples['token_type_ids'] = token_type_ids
+            samples['labels'] = labels
+            samples['video_fts'] = video_fts
+            samples['attn_mask'] = attn_mask
+            
+            return samples  
+
+Note that in a dataset subclass, if methods such as ``__getitem__`` and ``collater`` are not defined, the same functions from the corresponding superclass will be used. 
+For instance, by default, we always use the collater from the ``BaseDataset`` class to collate data samples. 
+
+Dataset Builder ``lavis.datasets.builders``
+**************************************************************
+Dataset Builder is the data processing module that controls the dataset classes (by training or evaluation split) and associates the specific dataset configurations to these dataset classes. 
+
+Base Dataset Builder ``lavis.datasets.builders.base_dataset_builder``
+======================================================================
+
+Note that any new builder class definition should inherit the base dataset builder class ``lavis.datasets.builders.base_dataset_builder``:
+
+.. code-block:: python
+
+    class BaseDatasetBuilder:
+        train_dataset_cls, eval_dataset_cls = None, None
+        ...
+
+This allows us to standardize the operations of dataset builders across all builder classes. We advise the users to carefully review the standard methods defined in the base builder class, including methods such as ``_download_data`` and ``build_dataset`` that will load download the data and create instances of dataset classes: 
+
+.. code-block:: python
+
+    class BaseDatasetBuilder:
+    ...
+
+        def build_datasets(self):
+            # download, split, etc...
+            # only called on 1 GPU/TPU in distributed
+    
+            if is_main_process():
+                self._download_data()
+    
+            if is_dist_avail_and_initialized():
+                dist.barrier()
+    
+            # at this point, all the annotations and image/videos should be all downloaded to the specified locations.
+            logging.info("Building datasets...")
+            datasets = self.build()  # dataset['train'/'val'/'test']
+            
+            return datasets
+    
+        def _download_data(self):
+            self._download_ann()
+            self._download_vis()
+    
+We encourage users not to modify the implementation of the base dataset builder class as this will affect all existing dataset builder subclasses.
+
+Dialogue Dataset Builder ``lavis.datasets.builders.dialogue_builder``
+======================================================================
+We can define any new builder subclass and associate this builder with the corresponding dataset classes and dataset configurations. 
+For instance, for the AVSD dataset, we can define a builder ``lavis.datasets.builders.dialogue_builder`` for dialogue-based datasets as follows: 
+
+.. code-block:: python
+
+    from lavis.datasets.builders.base_dataset_builder import BaseDatasetBuilder
+    from lavis.datasets.datasets.avsd_dialogue_datasets import (
+        AVSDDialDataset, 
+        AVSDDialEvalDataset 
+    )
+    
+    from lavis.common.registry import registry
+    
+    
+    @registry.register_builder("avsd_dialogue")
+    class AVSDDialBuilder(BaseDatasetBuilder):
+        train_dataset_cls = AVSDDialDataset 
+        eval_dataset_cls = AVSDDialEvalDataset 
+    
+        DATASET_CONFIG_DICT = {
+            "default": "configs/datasets/avsd/defaults_dial.yaml"
+        }
+
+Note that we chose to separately define the parameters ``train_dataset_cls`` and  ``eval_dataset_cls`` to consider cases where data is processed differently between training and test time. 
+For instance, in captioning tasks, during test time, each data sample often includes multiple ground-truth captions rather than just a single ground-truth during training time. 
+If the data processing is the same in both training and test time, the two parameters can be linked to the same dataset class. 
+
+Finally, define ``DATASET_CONFIG_DICT`` to associate the dataset configurations to the assigned dataset classes. 
+
+Registering Builder ``lavis.datasets.builders.__init__``
+======================================================================
+
+To add a new builder class, ensure to first include the class within the ``__init__.py``. For instance, to define a new builder for the AVSD dataset: 
+
+.. code-block:: python
+
+    from lavis.datasets.builders.dialogue_builder import (
+        AVSDDialBuilder
+    )
+    
+    __all__ = [
+        ...,
+        "AVSDDialBuilder"
+    ]
+
+Assigning Builder 
+======================================================================
+Note that during data loading and processing, the builder being assigned must have the correct registry to be able to load it properly. 
+For instance, the following should be specified in a configuration file e.g. ``dialogue_avsd_ft.yaml``: 
+
+.. code-block:: yaml
+
+    datasets:
+      avsd_dialogue: # name of the dataset builder
+        ...
+        # processor configuration 
+        ...
+
+Subsequently, any processes (e.g. training) should load this configuration file to assign the correct builder which will then associate the correct dataset classes to construct data samples. 
+
+.. code-block:: sh
+
+    python train.py --cfg-path dialogue_avsd_ft.yaml

docs/tutorial.evaluation.rst

@@ -0,0 +1,40 @@
+Evaluating Pre-trained Models on Task Datasets
+###############################################
+LAVIS provides pre-trained and finetuned model for off-the-shelf evaluation on task dataset. 
+Let's now see an example to evaluate BLIP model on the captioning task, using MSCOCO dataset.
+
+.. _prep coco:
+
+Preparing Datasets
+******************
+First, let's download the dataset. LAVIS provides `automatic downloading scripts` to help prepare 
+most of the public dataset, to download MSCOCO dataset, simply run
+
+.. code-block:: bash
+
+    cd lavis/datasets/download_scripts && python download_coco.py
+
+This will put the downloaded dataset at a default cache location ``cache`` used by LAVIS.
+
+If you want to use a different cache location, you can specify it by updating ``cache_root`` in ``lavis/configs/default.yaml``.
+
+If you have a local copy of the dataset, it is recommended to create a symlink from the cache location to the local copy, e.g.
+
+.. code-block:: bash
+
+    ln -s /path/to/local/coco cache/coco
+
+Evaluating pre-trained models
+******************************
+
+To evaluate pre-trained model, simply run
+
+.. code-block:: bash
+
+    bash run_scripts/blip/eval/eval_coco_cap.sh
+
+Or to evaluate a large model:
+
+.. code-block:: bash
+
+    bash run_scripts/blip/eval/eval_coco_cap_large.sh

docs/tutorial.models.rst

@@ -0,0 +1,245 @@
+Adding Models
+####################################
+
+This is a tutorial on adding new models using ``lavis.models`` module.
+
+The LAVIS library includes a standard model module that builds the foundation for many major language-vision models such as `ALBEF <https://arxiv.org/pdf/2107.07651.pdf>`_,
+`BLIP <https://arxiv.org/pdf/2201.12086.pdf>`_, `ALPRO <https://arxiv.org/pdf/2112.09583.pdf>`_, and `CLIP <https://arxiv.org/pdf/2103.00020.pdf>`_. 
+The ``lavis.models`` module is designed such that any new models can be added and integrated into the LAVIS library, with minimal steps to develop training and testing procedures. 
+In this tutorial, we will replicate the steps to add a GPT-style model specifically for `video-grounded dialogue tasks <https://arxiv.org/pdf/1901.09107.pdf>`_. 
+
+Base Model ``lavis.models.base_model``
+**************************************************************
+
+Note that any new model definition should inherit the base model class ``BaseModel``:
+
+.. code-block:: python
+
+    from omegaconf import OmegaConf
+    
+    import numpy as np
+    
+    import torch
+    import torch.nn as nn
+    
+    from lavis.common.utils import get_abs_path
+    
+    class BaseModel(nn.Module):
+        """Base class for models."""
+    
+        def __init__(self):
+            super().__init__()
+    
+        def forward_features(self, *args, **kwargs):
+            """Similar to *forward* but only return features."""
+            raise NotImplementedError
+    
+        def load_from_pretrained(self, url_or_filename):
+            raise NotImplementedError
+    
+        @classmethod
+        def _from_config(cls, cfg=None, model_type="base"):
+            if not cfg:
+                # useful when building model without a provided configuration file
+                cfg = OmegaConf.load(cls.default_config_path(model_type)).model
+    
+            return cls.from_config(cfg)
+    
+        @classmethod
+        def from_pretrained(cls, model_type="base"):
+            """
+            Build a pretrained model from the default configuration file, specified by model_type.
+            """
+            return cls._from_config(cfg=None, model_type=model_type)
+    
+        @property
+        def device(self):
+            return list(self.parameters())[0].device
+    
+        @classmethod
+        def default_config_path(cls, model_type="base"):
+            assert (
+                model_type in cls.PRETRAINED_MODEL_CONFIG_DICT
+            ), "Unknown model type {}".format(model_type)
+            return get_abs_path(cls.PRETRAINED_MODEL_CONFIG_DICT[model_type])
+    
+        def before_evaluation(self, **kwargs):
+            pass
+    
+        def show_n_params(self, return_str=True):
+            tot = 0
+            for p in self.parameters():
+                w = 1
+                for x in p.shape:
+                    w *= x
+                tot += w
+            if return_str:
+                if tot >= 1e6:
+                    return "{:.1f}M".format(tot / 1e6)
+                else:
+                    return "{:.1f}K".format(tot / 1e3)
+            else:
+                return tot
+
+
+In this base model, we already declare and standardize many common methods such as ``_from_config`` and ``_from_pretrained``. 
+Inheriting this base model class allows us to standardize operations of models across all model classes while still allowing customizations. 
+We advise users not to change the implementation of the base model class as this will affect all existing model subclasses.
+
+GPT-style Video-grounded Dialogue Model ``lavis.models.gpt_models.gpt_dialogue``
+********************************************************************************
+
+In this step, we can define a new model class, e.g. under ``lavis.models.gpt_models.gpt_dialogue``, for GPT-based dialogue models designed specifically for video-grounded dialogues. 
+Note that we assume the model class inherits from the standard model super class ``GPT2LMHeadModel`` from the ``transformers`` `library <https://huggingface.co/docs/transformers/index>`_.
+We also enforce model integration to the LAVIS framework through the inheritance of the ``BaseModel`` from the LAVIS library, as the secondary super class.
+
+.. code-block:: python
+
+    import torch
+    from lavis.common.registry import registry
+    from lavis.models.base_model import BaseModel
+    
+    from transformers import GPT2Model, GPT2LMHeadModel
+    from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+    import math
+    import torch
+    import torch.nn as nn
+    from torch.nn import CrossEntropyLoss, MSELoss
+        
+    @registry.register_model("gpt_dialogue")
+    class GPTDialogue(GPT2LMHeadModel, BaseModel):
+        ...
+ 
+Next, we can modify the architecture of the model during model initialization to fit the tasks of interest, i.e. video-grounded dialogues. 
+In this case, we want to add additional model parameters for a linear network to transform the video feature representations to the model dimension. 
+
+.. code-block:: python
+
+    class GPTDialogue(GPT2LMHeadModel, BaseModel):
+
+        def __init__(self, config, len_video_ft=4224):
+            
+            super().__init__(config)
+            
+            self.video_ff = nn.Linear(len_video_ft, config.n_embd)
+       
+            # Model parallel
+            self.model_parallel = False
+            self.device_map = None
+    
+            # Initialize weights and apply final processing
+            self.post_init()
+    
+Note that for each new model class, we advise redefining the ``from_config`` method which is inherited from the ``BaseModel`` class.
+As each model usually has its own unique configurations, redefining the method will ensure the model instances are created properly. 
+For instance, ``GPTDialogue`` requires an additional parameter of video feature length (``len_video_ft``) which should be part of the model initialization procedure. 
+Another additional parameter is the number of tokens/words (as we include additional special tokens in the vocabulary for dialogue tasks). 
+
+.. code-block:: python
+
+    class GPTDialogue(GPT2LMHeadModel, BaseModel):
+        ...
+        @classmethod
+        def from_config(cls, cfg):
+            model = cls.from_pretrained('gpt2', len_video_ft=cfg['len_video_ft']) 
+            model.resize_token_embeddings(cfg['len_tokenizer'])
+            return model
+
+Other basic methods should also be defined explicitly in the new model class, including the ``forward`` function. 
+For instance, in GPT models for video-grounded dialogue tasks, we want the forward operation also includes the transformation and integration of video features before passing the representations to the Transformer layers. 
+
+.. code-block:: python
+
+    class GPTDialogue(GPT2LMHeadModel, BaseModel):
+        ...
+
+        def forward(self, samples, 
+                    past_key_values=None,
+                    position_ids=None,
+                    head_mask=None,
+                    encoder_hidden_states=None,
+                    encoder_attention_mask=None,
+                    use_cache=None,
+                    output_attentions=None,
+                    output_hidden_states=None,
+                    return_dict=None):        
+                
+                input_embs = self.transformer.wte(samples['input_ids'])
+                video_embs = self.video_ff(samples['video_fts'])
+                input_embs = torch.cat([video_embs, input_embs], dim=1)
+                        
+                transformer_outputs = self.transformer(
+                    attention_mask=samples['attn_mask'],
+                    token_type_ids=samples['token_type_ids'],
+                    inputs_embeds=input_embs,
+                    position_ids=position_ids,
+                    head_mask=head_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    return_dict=return_dict,
+                )
+                hidden_states = transformer_outputs[0]
+            
+                lm_logits = self.lm_head(hidden_states)
+                ...
+
+Registering New Model ``lavis.models.__init__``
+********************************************************************************
+
+Any new model must be officially registered as part of the ``lavis.models`` module. 
+For instance, to add a model class for GPT-based dialogue models, we can modify the ``__init__.py`` as follows:
+
+.. code-block:: python
+
+    from lavis.models.gpt_models.gpt_dialogue import GPTDialogue
+    
+    __all__ = [
+        ...
+        "GPTDialogue"
+    ]
+
+Assigning Model
+********************************************************************************
+
+From the above example of a model class, note that we define a ``from_config method`` for the new model class. 
+This method will process a configuration file and pass specific parameters to initialize the model classes properly. 
+To do this, we can assign/ associate the correct registry of model classes in a configuration file. 
+For instance, the following should be specified in a configuration file e.g. ``dialogue_avsd_ft.yaml``:
+
+.. code-block:: yaml
+
+    model:
+      arch: gpt_dialogue # name of the model 
+      model_type: base
+
+
+Subsequently, any processes (e.g. training) should load this configuration file to assign the correct model.
+
+.. code-block:: sh
+
+    python train.py --cfg-path dialogue_avsd_ft.yaml
+
+Note that to simplify the model configuration, we only enable two main parameters here: ``arch`` and ``model_type``. ``arch`` refers to the model class registry, and ``model_type`` is the corresponding model type under this model family.
+For instance, with ``gpt_dialogue``, we have a model ``base`` which has its own configuration in a separate configuration file e.g. ``gpt_dialogue_base.yaml``:
+
+.. code-block:: yaml
+
+    model:
+      arch: gpt_dialogue
+      len_tokenizer: 50264 # 50257 tokens from gpt2 default tokenizer + additional special tokens       
+      len_video_ft: 4224 # i3d_rgb: 2048 i3d_flow: 2048 vggish: 128 
+
+We can pass load this configuration and pass the parameters to the above ``from_config`` method to initialize the model accordingly. 
+We advise the users to maintain a dictionary that contains default paths to model configurations, in the model class definition. 
+By default, the LAVIS framework will search for configurations from each model class defined as ``model.PRETRAINED_MODEL_CONFIG_DICT``.
+
+.. code-block:: python
+
+    class GPTDialogue(GPT2LMHeadModel, BaseModel):
+        PRETRAINED_MODEL_CONFIG_DICT = {
+                "base": "configs/models/gpt_dialogue_base.yaml"
+            }
+        ...

docs/tutorial.processors.rst

@@ -0,0 +1,233 @@
+Adding Processors
+################################################
+
+This is a tutorial on adding new processors using ``lavis.processors`` module. 
+
+The LAVIS library includes a standard processor module that preprocesses data e.g. image transformation and sequence concatenation.
+The ``lavis.processors`` module is designed such that any processors can be added, specifically to the requirements of corresponding models of interest. 
+In this tutorial, we will replicate the steps to add visual and textual processors specifically for `video-grounded dialogue tasks <https://arxiv.org/pdf/1901.09107.pdf>`_. 
+In addition, we also want the processors to have processing features to make the data samples compatible with GPT-style models.
+
+Base Processor ``lavis.processors.base_processors``
+*****************************************************
+
+Note that any new processor definition should inherit the base processor class ``BaseProcessor``:
+
+.. code-block:: python
+
+    from omegaconf import OmegaConf
+    
+    class BaseProcessor:
+        def __init__(self):
+            self.transform = lambda x: x
+            return
+    
+        def __call__(self, item):
+            return self.transform(item)
+    
+        @classmethod
+        def from_config(cls, cfg=None):
+            return cls()
+    
+        def build(self, **kwargs):
+            cfg = OmegaConf.create(kwargs)
+    
+            return self.from_config(cfg)
+
+This allows us to standardize operations of processors across all processor classes while still allowing customization of processors specifically to data and model types. 
+We encourage users not to modify the implementation of the base processor class as this will have an impact on all existing processor subclasses.
+
+GPT-style Processors ``lavis.processors.gpt_processors``
+**************************************************************
+In this step, we can define new processor classes, e.g. under ``lavis.processors.gpt_processors``, for GPT models designed specifically for video-grounded dialogues. 
+First, we want to process video features by defining ``GPTVideoFeatureProcessor`` class.
+In this tutorial, we assume video features are extracted beforehand and this processor simply loads the features from ``npy`` files.
+Other methods that are specifically defined are ``padding`` (which is used by dataset instances to pad multiple video samples) and ``get_attention_mask`` (which creates an attention mask for Transformer attention in GPT models). 
+
+.. code-block:: python 
+
+    SPECIAL_TOKENS_DICT = {'bos_token': "<bos>", 'eos_token': "<eos>", 'additional_special_tokens': ["<speaker1>", "<speaker2>", "<video>", "<cap>"], 'pad_token': "<pad>"}
+    ...
+
+    @registry.register_processor("gpt_video_ft")
+    class GPTVideoFeatureProcessor(BaseProcessor):
+        def __init__(self, visual_ft, audio_ft):
+
+            self.visual_ft = visual_ft
+            self.audio_ft = audio_ft
+
+            self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+            self.tokenizer.add_special_tokens(SPECIAL_TOKENS_DICT) 
+                    
+        def padding(self, seq):
+            padded_seq = torch.nn.utils.rnn.pad_sequence(seq, batch_first=True, padding_value=1.0) 
+            return padded_seq
+        
+        def get_attention_mask(self, seq):
+            return torch.sum(seq != 1, dim=2) != 0
+    
+        def __call__(self, ft_root, vname):
+            all_ft = []
+            
+            for ft_name in self.visual_ft:
+                ft_path = os.path.join(ft_root, ft_name, vname)
+                all_ft.append(np.load(ft_path + '.npy'))
+            
+            for ft_name in self.audio_ft: 
+                ft_path = os.path.join(ft_root, ft_name, vname)
+                all_ft.append(np.load(ft_path + '.npy'))
+            
+            min_len = min([len(ft) for ft in all_ft])
+            
+            sampled_ft = [ft[:min_len] for ft in all_ft]
+            sampled_ft = np.concatenate(sampled_ft, axis=1)
+            item = {} 
+            item['video_fts'] = torch.Tensor(sampled_ft) 
+            
+            video_type_token = self.tokenizer.convert_tokens_to_ids('<video>')
+            item['token_type_ids'] = torch.Tensor([video_type_token] * len(sampled_ft)).long() 
+            
+            return item 
+    
+        @classmethod
+        def from_config(cls, cfg=None):
+            if cfg is None:
+                cfg = OmegaConf.create()
+            
+            visual_ft = cfg.get("visual_ft", ["i3d_rgb"])
+            audio_ft = cfg.get("audio_ft", ["vggish"])
+            
+            return cls(
+                visual_ft=visual_ft,
+                audio_ft=audio_ft
+            )
+
+Another processor class that will be useful to have is to process dialogue data. Here we can define a ``GPTDialogueProcessor`` class.
+This processor class receives raw annotations and constructs inputs as a concatenation of input sequences (questions, dialogue contexts, and responses) to facilitate application in GPT models. 
+Other methods that are specifically defined are ``padding`` (which is used by dataset instances to pad multiple sequence samples) and ``get_attention_mask`` (which creates an attention mask for Transformer attention in GPT models). 
+
+.. code-block:: python 
+
+    SPECIAL_TOKENS_DICT = {'bos_token': "<bos>", 'eos_token': "<eos>", 'additional_special_tokens': ["<speaker1>", "<speaker2>", "<video>", "<cap>"], 'pad_token': "<pad>"}
+    ...
+
+    @registry.register_processor("gpt_dialogue")
+    class GPTDialogueProcessor(BaseProcessor):
+        def __init__(self, max_turns=3, use_caption=True):
+            self.max_turns = max_turns 
+            self.use_caption = use_caption 
+            self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+            self.tokenizer.add_special_tokens(SPECIAL_TOKENS_DICT) 
+            
+        def sample_sequence(self, caption, history, answer):
+            bos, eos, speaker1, speaker2, cap = self.tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS[:-2])
+            instance = {}
+            sequence = [caption] + history + [answer]
+            sequence = [s + [eos] for s in sequence] 
+    
+            instance["input_ids"] = list(chain(*sequence))
+            instance["token_type_ids"] = [cap] * len(sequence[0]) + [speaker2 if i % 2 else speaker1 for i, s in enumerate(sequence[1:]) for _ in s]
+            instance["labels"] = ([-1]*sum(len(s) for s in sequence[:-1])) + sequence[-1]
+            
+            assert len(instance["input_ids"])==len(instance["token_type_ids"])
+            assert len(instance["token_type_ids"])==len(instance["labels"])
+            
+            for k,v in instance.items():
+                instance[k] = torch.Tensor(v).long() 
+            
+            return instance 
+        
+        def padding(self, seq, pad_token=-1):
+            if pad_token==-1: pad_token = self.tokenizer.pad_token_id 
+            padded_seq = torch.nn.utils.rnn.pad_sequence(seq, batch_first=True, padding_value=pad_token) 
+            return padded_seq
+        
+        def get_attention_mask(self, seq, pad_token=-1):
+            if pad_token==-1: pad_token = self.tokenizer.pad_token_id 
+            return seq != pad_token
+        
+        def __call__(self, ann):
+            if self.use_caption:
+                caption = ' '.join([ann['caption'], ann['summary']])
+                caption = self.tokenizer.encode(caption)
+            else:
+                caption = []
+                
+            dial_history = []
+            for turn in ann['dialog'][-self.max_turns:]:
+                dial_history.append(turn['question'])
+                dial_history.append(turn['answer'])
+            dial_history.append(ann['question'])
+            dial_history = [self.tokenizer.encode(t) for t in dial_history]
+            
+            answer = self.tokenizer.encode(ann['answer'])
+            
+            item = self.sample_sequence(caption, dial_history, answer)
+            
+            return item 
+    
+        @classmethod
+        def from_config(cls, cfg=None):
+            if cfg is None:
+                cfg = OmegaConf.create()
+    
+            use_caption = cfg.get("use_caption", True)
+            max_turns = cfg.get("max_turns", 3)
+    
+            return cls(max_turns=max_turns, use_caption=use_caption)
+
+Registering New Processors ``lavis.processors.__init__``
+**************************************************************
+
+Finally, any new processor must be officially registered as part of the ``lavis.processors`` module. 
+For instance, to add processor classes for GPT-based dialogue models, including one for dialogue data ``GPTDialogueProcessor`` and one for video features ``GPTVideoFeatureProcessor``, we can modify the ``__init__.py`` as follows: 
+
+.. code-block:: python
+
+    from lavis.processors.gpt_processors import (
+        GPTVideoFeatureProcessor,
+        GPTDialogueProcessor,
+    )
+    
+    __all__ = [
+        ...
+        # GPT
+        "GPTVideoFeatureProcessor",
+        "GPTDialogueProcessor"
+    ]
+
+Assigning Processors 
+**************************************************************
+From the above example of processor classes, note that we define a ``from_config`` method for each class. 
+This method will process a configuration file and pass specific parameters e.g. ``max_turns``, ``visual_ft``, to initialize the processor classes properly. 
+To do this, we can assign/ associate the correct registry of processor classes in a configuration file.
+For instance, the following should be specified in a configuration file e.g. ``dialogue_avsd_ft.yaml``:
+
+.. code-block:: yaml 
+
+    datasets:
+      avsd_dialogue: # name of the dataset builder
+        vis_processor:
+            train:
+              name: "gpt_video_ft" # name of the visual processor for training data
+              visual_ft: ["i3d_flow", "i3d_rgb"]  
+              audio_ft: ["vggish"]    
+            eval:
+              name: "gpt_video_ft" # name of the visual processor for evaluation data
+              visual_ft: ["i3d_flow", "i3d_rgb"]  
+              audio_ft: ["vggish"]   
+        text_processor:
+            train:
+              name: "gpt_dialogue" # name of the textual processor for training data
+              max_turns:  3
+              use_caption: True 
+            eval:
+              name: "gpt_dialogue" # name of the textual processor for evaluation data
+              max_turns:  3
+              use_caption: True 
+
+Subsequently, any processes (e.g. training) should load this configuration file to assign the correct processors.
+
+.. code-block:: sh
+
+    python train.py --cfg-path dialogue_avsd_ft.yaml

docs/tutorial.rst

@@ -0,0 +1,13 @@
+Tutorials
+==============================
+
+.. toctree::
+   :maxdepth: 1
+
+   tutorial.evaluation
+   tutorial.training-example
+   tutorial.configs
+   tutorial.datasets
+   tutorial.processors
+   tutorial.models
+   tutorial.tasks

docs/tutorial.tasks.rst

@@ -0,0 +1,184 @@
+Adding Tasks
+####################################
+
+This is a tutorial on adding new machine learning tasks using ``lavis.tasks`` module.
+
+The LAVIS library includes a standard task module that centralizes the model training and evaluation procedure of machine learning tasks. 
+The ``lavis.tasks`` module is designed such that any new tasks can be added and integrated, catering to any customization in the training and testing procedures. 
+In this tutorial, we will replicate the steps to add a new task into LAVIS for the `video-grounded dialogue tasks <https://arxiv.org/pdf/1901.09107.pdf>`_. 
+
+Base Task ``lavis.tasks.base_task``
+********************************************************************************
+
+Note that any new model definition should inherit the base task class ``BaseTask``:
+
+.. code-block:: python
+
+    import logging
+    import os
+    
+    import torch.distributed as dist
+    from lavis.common.dist_utils import get_rank, get_world_size, is_main_process
+    from lavis.common.logger import MetricLogger, SmoothedValue
+    from lavis.common.registry import registry
+    from lavis.datasets.data_utils import prepare_sample
+    
+    class BaseTask:
+        def __init__(self, **kwargs):
+            super().__init__()
+    
+            self.inst_id_key = "instance_id"
+    
+        @classmethod
+        def setup_task(cls, **kwargs):
+            return cls()
+    
+        def build_model(self, cfg):
+            model_config = cfg.model_cfg
+    
+            model_cls = registry.get_model_class(model_config.arch)
+            return model_cls.from_config(model_config)
+    
+        def build_datasets(self, cfg):
+            """
+            Build a dictionary of datasets, keyed by split 'train', 'valid', 'test'.
+            Download dataset and annotations automatically if not exist.
+    
+            Args:
+                cfg (common.config.Config): _description_
+    
+            Returns:
+                dict: Dictionary of torch.utils.data.Dataset objects by split.
+            """
+    
+            datasets = dict()
+    
+            datasets_config = cfg.datasets_cfg
+    
+            assert len(datasets_config) > 0, "At least one dataset has to be specified."
+    
+            for name in datasets_config:
+                dataset_config = datasets_config[name]
+    
+                builder = registry.get_builder_class(name)(dataset_config)
+                dataset = builder.build_datasets()
+    
+                datasets[name] = dataset
+    
+            return datasets
+    
+        def train_step(self, model, samples):
+            loss = model(samples)["loss"]
+            return loss
+    
+        ...
+
+In this base task, we already declare and standardize many common methods such as ``train_step``, ``build_model``, and ``build_datasets``. 
+Inheriting this base task class allows us to standardize operations of tasks across all task classes.
+We recommend users not change the implementation of the base task class as this will have an impact on all existing task subclasses.
+
+Dialogue Task ``lavis.tasks.dialogue``
+********************************************************************************
+
+In this step, we can define a new task class, e.g. under ``lavis.tasks.dialogue``, for video-grounded dialogues.
+For instance, we define a new task class ``DialogueTask`` that inherits the super task class ``BaseTask``.
+
+.. code-block:: python
+
+    import json
+    import os
+    
+    from lavis.common.dist_utils import main_process
+    from lavis.common.logger import MetricLogger
+    from lavis.common.registry import registry
+    from lavis.tasks.base_task import BaseTask
+    from lavis.datasets.data_utils import prepare_sample
+    
+    import numpy as np 
+    
+    @registry.register_task("dialogue")
+    class DialogueTask(BaseTask):
+        def __init__(self, num_beams, max_len, min_len, evaluate, report_metric=True):
+            super().__init__()
+    
+            self.num_beams = num_beams
+            self.max_len = max_len
+            self.min_len = min_len
+            self.evaluate = evaluate
+    
+            self.report_metric = report_metric
+    
+        @classmethod
+        def setup_task(cls, cfg):
+            run_cfg = cfg.run_cfg
+    
+            num_beams = run_cfg.num_beams
+            max_len = run_cfg.max_len
+            min_len = run_cfg.min_len
+            evaluate = run_cfg.evaluate
+    
+            report_metric = run_cfg.get("report_metric", True)
+    
+            return cls(
+                num_beams=num_beams,
+                max_len=max_len,
+                min_len=min_len,
+                evaluate=evaluate,
+                report_metric=report_metric,
+            )
+    
+        def valid_step(self, model, samples):
+            results = []        
+            loss = model(samples)["loss"].item() 
+            
+            return [loss] 
+        ...
+
+Note that for any new task, we advise the users to review carefully the functions implemented within ``BaseTask`` and consider which methods should be modified. 
+For instance, the base task class already contains a standard implementation of model training steps that are common among machine learning steps. 
+Some major methods we want to emphasize and should be customized by each task are the ``valid_step`` and ``evaluation``. 
+These operations were not fully implemented in the base task class due to the differences in evaluation procedures among many machine learning tasks. 
+Another method that should be considered is the ``setup_task`` method. 
+This method will receive configurations that set task-specific parameters to initialize any task instance.
+
+Registering New Task ``lavis.tasks.__init__`` 
+********************************************************************************
+
+Any new task must be officially registered as part of the ``lavis.tasks`` module. For instance, to add a new task for video-grounded dialogues, we can modify the ``__init__.py`` as follows:
+
+.. code-block:: python
+
+    from lavis.tasks.dialogue import DialogueTask
+    
+    ...
+    __all__ = [
+        ...
+        "DialogueTask"
+    ]
+
+Assigning Task 
+***************
+
+From the above example of task class, note that we define a ``setup_task`` method for each task class. 
+This method will process a configuration file and pass specific parameters e.g. ``num_beams`` (for beam search generative tasks during the inference stage), to initialize the task classes properly. 
+To assign and associate any task, we need to specify the correct registry of task classes in a configuration file. 
+For instance, the following should be specified in a configuration file e.g. ``dialogue_avsd_ft.yaml``:
+
+.. code-block:: yaml
+
+    run:
+      task: dialogue # name of the task 
+      
+      # optimizer
+      ...
+    
+      max_len: 20
+      min_len: 5
+      num_beams: 3    
+      ...
+    
+Subsequently, any processes (e.g. training) should load this configuration file to assign the correct task.
+
+.. code-block:: sh
+
+    python train.py --cfg-path dialogue_avsd_ft.yaml

docs/tutorial.training-example.rst

@@ -0,0 +1,145 @@
+Example on Finetuning BLIP on COCO-Captioning
+################################################
+
+To finetune BLIP model on the coco caption dataset, first refer to :ref:`prep coco` to prepare the dataset if you have not done so.
+
+To finetune the model, we have prepared a run script for you, which can run as follows:
+
+.. code-block:: bash
+
+    bash run_scripts/blip/train/train_caption_coco_large.sh
+
+This will finetune the pre-trained BLIP large model into a new model that can be used for captioning.
+
+Deep Dive
+**********
+Now let's take a closer look at the script and see what it does.
+
+.. code-block:: bash
+
+    python -m torch.distributed.run --nproc_per_node=8 train.py --cfg-path lavis/projects/blip/train/caption_coco_large_ft.yaml
+
+As can be seen, the script simply calls the :code:`train.py` with PyTorch distributed training enabled.
+The :code:`--cfg-path` argument specifies the **runtime config** file to use. The config file is a YAML file that specifies the training parameters, shown as follows:
+
+.. literalinclude:: ../lavis/projects/blip/train/caption_coco_large_ft.yaml
+    :language: yaml
+    :linenos:
+
+The runtime config file is divided into 3 sections:
+    - :code:`model`: specifies the model architecture and type to use.
+    - :code:`data`: specifies the dataset to use.
+    - :code:`run`: specifies the runner arguments, such as tasks, optimizer, learning rate scheduler, etc.
+
+We describe each section in detail below.
+
+Model configurations
+=====================
+
+.. literalinclude:: ../lavis/projects/blip/train/caption_coco_large_ft.yaml
+    :language: yaml
+    :linenos:
+    :lines: 6-10
+
+The :code:`arch` argument specifies the model architecture to use. In this case, we use the :code:`blip_caption` architecture.
+You can find available architectures by inspecting the :code:`model_zoo`.
+Once the architecture is specified, the runner will look for the model class registered with the name and try to instantiate a model instance.
+In this case :code:`BlipCaption` is the model registered with the name :code:`blip_caption`.
+
+The registry maintains a mapping from the name string to the model class.
+This allows the runner to find the model class dynamically based on the name string from the config file. 
+The following segment in :code:`lavis/models/blip_models/blip_caption.py` shows how :code:`BlipCaption` is registered with the name string :code:`blip_caption`:
+
+.. literalinclude:: ../lavis/models/blip_models/blip_caption.py
+    :language: python
+    :linenos:
+    :lines: 20-38
+
+One same model architecture may be pre-trained or finetuned on different datasets or have different model configurations.
+For example, :code:`BlipCaption` have:
+
+    - :code:`base_coco`: pre-trained base BLIP model adapated for COCO captioning finetuning.
+
+    - :code:`large_coco`: pre-trained large BLIP model adapated for COCO captioning finetuning.
+
+Therefore, we also need to specify :code:`model_type`. Here we use :code:`large_coco`.
+And we set :code:`load_finetuned` to :code:`False` to indicate that we are finetuning the model from the pre-trained weights.
+If :code:`load_finetuned` set to :code:`True` as by default, the model will load finetuned weights on coco captioning.
+
+Given the model architecture and type, the library will then look for the default model config for :code:`large_coco` in :code:`lavis/models/blip_models/blip_caption.py`.
+As can be seen in the above code snippet, the corresponding config path is stored in :code:`BlipCaption.PRETRAINED_MODEL_CONFIG_DICT`. 
+Then the library will load :code:`lavis/configs/models/blip_caption_large_coco.yaml` as the configuration to build the model.
+
+*Priority of Configs*: Note that the priority of the run config is higher than the default model config, meaning that arguments in the run config will override the default model config.
+For example, in the default model config, :code:`load_finetuned` is set to :code:`True` by default, while in the run config, we set it to :code:`False` and finetuning from the pre-trained weights only.
+
+
+Dataset configurations
+=========================
+
+The second section of the config file specifies the dataset(s) to use.
+
+.. literalinclude:: ../lavis/projects/blip/train/caption_coco_large_ft.yaml
+    :language: yaml
+    :linenos:
+    :lines: 12-24
+
+We associate each dataset with a :code:`vis_processor` and a :code:`text_processor`, responsible for processing the visual and textual input respectively.
+Here we again use the registry mechanism to dynamically load the processor class based on the name string.
+For example, :code:`blip_image_train` is the name string for the :code:`BlipImageTrainProcessor` class, which is registered in :code:`lavis/processors/blip_processors.py`.
+
+Similarly, the dataset name string is also registered in the registry, pointing to a dataset builder :code:`COCOCapBuilder` class.
+By default, the builder will load the default dataset configuration as in :code:`DATASET_CONFIG_DICT`. You may also add new dataset types by adding new entries to the dictionary.
+
+The dataset configuration used here is:
+
+.. literalinclude:: ../lavis/configs/datasets/coco/defaults_cap.yaml
+    :language: yaml
+    :linenos:
+    :lines: 6-28
+
+In this configuration file, we specify the dataset name and mainly its building information.
+The build information is divided into two parts: :code:`annotation` and :code:`images`. The annotation files will be automatically downloaded upon loading the dataset for the first time.
+The :code:`images` part specifies the image root directory. This is a relative path to the cache directory, which is :code:`cache` by default. If you have a local copy of the dataset, you can specify the path to the local copy by
+overwriting the :code:`images` part in the runtime config file. For example, you may alter the run config as below to use your local dataset copy:
+
+.. code:: yaml
+
+    datasets:
+        coco_caption: # name of the dataset builder
+            vis_processor:
+                train:
+                name: "blip_image_train"
+                eval:
+                name: "blip_image_eval"
+            text_processor:
+                train:
+                name: "blip_caption"
+                prompt: "a picture of "
+                eval:
+                name: "blip_caption"
+            images:
+                YOUR_LOCAL_IMAGE_ROOT_DIR
+
+LAVIS supports using multiple datasets for training. See an example in :code:`lavis/projects/blip/train/pretrain_14m.yaml`.
+
+
+Runner configurations
+=========================
+The last section of the config file specifies the arguments for the runner, shown below:
+
+.. literalinclude:: ../lavis/projects/blip/train/caption_coco_large_ft.yaml
+    :language: yaml
+    :linenos:
+    :lines: 26-56
+
+Here we specify runner-related arguments, including
+    - task-specific arguments, such as :code:`task`, :code:`max_len`, :code:`min_len`, etc.
+    - learning rate schedulers, optimizer;
+    - distributed training settings;
+    - logging and checkpointing settings.
+
+Available Configurations
+#########################
+
+See :ref:`config` for the full list of available configurations and their descriptions.

examples/blip2_itm.py

@@ -0,0 +1,520 @@
+import re
+
+import torch
+from PIL import Image
+
+from lavis.models import load_model_and_preprocess
+from lavis.processors import load_processor
+from lavis.common.registry import registry
+from torch.nn import functional as F
+from lavis.models.base_model import all_gather_with_grad, concat_all_gather
+import numpy as np
+import pandas as pd
+import time
+from fuzzywuzzy import process
+from multiprocessing import Pool, Queue, Process
+import difflib
+import Levenshtein
+import os
+# import obonet
+
+
+def fuzzy_match(texts):
+    text_dict = {}
+    for context in texts:
+        if context not in choices:
+            # txt_dict[txt] = process.extractOne(txt, choices)[0]
+            text_dict[context] = difflib.get_close_matches(context, choices, n=1, cutoff=0.)[0]
+    return text_dict
+
+
+def txt_map(x, txt_dict):
+    if type(x) == str:
+        x = eval(x)
+    x_ = []
+    for i in x:
+        if i in txt_dict:
+            x_.append(txt_dict[i])
+        else:
+            x_.append(i)
+    return x_
+
+
+def levenshtein_sim(text, label):
+    all_s = []
+    for x in label:
+        s = 0
+        for y in text:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+def func(text, label):
+    all_s = []
+    for x in label:
+        s = 0
+        for y in text:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+
+def stage2_output(df_test):
+    config = {'arch': 'blip2_protein_opt', 'load_finetuned': False,
+              'pretrained': '/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage2/20230924220/checkpoint_5.pth',
+              'finetuned': '', 'num_query_token': 32, 'opt_model': 'facebook/opt-2.7b', 'prompt': '',
+              'model_type': 'pretrain_protein_opt2.7b', 'load_pretrained': True, 'freeze_vit': True,
+              'max_protein_len': 600,
+              'max_txt_len': 25}
+
+    model_cls = registry.get_model_class(config['arch'])
+    model = model_cls.from_config(config)
+    model.to(device)
+    model.eval()
+
+    images = df_test['protein'].tolist()
+    n = len(images)
+    bsz = 12
+    iter = n // bsz + 1
+
+    for i in range(iter):
+        image = images[i*bsz: min(n, (i+1)*bsz)]
+        image = [('protein{}'.format(i), x) for i, x in enumerate(image)]
+
+        with model.maybe_autocast():
+            _, _, batch_tokens = model.visual_encoder(image)
+            image_embeds = model.ln_vision(batch_tokens.to(device), repr_layers=[model.vis_layers], return_contacts=True)["representations"][model.vis_layers].contiguous()
+
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+
+        query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+        query_output = model.Qformer.bert(
+            query_embeds=query_tokens,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_atts,
+            return_dict=True,
+        )
+
+        inputs_opt = model.opt_proj(query_output.last_hidden_state)
+        atts_opt = torch.ones(inputs_opt.size()[:-1], dtype=torch.long).to(device)
+
+        model.opt_tokenizer.padding_side = "right"
+
+        text = ['' for i in range(len(image))]
+        opt_tokens = model.opt_tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            truncation=True,
+            max_length=model.max_txt_len,
+        ).to(device)
+        inputs_embeds = model.opt_model.model.decoder.embed_tokens(opt_tokens.input_ids)
+        inputs_embeds = torch.cat([inputs_opt, inputs_embeds], dim=1)
+        attention_mask = torch.cat([atts_opt, opt_tokens.attention_mask], dim=1)
+        num_txt = 10
+        return_num_txt = 5
+        with model.maybe_autocast():
+            outputs = model.opt_model.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, min_length=3,
+                                               max_length=30,
+                                               repetition_penalty=5., num_beams=num_txt, eos_token_id=50118,
+                                               length_penalty=1., num_return_sequences=return_num_txt, temperature=1.)
+        output_text = model.opt_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        output_text = [text.strip() for text in output_text]
+        output_text_ = []
+        for i in range(len(image)):
+            output_text_.append(';'.join(output_text[i * return_num_txt:(i + 1) * return_num_txt]))
+        with open('/cluster/home/wenkai/LAVIS/output/output{}.txt'.format(fix), 'a+') as f:
+            for i in range(len(image)):
+                f.write(image[i][1] + "|" + output_text_[i] + '\n')
+
+
+cat = 'mf'
+fix = '_mf'
+if cat == 'bp':
+    fix = '_bp'
+if cat == 'cc':
+    fix = '_cc'
+
+# model_pth = {'mf': 'uniprot_swissprot_mf_stage1_epo19.pth', 'bp': 'checkpoint17_GO_swissprot_reviewed_bp_stage1.pth', 'cc': ''}
+
+# graph = obonet.read_obo("http://purl.obolibrary.org/obo/go.obo")
+
+# setup device to use
+device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+# device = 'cpu'
+
+### Levenshtein similarity
+test = pd.read_csv('/cluster/home/wenkai/LAVIS/data/sim_split/test{}.csv'.format(fix), sep='|')[:10000]
+test['function'] = test['function'].apply(lambda x: x.lower())
+
+
+if os.path.exists('/cluster/home/wenkai/LAVIS/output/output{}.txt'.format(fix)):
+    os.remove('/cluster/home/wenkai/LAVIS/output/output{}.txt'.format(fix))
+print("stage 2 predict starting")
+stage2_output(test)
+print("stage 2 predict completed")
+
+
+
+df_pred = pd.read_csv('/cluster/home/wenkai/LAVIS/output/output{}.txt'.format(fix), sep='|', header=None, on_bad_lines='warn')
+df_pred.columns = ['protein', 'function']
+df_pred = df_pred.drop_duplicates()
+df_pred['function'] = df_pred['function'].apply(lambda x: str(x).split(';'))
+df_pred['function'] = df_pred['function'].apply(lambda x: [i.strip() for i in list(set(x))])
+
+test.columns
+test_g = test.groupby(['protein']).agg({'function': lambda x: list(x)}).reset_index()
+test_g.columns = ['protein', 'label']
+
+data = pd.merge(df_pred, test_g, on='protein', how='left')
+data = data[data['label'].notnull()]
+
+sim = []
+for text, label in zip(data['function'].tolist(), data['label'].tolist()):
+    sim.append(func(text, label))
+
+data['sim'] = sim
+data['avg_score'] = data['sim'].apply(lambda x: round(np.mean(x), 3))
+print("average similarity score: {}".format(round(data['avg_score'].mean(), 3)))
+# data.to_csv('/home/nilin/LAVIS/predict_{}.csv'.format(cat), index=False, sep='|')
+
+
+test = pd.read_csv('/cluster/home/wenkai/LAVIS/data/sim_split/test{}.csv'.format(fix), sep='|', usecols=['function', 'GO_label'])
+test['function'] = test['function'].apply(lambda x: x.lower())
+test = test.drop_duplicates()
+test_dict = dict(zip(test['function'], test['GO_label']))
+val = pd.read_csv('/cluster/home/wenkai/LAVIS/data/sim_split/val{}.csv'.format(fix), sep='|', usecols=['function', 'GO_label'])
+val['function'] = val['function'].apply(lambda x: x.lower())
+val = val.drop_duplicates()
+val_dict = dict(zip(val['function'], val['GO_label']))
+train = pd.read_csv('/cluster/home/wenkai/LAVIS/data/sim_split/train{}.csv'.format(fix), sep='|', usecols=['function', 'GO_label'])
+train['function'] = train['function'].apply(lambda x: x.lower())
+train = train.drop_duplicates()
+train_dict = dict(zip(train['function'], train['GO_label']))
+
+
+# go_des = pd.read_csv('/home/nilin/LAVIS/data/go_descriptions_new.txt', sep='|', header=None)
+# # go_des = pd.read_csv('/home/nilin/LAVIS/data/go_descriptions.txt', sep='|', header=None)
+# go_des.columns = ['GO', 'function']
+# go_des = go_des[go_des['function'].notnull()]
+# go_des['function'] = go_des['function'].apply(lambda x: x.lower())
+# GO_dict = dict(zip(go_des['function'], go_des['GO']))
+GO_dict = {}
+GO_dict.update(train_dict)
+GO_dict.update(val_dict)
+GO_dict.update(test_dict)
+choices = list(GO_dict.keys())
+
+
+
+# data = pd.read_csv('/home/nilin/LAVIS/predict_{}.csv'.format(cat), sep='|')
+data = data.sort_values(by='protein')
+data = data.drop_duplicates('protein')
+# data = data.sample(1000)
+
+### 预测的文本如果不在GO标签词中，则算作最相似的GO标签
+t0 = time.time()
+txt_dict = {}
+
+all_txt = []
+for txt in data['function']:
+    if type(txt) == str:
+        all_txt.extend(eval(txt))
+    else:
+        all_txt.extend(txt)
+all_txt = list(set(all_txt))
+
+n = len(all_txt)
+thread = 20
+size = int(n/thread)
+inds = list(range(0, n, size))
+inds.append(n)
+all_txt_sep = [all_txt[i: min(i+size, n)] for i in inds[:-1]]
+
+with Pool(processes=thread) as pool:
+    result = pool.map(fuzzy_match, all_txt_sep)
+pool.close()
+pool.join()
+for d in result:
+    txt_dict.update(d)
+
+# for txt in all_txt[:10]:
+#     fuzzy_match(txt)
+
+data['function'] = data['function'].apply(lambda x: txt_map(x, txt_dict))
+data['function'] = data['function'].apply(lambda x: list(set(x)))
+print("fuzzy matching time: {}".format(time.time() - t0))
+
+
+
+
+### Find the generated GO text that not included in the ground truth. Then generate pairs between them.
+# pair_a, pair_b = [], []
+# for preds, labels in zip(data['function'], data['label']):
+#     if type(preds) == str:
+#         preds = eval(preds)
+#     if type(labels) == str:
+#         labels = eval(labels)
+#     l = len(labels)
+#     for pred in preds:
+#         if pred not in labels:
+#             pair_a.extend([pred]*l)
+#             pair_b.extend(labels[:])
+# pair_a = [re.sub('_', ':', GO_dict[i]) for i in pair_a]
+# pair_b = [re.sub('_', ':', GO_dict[i]) for i in pair_b]
+# with open('/home/nilin/LAVIS/examples/GO_pair{}.txt'.format(fix), 'w+') as f:
+#     for i, j in zip(pair_a, pair_b):
+#         f.write(i+' '+j+'\n')
+
+
+# load model
+model_config = {'arch': 'blip2_protein', 'load_finetuned': False,
+                'pretrained': '/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage1/20230922185/checkpoint_15.pth',
+                'finetuned': '', 'num_query_token': 32, 'prompt': '',
+                'model_type': 'pretrain', 'load_pretrained': True, 'freeze_vit': False,
+                'max_protein_len': 512, 'max_txt_len': 25}
+
+model_cls = registry.get_model_class(model_config['arch'])
+model = model_cls.from_config(model_config)
+model = model.to(device)
+model.eval()
+
+# evaluate
+t0 = time.time()
+proteins = list(data['protein'])
+txts = list(data['function'])
+scores = []
+for seq, txt in zip(proteins, txts):
+    image = [('protein1', seq)]
+    _, _, batch_tokens = model.visual_encoder(image)
+    image_embeds = model.ln_vision(batch_tokens.to(device), repr_layers=[30], return_contacts=True)["representations"][
+        30].contiguous()
+
+    image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+
+    query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+
+    query_output = model.Qformer.bert(
+        query_embeds=query_tokens,
+        encoder_hidden_states=image_embeds,
+        encoder_attention_mask=image_atts,
+        use_cache=True,
+        return_dict=True,
+    )
+
+    image_feats = F.normalize(model.vision_proj(query_output.last_hidden_state), dim=-1)
+
+    image_feats_all = concat_all_gather(image_feats)
+
+    if type(txt) == str:
+        txt = eval(txt)
+    length = len(txt)
+    with torch.no_grad():
+        text_tokens = model.tokenizer(
+            txt,
+            padding="max_length",
+            truncation=True,
+            max_length=model.max_txt_len,
+            return_tensors="pt",
+        ).to(device)
+        text_output = model.Qformer.bert(
+            text_tokens.input_ids,
+            attention_mask=text_tokens.attention_mask,
+            return_dict=True,
+        )
+
+        text_feat = F.normalize(
+            model.text_proj(text_output.last_hidden_state[:, 0, :]), dim=-1
+        )
+
+        text_feat_all = concat_all_gather(text_feat)
+        sim_q2t = torch.matmul(image_feats.unsqueeze(1), text_feat_all.unsqueeze(-1)).squeeze()
+        sim_i2t, _ = sim_q2t.max(-1)
+        # print('sim_i2t: {}'.format(sim_i2t))
+    if length > 1:
+        scores.append(list(sim_i2t.detach().cpu().numpy()))
+    else:
+        scores.append([sim_i2t.item()])
+print("model evaluate time: {}".format(time.time() - t0))
+data['score'] = scores
+
+# precision and recall top-k
+topk = 2
+threshould = 0.1
+labels = []
+pred_labels = []
+for l in data['label']:
+    if type(l) == str:
+        l = eval(l)
+    labels.extend(l)
+
+labels = list(set(labels))
+total = len(labels)
+for topk in range(1,7):
+    for threshould in range(1, 25, 1):
+        threshould /= 100
+        filter_txts = []
+        recalls = []
+        precisions = []
+        f1 = []
+        tp_dict, fp_dict, fn_dict = dict(zip(labels, [0]*len(labels))), dict(zip(labels, [0]*len(labels))), dict(zip(labels, [0]*len(labels)))
+        for txts, scores, label in zip(data['function'], data['score'], data['label']):
+            if type(label) == str:
+                label = eval(label)
+            txts_ = np.array(txts)
+            scores = np.array(scores)
+            txts = txts_[scores > threshould]
+            if len(txts) < 1:
+                txts = txts_[np.argmax(scores)]
+            scores = scores[scores > threshould]
+            
+            l = len(scores)
+            ll = len(label)
+            if l <= topk:
+                filter_txts.append(list(txts))
+            else:
+                ind = np.argpartition(scores, -topk)[-topk:]
+                txts = txts[ind]
+                filter_txts.append(list(txts))
+                l = topk
+            for t in label:
+                if t in txts:
+                    tp_dict[t] += 1
+                else:
+                    fn_dict[t] += 1
+            for p in txts:
+                if p not in label:
+                    if p in fp_dict:
+                        fp_dict[p] += 1
+                    else:
+                        fp_dict[p] = 1
+            pred_labels.extend(txts)
+        p_total = len(set(pred_labels))
+        re, pr = 0., 0.
+        for x in labels:
+            re += tp_dict[x] / (1.0 * (tp_dict[x] + fn_dict[x] + 1e-8))
+            pr += tp_dict[x] / (1.0 * (tp_dict[x] + fp_dict[x]+1e-8))
+        r = re / total
+        p = pr / total
+        f1 = 2 * p * r / (p + r)
+        print("Topk: {}, threshould: {}, macro_recall: {}, macro_precision: {}, micro_f1: {}".format(topk, threshould, r, p, f1))
+        #     num_r = 0
+        #     num_p = 0
+        #     for x in label:
+        #         if x in txts:
+        #             num_r += 1
+        #     for x in txts:
+        #         if x in label:
+        #             num_p += 1
+        #     recall = num_r/ll
+        #     precision = num_p/(l+0.0001)
+        #     recalls.append(recall)
+        #     precisions.append(precision)
+        #     f1.append((2*recall*precision)/(recall+precision+0.0001))
+        #
+        # data['predict'] = filter_txts
+        # data['precision'] = precisions
+        # data['recall'] = recalls
+        # data['f1'] = f1
+        # print("Topk: {}, threshould: {}, macro_recall: {}, macro_precision: {}, micro_f1: {}".format(topk, threshould, round(data['recall'].mean(), 4), round(data['precision'].mean(), 4), round(data['f1'].mean(), 4)))
+    
+
+
+
+
+
+# sim = []
+# for text, label in zip(data['predict'].tolist(), data['label'].tolist()):
+#     sim.append(levenshtein_sim(text, label))
+#
+# data['sim_filter'] = sim
+# data['avg_score'] = data['sim_filter'].apply(lambda x: round(np.mean(x), 3))
+
+
+# data['function'] = data['function'].apply(lambda x: eval(re.sub(';', ',', str(x))))
+# data['label'] = data['label'].apply(lambda x: eval(re.sub(';', ',', str(x))))
+# data['sim'] = data['sim'].apply(lambda x: eval(re.sub(';', ',', str(x))))
+#
+# data['function'] = data['function'].apply(lambda x: re.sub(',', ';', str(x)))
+# data['label'] = data['label'].apply(lambda x: re.sub(',', ';', str(x)))
+# data['sim'] = data['sim'].apply(lambda x: re.sub(',', ';', str(x)))
+# data['predict'] = data['predict'].apply(lambda x: re.sub(',', ';', str(x)))
+# data['sim_filter'] = data['sim_filter'].apply(lambda x: re.sub(',', ';', str(x)))
+
+data.to_csv('/cluster/home/wenkai/LAVIS/output/predict_sim{}.csv'.format(fix), sep='|', index=False)
+# data = pd.read_csv('/cluster/home/wenkai/LAVIS/output/predict_sim{}.csv'.format(fix), sep='|')
+
+
+
+
+
+
+
+
+#
+# # example
+# image = ['MIELKHVTFGYNKKQMVLQDINITIPDGENVGILGESGCGKSTLASLVLGLFKPVKGEIYLSDNAVLTIFQHPLTSFNPDWTIETSLKEALYYYRGLTDNTAQDQLLLQHLSTFELNAQLLTKLPSEVSGGQLQRFNVMRSLLAQPRVLICDEITSNLDVIAEQNVINILKAQTITNLNHFIVISHDLSVLQRLVNRIIVLKDGMIVDDFAIEELFNVDRHPYTKELVQTFSY']
+# image = [('protein{}'.format(i), x) for i, x in enumerate(image)]
+#
+# _, _, batch_tokens = model.visual_encoder(image)
+# image_embeds = model.ln_vision(batch_tokens.to(device), repr_layers=[30], return_contacts=True)["representations"][30].contiguous()
+#
+# image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+#
+# query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+#
+# query_output = model.Qformer.bert(
+#     query_embeds=query_tokens,
+#     encoder_hidden_states=image_embeds,
+#     encoder_attention_mask=image_atts,
+#     use_cache=True,
+#     return_dict=True,
+# )
+#
+# image_feats = F.normalize(model.vision_proj(query_output.last_hidden_state), dim=-1)
+#
+# image_feats_all = concat_all_gather(image_feats)
+#
+# functions = ['transmembrane transporter activity', 'nickel cation transmembrane transporter activity', 'nickel cation binding', 'atp hydrolysis activity', 'atp hydrolysis', 'cadmium binding', 'abc-type nickel transmembrane transporter activity', 'abc-type nickel transporter activity', 'nickel transmembrane transporter activity', 'atp binding']
+# for text in functions:
+#     with torch.no_grad():
+#         # text = 'flavin adenine dinucleotide binding'
+#         text_tokens = model.tokenizer(
+#             text,
+#             padding="max_length",
+#             truncation=True,
+#             max_length=model.max_txt_len,
+#             return_tensors="pt",
+#         ).to(device)
+#         text_output = model.Qformer.bert(
+#             text_tokens.input_ids,
+#             attention_mask=text_tokens.attention_mask,
+#             return_dict=True,
+#         )
+#
+#         text_feat = F.normalize(
+#             model.text_proj(text_output.last_hidden_state[:, 0, :]), dim=-1
+#         )
+#
+#         text_feat_all = concat_all_gather(text_feat)
+#         sim_q2t = torch.matmul(image_feats.unsqueeze(1), text_feat_all.unsqueeze(-1)).squeeze()
+#         sim_i2t, _ = sim_q2t.max(-1)
+#         print('sim_i2t: {}'.format(sim_i2t))
+#
+#         # # text-query similarity: [batch_size, batch_size*num_gpu, num_query_tokens]
+#         # sim_t2q = torch.matmul(
+#         #     text_feat.unsqueeze(1).unsqueeze(1), image_feats_all.permute(0, 2, 1)
+#         # ).squeeze()
+#         #
+#         # # text-image similarity: aggregate across all query tokens
+#         # sim_t2i, _ = sim_t2q.max(-1)
+#         # print('sim_t2i: {}'.format(sim_t2i))
+
+
+

examples/blip2_predict_func.py

@@ -0,0 +1,178 @@
+import os
+import re
+
+import torch
+from PIL import Image
+
+from lavis.models import load_model_and_preprocess
+from lavis.processors import load_processor
+from lavis.common.registry import registry
+from torch.nn import functional as F
+from lavis.models.base_model import all_gather_with_grad, concat_all_gather
+import numpy as np
+import pandas as pd
+import time
+from fuzzywuzzy import process
+from multiprocessing import Pool, Queue, Process
+import difflib
+import Levenshtein
+# import obonet
+
+
+# setup device to use
+device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+# device = 'cpu'
+
+
+def txt_map(x, txt_dict):
+    if type(x) == str:
+        x = eval(x)
+    x_ = []
+    for i in x:
+        if i in txt_dict:
+            x_.append(txt_dict[i])
+        else:
+            x_.append(i)
+    return x_
+
+
+def levenshtein_sim(text, label):
+    all_s = []
+    for x in label:
+        s = 0
+        for y in text:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+def func(text, label):
+    all_s = []
+    for x in text:
+        s = 0
+        for y in label:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+
+def stage2_output(df_test, return_num_txt=1):
+    config = {'arch': 'blip2_protein_opt', 'load_finetuned': False,
+              'pretrained': '/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage2/20230924220/checkpoint_5.pth',
+              'finetuned': '', 'num_query_token': 32, 'opt_model': 'facebook/opt-2.7b', 'prompt': '',
+              'model_type': 'pretrain_protein_opt2.7b', 'load_pretrained': True, 'freeze_vit': True,
+              'max_protein_len': 600,
+              'max_txt_len': 25}
+
+    model_cls = registry.get_model_class(config['arch'])
+    model = model_cls.from_config(config)
+    model.to(device)
+    model.eval()
+
+    images = df_test['protein'].tolist()
+    n = len(images)
+    bsz = 12
+    iter = n // bsz + 1
+
+    for i in range(iter):
+        image = images[i*bsz: min(n, (i+1)*bsz)]
+        image = [('protein{}'.format(i), x) for i, x in enumerate(image)]
+
+        with model.maybe_autocast():
+            _, _, batch_tokens = model.visual_encoder(image)
+            image_embeds = model.ln_vision(batch_tokens.to(device), repr_layers=[model.vis_layers], return_contacts=True)["representations"][model.vis_layers].contiguous()
+
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+
+        query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+        query_output = model.Qformer.bert(
+            query_embeds=query_tokens,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_atts,
+            return_dict=True,
+        )
+
+        inputs_opt = model.opt_proj(query_output.last_hidden_state)
+        atts_opt = torch.ones(inputs_opt.size()[:-1], dtype=torch.long).to(device)
+
+        model.opt_tokenizer.padding_side = "right"
+
+        text = ['' for i in range(len(image))]
+        opt_tokens = model.opt_tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            truncation=True,
+            max_length=model.max_txt_len,
+        ).to(device)
+        inputs_embeds = model.opt_model.model.decoder.embed_tokens(opt_tokens.input_ids)
+        inputs_embeds = torch.cat([inputs_opt, inputs_embeds], dim=1)
+        attention_mask = torch.cat([atts_opt, opt_tokens.attention_mask], dim=1)
+        num_txt = 6
+        with model.maybe_autocast():
+            outputs = model.opt_model.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, min_length=3,
+                                               max_length=30,
+                                               repetition_penalty=1., num_beams=num_txt, eos_token_id=50118,
+                                               length_penalty=1., num_return_sequences=return_num_txt, temperature=1.)
+        output_text = model.opt_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        output_text = [text.strip() for text in output_text]
+        output_text_ = []
+        for i in range(len(image)):
+            output_text_.append(';'.join(output_text[i * return_num_txt:(i + 1) * return_num_txt]))
+        with open('/cluster/home/wenkai/LAVIS/output/output{}.txt'.format(fix), 'a+') as f:
+            for i in range(len(image)):
+                f.write(image[i][1] + "|" + output_text_[i] + '\n')
+
+
+cat = 'mf'
+fix = '_mf'
+if cat == 'bp':
+    fix = '_bp'
+if cat == 'cc':
+    fix = '_cc'
+
+return_num_txt = 1
+# graph = obonet.read_obo("http://purl.obolibrary.org/obo/go.obo")
+
+### Levenshtein similarity
+test = pd.read_csv('/cluster/home/wenkai/LAVIS/data/sim_split/test{}.csv'.format(fix), sep='|')
+test['function'] = test['function'].apply(lambda x: x.lower())
+
+
+if os.path.exists('/cluster/home/wenkai/LAVIS/output/output{}.txt'.format(fix)):
+    os.remove('/cluster/home/wenkai/LAVIS/output/output{}.txt'.format(fix))
+print("stage 2 predict starting")
+stage2_output(test)
+print("stage 2 predict completed")
+
+df_pred = pd.read_csv('/cluster/home/wenkai/LAVIS/output/output{}.txt'.format(fix), sep='|', header=None, on_bad_lines='warn')
+df_pred.columns = ['protein', 'function']
+df_pred = df_pred.drop_duplicates()
+df_pred['function'] = df_pred['function'].apply(lambda x: str(x).split(';'))
+df_pred['function'] = df_pred['function'].apply(lambda x: [i.strip() for i in list(set(x))])
+
+test_g = test.groupby(['protein']).agg({'function': lambda x: list(x)}).reset_index()
+test_g.columns = ['protein', 'label']
+
+data = pd.merge(df_pred, test_g, on='protein', how='left')
+data = data[data['label'].notnull()]
+
+sim = []
+for text, label in zip(data['function'].tolist(), data['label'].tolist()):
+    sim.append(func(text, label))
+
+data['sim'] = sim
+data['avg_score'] = data['sim'].apply(lambda x: round(np.mean(x), 3))
+data['count'] = data['sim'].apply(lambda x: x.count(1.))
+print("average similarity score: {}".format(round(data['avg_score'].mean(), 3)))
+print("Return texts: {}; Accuracy: {}".format(return_num_txt, data['count'].sum()/(return_num_txt*data.shape[0])))
+data.to_csv('/cluster/home/wenkai/LAVIS/output/predict_{}.csv'.format(cat), index=False, sep='|')
+
+
+
+

examples/blip2_predict_func_concat.py

@@ -0,0 +1,193 @@
+import os
+import re
+
+import torch
+from PIL import Image
+
+from lavis.models import load_model_and_preprocess
+from lavis.processors import load_processor
+from lavis.common.registry import registry
+from torch.nn import functional as F
+from lavis.models.base_model import all_gather_with_grad, concat_all_gather
+import numpy as np
+import pandas as pd
+import time
+from fuzzywuzzy import process
+from multiprocessing import Pool, Queue, Process
+import difflib
+import Levenshtein
+
+# import obonet
+
+
+# setup device to use
+device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+
+
+# device = torch.device("cuda")
+
+
+def txt_map(x, txt_dict):
+    if type(x) == str:
+        x = eval(x)
+    x_ = []
+    for i in x:
+        if i in txt_dict:
+            x_.append(txt_dict[i])
+        else:
+            x_.append(i)
+    return x_
+
+
+def levenshtein_sim(text, label):
+    all_s = []
+    for x in label:
+        s = 0
+        for y in text:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+
+def func(text, label):
+    all_s = []
+    for x in text:
+        s = 0
+        for y in label:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+
+def stage2_output(df_test, return_num_txt=1):
+    config = {'arch': 'blip2_protein_opt', 'load_finetuned': False,
+              'pretrained': '/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage2/20231029182/checkpoint_0.pth',
+              'finetuned': '', 'num_query_token': 32, 'opt_model': 'facebook/opt-2.7b', 'prompt': '',
+              'model_type': 'pretrain_protein_opt2.7b', 'load_pretrained': True, 'freeze_vit': True,
+              'max_protein_len': 600,
+              'max_txt_len': 256}
+
+    model_cls = registry.get_model_class(config['arch'])
+    model = model_cls.from_config(config)
+    model.to(device)
+    model.eval()
+
+    images = df_test['protein'].tolist()
+    n = len(images)
+    bsz = 8
+    iter = n // bsz + 1
+    with open('/cluster/home/wenkai/LAVIS/output/output_concat_{}{}{}.txt'.format(split, fix, type_fix), 'a+') as f:
+        for i in range(iter):
+            image = images[i * bsz: min(n, (i + 1) * bsz)]
+            image = [('protein{}'.format(i), x) for i, x in enumerate(image)]
+
+            with model.maybe_autocast():
+                _, _, batch_tokens = model.visual_encoder(image)
+                image_embeds = \
+                model.ln_vision(batch_tokens.to(device), repr_layers=[model.vis_layers], return_contacts=True)[
+                    "representations"][model.vis_layers].contiguous()
+
+            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+
+            query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+            query_output = model.Qformer.bert(
+                query_embeds=query_tokens,
+                encoder_hidden_states=image_embeds,
+                encoder_attention_mask=image_atts,
+                return_dict=True,
+            )
+
+            inputs_opt = model.opt_proj(query_output.last_hidden_state)
+            atts_opt = torch.ones(inputs_opt.size()[:-1], dtype=torch.long).to(device)
+
+            model.opt_tokenizer.padding_side = "right"
+
+            text = ['' for i in range(len(image))]
+            opt_tokens = model.opt_tokenizer(
+                text,
+                return_tensors="pt",
+                padding="longest",
+                truncation=True,
+                max_length=model.max_txt_len,
+            ).to(device)
+            inputs_embeds = model.opt_model.model.decoder.embed_tokens(opt_tokens.input_ids)
+            inputs_embeds = torch.cat([inputs_opt, inputs_embeds], dim=1)
+            attention_mask = torch.cat([atts_opt, opt_tokens.attention_mask], dim=1)
+            num_txt = 5
+            with model.maybe_autocast():
+                outputs = model.opt_model.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, min_length=1,
+                                                   max_length=256,
+                                                   repetition_penalty=1., num_beams=num_txt, eos_token_id=50118,
+                                                   length_penalty=1., num_return_sequences=return_num_txt, temperature=1.)
+            output_text = model.opt_tokenizer.batch_decode(outputs)
+
+            output_text = [re.sub('\t', '', str(x)) for x in output_text]
+            output_text = [text.strip() for text in output_text]
+            output_text_ = []
+            for i in range(len(image)):
+                output_text_.append(';'.join(output_text[i * return_num_txt:(i + 1) * return_num_txt]))
+
+            for i in range(len(image)):
+                f.write(image[i][1] + "|" + output_text_[i] + '\n')
+
+
+if __name__=="__main__":
+    split = 'test'
+    cat = 'bp'
+    fix = '_mf'
+    type_fix = ''
+    if cat == 'bp':
+        fix = '_bp'
+    if cat == 'cc':
+        fix = '_cc'
+
+    print(device)
+    return_num_txt = 1
+    # graph = obonet.read_obo("http://purl.obolibrary.org/obo/go.obo")
+
+    ### Levenshtein similarity
+    print("reading file ...")
+    test = pd.read_csv('/cluster/home/wenkai/LAVIS/data/sim_split_concat/{}{}.csv'.format(split, fix),
+                       usecols=['name', 'protein', 'function'], sep='|')
+    # test['function'] = test['function'].apply(lambda x: x.lower().split('; '))
+    test.columns = ['name', 'protein', 'label']
+
+    if os.path.exists('/cluster/home/wenkai/LAVIS/output/output_concat_{}{}{}.txt'.format(split, fix, type_fix)):
+        os.remove('/cluster/home/wenkai/LAVIS/output/output_concat_{}{}{}.txt'.format(split, fix, type_fix))
+    print("stage 2 predict starting")
+    stage2_output(test)
+    print("stage 2 predict completed")
+
+    df_pred = pd.read_csv('/cluster/home/wenkai/LAVIS/output/output_concat_{}{}{}.txt'.format(split, fix, type_fix), sep='|',
+                          header=None, on_bad_lines='warn')
+    df_pred.columns = ['protein', 'pred']
+    df_pred = df_pred.drop_duplicates()
+    # df_pred['function'] = df_pred['function'].apply(lambda x: str(x).split(';'))
+    # df_pred['function'] = df_pred['function'].apply(lambda x: [i.strip() for i in list(set(x))])
+
+
+    data = pd.merge(df_pred, test, on='protein', how='left')
+    data = data[data['label'].notnull()]
+
+    # sim = []
+    # for text, label in zip(data['function'].tolist(), data['label'].tolist()):
+    #    sim.append(func(text, label))
+
+    # data['sim'] = sim
+    # data['avg_score'] = data['sim'].apply(lambda x: round(np.mean(x), 3))
+    # data['count'] = data['sim'].apply(lambda x: x.count(1.))
+    # print("average similarity score: {}".format(round(data['avg_score'].mean(), 3)))
+    # print("Return texts: {}; Accuracy: {}".format(return_num_txt, data['count'].sum()/(return_num_txt*data.shape[0])))
+    data[['name', 'label', 'pred']].to_csv(
+        '/cluster/home/wenkai/LAVIS/output/predict_concat_{}{}{}.csv'.format(split, cat, type_fix), index=False, sep='|')
+
+
+
+
+

examples/blip2_predict_func_concat_pretrain.py

@@ -0,0 +1,197 @@
+import os
+import re
+
+import torch
+from PIL import Image
+
+from lavis.models import load_model_and_preprocess
+from lavis.processors import load_processor
+from lavis.common.registry import registry
+from torch.nn import functional as F
+from lavis.models.base_model import all_gather_with_grad, concat_all_gather
+import numpy as np
+import pandas as pd
+import time
+from fuzzywuzzy import process
+from multiprocessing import Pool, Queue, Process
+import difflib
+import Levenshtein
+
+# import obonet
+
+
+# setup device to use
+device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+
+
+# device = torch.device("cuda")
+
+
+def txt_map(x, txt_dict):
+    if type(x) == str:
+        x = eval(x)
+    x_ = []
+    for i in x:
+        if i in txt_dict:
+            x_.append(txt_dict[i])
+        else:
+            x_.append(i)
+    return x_
+
+
+def levenshtein_sim(text, label):
+    all_s = []
+    for x in label:
+        s = 0
+        for y in text:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+
+def func(text, label):
+    all_s = []
+    for x in text:
+        s = 0
+        for y in label:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+
+def stage2_output(df_test, return_num_txt=1):
+    config = {'arch': 'blip2_protein_opt', 'load_finetuned': False,
+              'pretrained': '/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage2/20231029182/checkpoint_0.pth',
+              'finetuned': '', 'num_query_token': 32, 'opt_model': 'facebook/opt-2.7b', 'prompt': '',
+              'model_type': 'pretrain_protein_opt2.7b', 'load_pretrained': True, 'freeze_vit': True,
+              'max_protein_len': 600,
+              'max_txt_len': 256}
+
+    model_cls = registry.get_model_class(config['arch'])
+    model = model_cls.from_config(config)
+    model.to(device)
+    model.eval()
+
+    images = df_test['protein'].tolist()
+    n = len(images)
+    bsz = 8
+    iter = n // bsz + 1
+    if n > 0:
+        for i in range(iter):
+            image = images[i * bsz: min(n, (i + 1) * bsz)]
+            image = [('protein{}'.format(i), x) for i, x in enumerate(image)]
+
+            with model.maybe_autocast():
+                _, _, batch_tokens = model.visual_encoder(image)
+                image_embeds = \
+                model.ln_vision(batch_tokens.to(device), repr_layers=[model.vis_layers], return_contacts=True)[
+                    "representations"][model.vis_layers].contiguous()
+
+            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+
+            query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+            query_output = model.Qformer.bert(
+                query_embeds=query_tokens,
+                encoder_hidden_states=image_embeds,
+                encoder_attention_mask=image_atts,
+                return_dict=True,
+            )
+
+            inputs_opt = model.opt_proj(query_output.last_hidden_state)
+            atts_opt = torch.ones(inputs_opt.size()[:-1], dtype=torch.long).to(device)
+
+            model.opt_tokenizer.padding_side = "right"
+
+            text = ['' for i in range(len(image))]
+            opt_tokens = model.opt_tokenizer(
+                text,
+                return_tensors="pt",
+                padding="longest",
+                truncation=True,
+                max_length=model.max_txt_len,
+            ).to(device)
+            inputs_embeds = model.opt_model.model.decoder.embed_tokens(opt_tokens.input_ids)
+            inputs_embeds = torch.cat([inputs_opt, inputs_embeds], dim=1)
+            attention_mask = torch.cat([atts_opt, opt_tokens.attention_mask], dim=1)
+            num_txt = 5
+            with model.maybe_autocast():
+                outputs = model.opt_model.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, min_length=1,
+                                                   max_length=256,
+                                                   repetition_penalty=1., num_beams=num_txt, eos_token_id=50118,
+                                                   length_penalty=1., num_return_sequences=return_num_txt, temperature=1.)
+            output_text = model.opt_tokenizer.batch_decode(outputs)
+
+            output_text = [re.sub('\t', '', str(x)) for x in output_text]
+            output_text = [text.strip() for text in output_text]
+            output_text_ = []
+            for i in range(len(image)):
+                output_text_.append(';'.join(output_text[i * return_num_txt:(i + 1) * return_num_txt]))
+
+            f = open('/cluster/home/wenkai/LAVIS/output/output_concat_{}{}{}.txt'.format(split, fix, type_fix), 'a+')
+            for i in range(len(image)):
+                f.write(image[i][1] + "|" + output_text_[i] + '\n')
+            f.close()
+
+                
+
+
+if __name__=="__main__":
+    split = 'test'
+    cat = ''
+    fix = ''
+    type_fix = '_pretrain'
+    if cat == 'bp':
+        fix = '_bp'
+    if cat == 'cc':
+        fix = '_cc'
+
+    print(device)
+    return_num_txt = 1
+    # graph = obonet.read_obo("http://purl.obolibrary.org/obo/go.obo")
+
+    ### Levenshtein similarity
+    print("reading file ...")
+    test = pd.read_csv('/cluster/home/wenkai/LAVIS/data/pretrain/{}_sample10000.csv'.format(split),
+                       usecols=['name', 'protein', 'function'], sep='|')
+    # test['function'] = test['function'].apply(lambda x: x.lower().split('; '))
+    test.columns = ['name', 'protein', 'label']
+
+    if os.path.exists('/cluster/home/wenkai/LAVIS/output/output_concat_{}{}{}.txt'.format(split, fix, type_fix)):
+        os.remove('/cluster/home/wenkai/LAVIS/output/output_concat_{}{}{}.txt'.format(split, fix, type_fix))
+    print("stage 2 predict starting")
+    stage2_output(test)
+    print("stage 2 predict completed")
+
+    df_pred = pd.read_csv('/cluster/home/wenkai/LAVIS/output/output_concat_{}{}{}.txt'.format(split, fix, type_fix), sep='|',
+                          header=None, on_bad_lines='warn')
+    df_pred.columns = ['protein', 'pred']
+    df_pred = df_pred.drop_duplicates()
+    # df_pred['function'] = df_pred['function'].apply(lambda x: str(x).split(';'))
+    # df_pred['function'] = df_pred['function'].apply(lambda x: [i.strip() for i in list(set(x))])
+
+
+    data = pd.merge(df_pred, test, on='protein', how='left')
+    data = data[data['label'].notnull()]
+
+    # sim = []
+    # for text, label in zip(data['function'].tolist(), data['label'].tolist()):
+    #    sim.append(func(text, label))
+
+    # data['sim'] = sim
+    # data['avg_score'] = data['sim'].apply(lambda x: round(np.mean(x), 3))
+    # data['count'] = data['sim'].apply(lambda x: x.count(1.))
+    # print("average similarity score: {}".format(round(data['avg_score'].mean(), 3)))
+    # print("Return texts: {}; Accuracy: {}".format(return_num_txt, data['count'].sum()/(return_num_txt*data.shape[0])))
+    data[['name', 'label', 'pred']].to_csv(
+        '/cluster/home/wenkai/LAVIS/output/predict_concat_{}{}{}.csv'.format(split, cat, type_fix), index=False, sep='|')
+
+
+
+
+

examples/blip2_predict_func_concat_timesplit.py

@@ -0,0 +1,166 @@
+import os
+import re
+
+import torch
+from PIL import Image
+
+from lavis.models import load_model_and_preprocess
+from lavis.processors import load_processor
+from lavis.common.registry import registry
+from torch.nn import functional as F
+from lavis.models.base_model import all_gather_with_grad, concat_all_gather
+import numpy as np
+import pandas as pd
+import time
+from fuzzywuzzy import process
+from multiprocessing import Pool, Queue, Process
+import difflib
+import Levenshtein
+# import obonet
+
+
+# setup device to use
+device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+# device = 'cpu'
+
+
+def txt_map(x, txt_dict):
+    if type(x) == str:
+        x = eval(x)
+    x_ = []
+    for i in x:
+        if i in txt_dict:
+            x_.append(txt_dict[i])
+        else:
+            x_.append(i)
+    return x_
+
+
+def levenshtein_sim(text, label):
+    all_s = []
+    for x in label:
+        s = 0
+        for y in text:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+def func(text, label):
+    all_s = []
+    for x in text:
+        s = 0
+        for y in label:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+
+def stage2_output(df_test, return_num_txt=1):
+    config = {'arch': 'blip2_protein_opt', 'load_finetuned': False,
+              'pretrained': '/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage2/20231007085/checkpoint_19.pth',
+              'finetuned': '', 'num_query_token': 32, 'opt_model': 'facebook/opt-2.7b', 'prompt': '',
+              'model_type': 'pretrain_protein_opt2.7b', 'load_pretrained': True, 'freeze_vit': True,
+              'max_protein_len': 600,
+              'max_txt_len': 256}
+
+    model_cls = registry.get_model_class(config['arch'])
+    model = model_cls.from_config(config)
+    model.to(device)
+    model.eval()
+
+    images = df_test['protein'].tolist()
+    n = len(images)
+    bsz = 12
+    iter = n // bsz + 1
+
+    for i in range(iter):
+        image = images[i*bsz: min(n, (i+1)*bsz)]
+        image = [('protein{}'.format(i), x) for i, x in enumerate(image)]
+
+        with model.maybe_autocast():
+            _, _, batch_tokens = model.visual_encoder(image)
+            image_embeds = model.ln_vision(batch_tokens.to(device), repr_layers=[model.vis_layers], return_contacts=True)["representations"][model.vis_layers].contiguous()
+
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+
+        query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+        query_output = model.Qformer.bert(
+            query_embeds=query_tokens,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_atts,
+            return_dict=True,
+        )
+
+        inputs_opt = model.opt_proj(query_output.last_hidden_state)
+        atts_opt = torch.ones(inputs_opt.size()[:-1], dtype=torch.long).to(device)
+
+        model.opt_tokenizer.padding_side = "right"
+
+        text = ['' for i in range(len(image))]
+        opt_tokens = model.opt_tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            truncation=True,
+            max_length=model.max_txt_len,
+        ).to(device)
+        inputs_embeds = model.opt_model.model.decoder.embed_tokens(opt_tokens.input_ids)
+        inputs_embeds = torch.cat([inputs_opt, inputs_embeds], dim=1)
+        attention_mask = torch.cat([atts_opt, opt_tokens.attention_mask], dim=1)
+        num_txt = 5
+        with model.maybe_autocast():
+            outputs = model.opt_model.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, min_length=3,
+                                               max_length=256,
+                                               repetition_penalty=1., num_beams=num_txt, eos_token_id=50118,
+                                               length_penalty=1., num_return_sequences=return_num_txt, temperature=1.)
+        output_text = model.opt_tokenizer.batch_decode(outputs)
+        output_text = [re.sub('\t', '', x) for x in output_text]
+
+        output_text = [text.strip() for text in output_text]
+        output_text_ = []
+        for i in range(len(image)):
+            output_text_.append(';'.join(output_text[i * return_num_txt:(i + 1) * return_num_txt]))
+        with open('/cluster/home/wenkai/LAVIS/output/output_timeconcat{}.txt'.format(fix), 'a+') as f:
+            for i in range(len(image)):
+                f.write(image[i][1] + "|" + output_text_[i] + '\n')
+
+
+cat = 'mf'
+fix = '_mf'
+if cat == 'bp':
+    fix = '_bp'
+if cat == 'cc':
+    fix = '_cc'
+
+return_num_txt = 1
+# graph = obonet.read_obo("http://purl.obolibrary.org/obo/go.obo")
+
+### Levenshtein similarity
+test = pd.read_csv('/cluster/home/wenkai/LAVIS/data/review_time_concat/test{}.csv'.format(fix), usecols=['name', 'protein', 'function'], sep='|')
+#test['function'] = test['function'].apply(lambda x: x.lower().split('; '))
+test.columns = ['name', 'protein', 'label']
+
+if os.path.exists('/cluster/home/wenkai/LAVIS/output/output_timeconcat{}.txt'.format(fix)):
+    os.remove('/cluster/home/wenkai/LAVIS/output/output_timeconcat{}.txt'.format(fix))
+print("stage 2 predict starting")
+stage2_output(test)
+print("stage 2 predict completed")
+
+df_pred = pd.read_csv('/cluster/home/wenkai/LAVIS/output/output_timeconcat{}.txt'.format(fix), sep='|', header=None, on_bad_lines='warn')
+df_pred.columns = ['protein', 'pred']
+df_pred = df_pred.drop_duplicates()
+
+data = pd.merge(df_pred, test, on='protein', how='left')
+data = data[data['label'].notnull()]
+
+data[['name', 'label', 'pred']].to_csv('/cluster/home/wenkai/LAVIS/output/predict_timeconcat_{}.csv'.format(cat), index=False, sep='|')
+
+
+
+

examples/blip2_predict_names.py

@@ -0,0 +1,247 @@
+import os
+import re
+
+import torch
+from PIL import Image
+
+from lavis.models import load_model_and_preprocess
+from lavis.processors import load_processor
+from lavis.common.registry import registry
+from torch.nn import functional as F
+from lavis.models.base_model import all_gather_with_grad, concat_all_gather
+import numpy as np
+import pandas as pd
+import time
+from fuzzywuzzy import process
+from multiprocessing import Pool, Queue, Process
+import difflib
+import Levenshtein
+# import obonet
+
+
+# setup device to use
+device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+# device = 'cpu'
+
+
+def txt_map(x, txt_dict):
+    if type(x) == str:
+        x = eval(x)
+    x_ = []
+    for i in x:
+        if i in txt_dict:
+            x_.append(txt_dict[i])
+        else:
+            x_.append(i)
+    return x_
+
+
+def levenshtein_sim(text, label):
+    all_s = []
+    for x in label:
+        s = 0
+        for y in text:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+def func(text, label):
+    all_s = []
+    for x in label:
+        s = 0
+        for y in text:
+            temp = Levenshtein.ratio(x, y)
+            if temp > s:
+                s = temp
+        all_s.append(s)
+    all_s = [round(i, 3) for i in all_s]
+    return all_s
+
+
+def stage2_output(df_test):
+    config = {'arch': 'blip2_protein_opt', 'load_finetuned': False,
+              'pretrained': '/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage2/20230926091/checkpoint_3.pth',
+              'finetuned': '', 'num_query_token': 32, 'opt_model': 'facebook/opt-2.7b', 'prompt': '',
+              'model_type': 'pretrain_protein_opt2.7b', 'load_pretrained': True, 'freeze_vit': True,
+              'max_protein_len': 600,
+              'max_txt_len': 25}
+
+    model_cls = registry.get_model_class(config['arch'])
+    model = model_cls.from_config(config)
+    model.to(device)
+    model.eval()
+
+    images = df_test['protein'].tolist()
+    n = len(images)
+    bsz = 12
+    iter = n // bsz + 1
+
+    for i in range(iter):
+        image = images[i*bsz: min(n, (i+1)*bsz)]
+        image = [('protein{}'.format(i), x) for i, x in enumerate(image)]
+
+        with model.maybe_autocast():
+            _, _, batch_tokens = model.visual_encoder(image)
+            image_embeds = model.ln_vision(batch_tokens.to(device), repr_layers=[model.vis_layers], return_contacts=True)["representations"][model.vis_layers].contiguous()
+
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+
+        query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+        query_output = model.Qformer.bert(
+            query_embeds=query_tokens,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_atts,
+            return_dict=True,
+        )
+
+        inputs_opt = model.opt_proj(query_output.last_hidden_state)
+        atts_opt = torch.ones(inputs_opt.size()[:-1], dtype=torch.long).to(device)
+
+        model.opt_tokenizer.padding_side = "right"
+
+        text = ['' for i in range(len(image))]
+        opt_tokens = model.opt_tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            truncation=True,
+            max_length=model.max_txt_len,
+        ).to(device)
+        inputs_embeds = model.opt_model.model.decoder.embed_tokens(opt_tokens.input_ids)
+        inputs_embeds = torch.cat([inputs_opt, inputs_embeds], dim=1)
+        attention_mask = torch.cat([atts_opt, opt_tokens.attention_mask], dim=1)
+        num_txt = 5
+        return_num_txt = 2
+        with model.maybe_autocast():
+            outputs = model.opt_model.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, min_length=3,
+                                               max_length=30,
+                                               repetition_penalty=5., num_beams=num_txt, eos_token_id=50118,
+                                               length_penalty=1., num_return_sequences=return_num_txt, temperature=1.)
+        output_text = model.opt_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        output_text = [text.strip() for text in output_text]
+        output_text_ = []
+        for i in range(len(image)):
+            output_text_.append(';'.join(output_text[i * return_num_txt:(i + 1) * return_num_txt]))
+        with open('/cluster/home/wenkai/LAVIS/output/output_names.txt', 'a+') as f:
+            for i in range(len(image)):
+                f.write(image[i][1] + "|" + output_text_[i] + '\n')
+
+
+def evaluate_score(data):
+    model_config = {'arch': 'blip2_protein', 'load_finetuned': False,
+                    'pretrained': '/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage1/20230925102/checkpoint_6.pth',
+                    'finetuned': '', 'num_query_token': 32, 'prompt': '',
+                    'model_type': 'pretrain', 'load_pretrained': True, 'freeze_vit': False,
+                    'max_protein_len': 512, 'max_txt_len': 30}
+
+    model_cls = registry.get_model_class(model_config['arch'])
+    model = model_cls.from_config(model_config)
+    model = model.to(device)
+    model.eval()
+
+    # evaluate
+    t0 = time.time()
+    proteins = list(data['protein'])
+    txts = list(data['function'])
+    scores = []
+    for seq, txt in zip(proteins, txts):
+        image = [('protein1', seq)]
+        _, _, batch_tokens = model.visual_encoder(image)
+        image_embeds = \
+        model.ln_vision(batch_tokens.to(device), repr_layers=[30], return_contacts=True)["representations"][
+            30].contiguous()
+
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+
+        query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+
+        query_output = model.Qformer.bert(
+            query_embeds=query_tokens,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_atts,
+            use_cache=True,
+            return_dict=True,
+        )
+
+        image_feats = F.normalize(model.vision_proj(query_output.last_hidden_state), dim=-1)
+
+        image_feats_all = concat_all_gather(image_feats)
+
+        if type(txt) == str:
+            txt = eval(txt)
+        length = len(txt)
+        with torch.no_grad():
+            text_tokens = model.tokenizer(
+                txt,
+                padding="max_length",
+                truncation=True,
+                max_length=model.max_txt_len,
+                return_tensors="pt",
+            ).to(device)
+            text_output = model.Qformer.bert(
+                text_tokens.input_ids,
+                attention_mask=text_tokens.attention_mask,
+                return_dict=True,
+            )
+
+            text_feat = F.normalize(
+                model.text_proj(text_output.last_hidden_state[:, 0, :]), dim=-1
+            )
+
+            text_feat_all = concat_all_gather(text_feat)
+            sim_q2t = torch.matmul(image_feats.unsqueeze(1), text_feat_all.unsqueeze(-1)).squeeze()
+            sim_i2t, _ = sim_q2t.max(-1)
+            # print('sim_i2t: {}'.format(sim_i2t))
+        if length > 1:
+            scores.append(list(sim_i2t.detach().cpu().numpy()))
+        else:
+            scores.append([sim_i2t.item()])
+    print("model evaluate time: {}".format(time.time() - t0))
+    data['sim'] = scores
+    return data
+
+
+
+# graph = obonet.read_obo("http://purl.obolibrary.org/obo/go.obo")
+
+### Levenshtein similarity
+test = pd.read_csv('/cluster/home/wenkai/LAVIS/data/raw_time_split/reviewed//test.csv', sep='|')
+test['function'] = test['function'].apply(lambda x: x.lower())
+
+
+if os.path.exists('/cluster/home/wenkai/LAVIS/output/output_names.txt'):
+    os.remove('/cluster/home/wenkai/LAVIS/output/output_names.txt')
+print("stage 2 predict starting")
+stage2_output(test)
+print("stage 2 predict completed")
+
+df_pred = pd.read_csv('/cluster/home/wenkai/LAVIS/output/output_names.txt', sep='|', header=None, on_bad_lines='warn')
+df_pred.columns = ['protein', 'function']
+df_pred = df_pred.drop_duplicates()
+df_pred['function'] = df_pred['function'].apply(lambda x: str(x).split(';'))
+df_pred['function'] = df_pred['function'].apply(lambda x: [i.strip() for i in list(set(x))])
+
+test.columns
+test_g = test.groupby(['protein']).agg({'function': lambda x: list(x)}).reset_index()
+test_g.columns = ['protein', 'label']
+
+data = pd.merge(df_pred, test_g, on='protein', how='left')
+data = data[data['label'].notnull()]
+
+sim = []
+for text, label in zip(data['function'].tolist(), data['label'].tolist()):
+    sim.append(func(text, label))
+
+data['sim'] = sim
+data['avg_score'] = data['sim'].apply(lambda x: round(np.mean(x), 3))
+print("average similarity score: {}".format(round(data['avg_score'].mean(), 3)))
+data.to_csv('/cluster/home/wenkai/LAVIS/output/output_names.csv', index=False, sep='|')
+
+
+
+
+
+

examples/predict_test.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+#SBATCH -J infer_test
+#SBATCH -p gpu1
+#SBATCH -N 1
+#SBATCH -w node[84]
+#SBATCH --mem 80G
+#SBATCH --gres=gpu:1
+#SBATCH --output=log_predict_test.out
+#SBATCH --error=log_predict_test.err
+#SBATCH --cpus-per-task=8
+module load anaconda3/2021.05
+source activate LAVIS
+
+python blip2_predict_func_concat_pretrain.py

examples/predict_train.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+#SBATCH -J infer_cc
+#SBATCH -p gpu1
+#SBATCH -N 1
+#SBATCH -w node[84]
+#SBATCH --mem 80G
+#SBATCH --gres=gpu:1
+#SBATCH --output=log_predict.out
+#SBATCH --error=log_predict.err
+#SBATCH --cpus-per-task=8
+module load anaconda3/2021.05
+source activate LAVIS
+
+python blip2_predict_func_concat.py