feat: upload BeiT3 API for deployment

.env

@@ -0,0 +1,10 @@
+# PROJECT INFORMATION
+HOST=0.0.0.0
+PORT=7860
+CORS_HEADERS=["*"]
+CORS_ORIGINS=["*"]
+
+DEVICE="cpu" # ["cuda", "cpu"]
+
+INDEX_FILE_PATH="data/faiss-index/index_beit3_L01_to_L20.faiss"
+KEYFRAMES_GROUPS_JSON_PATH="data/config/keyframes_groups_L01_to_L20.json"

Dockerfile

@@ -0,0 +1,39 @@
+FROM python:3.8-slim
+
+RUN apt-get update && \
+    apt-get install git bash wget iputils-ping -y && \
+    apt clean && \
+    rm -rf /var/cache/apt/*
+
+WORKDIR /code
+
+COPY requirements.txt /code/requirements.txt
+
+# PYTHONDONTWRITEBYTECODE=1: Disables the creation of .pyc files (compiled bytecode)
+# PYTHONUNBUFFERED=1: Disables buffering of the standard output stream
+# PYTHONIOENCODING: specifies the encoding to be used for the standard input, output, and error streams
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PYTHONIOENCODING=utf-8
+
+RUN pip install -U pip && \
+    python -m pip install -r /code/requirements.txt
+
+RUN useradd -m -u 1000 user
+
+USER user
+
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user . $HOME/app
+
+# Download index
+# RUN mkdir ./data/faiss-index/ && \
+#     gsutil cp "gs://thangtd1/faiss-index/index_beit3_L01_to_L20.faiss" ./data/faiss-index/
+
+RUN bash $HOME/app/download_models.sh
+
+CMD python ./src/main.py

data/config/keyframes_groups_L01_to_L20.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e66e7bd8837d2b9dc2ce3324e19b3346335dc2437ecfd4279138485bc51422f
+size 26271112

data/faiss-index/index_beit3_L01_to_L20.faiss

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b5a31a0b7815943a055dfd5d39bf2be3ba33447586354cb2fa6505894b3c832
+size 620998701

download_models.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# Test network connection before download
+ping -c 4 8.8.8.8
+if [ $? -gt 0 ]
+then
+    exit 1
+fi
+# Download processor and beit-3 model from provided urls
+wget "https://conversationhub.blob.core.windows.net/beit-share-public/beit3/sentencepiece/beit3.spm?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D" -O ./src/itr/beit3_model/beit3.spm
+wget "https://conversationhub.blob.core.windows.net/beit-share-public/beit3/f30k_retrieval/beit3_base_patch16_384_f30k_retrieval.pth?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D" -O ./src/itr/beit3_model/beit3_base_patch16_384_f30k_retrieval.pth

requirements.dev.txt

@@ -0,0 +1,21 @@
+fastapi==0.103.1
+uvicorn==0.23.2
+pydantic-settings==2.0.3
+
+
+# Models
+torch==2.0.0
+torchvision==0.15.1
+torchscale==0.2.0
+ftfy==6.1.1
+regex
+tqdm==4.66.1
+transformers==4.33.1
+timm==0.4.12
+sentencepiece==0.1.99
+
+# Vector Database
+faiss-cpu
+
+# Project settings
+pre-commit

requirements.txt

@@ -0,0 +1,18 @@
+fastapi==0.103.1
+uvicorn==0.23.2
+pydantic-settings==2.0.3
+
+
+# Models
+torch==2.0.0
+torchvision==0.15.1
+torchscale==0.2.0
+ftfy==6.1.1
+regex
+tqdm==4.66.1
+transformers==4.33.1
+timm==0.4.12
+sentencepiece==0.1.99
+
+# Vector Database
+faiss-cpu

src/config.py

@@ -0,0 +1,28 @@
+from pathlib import Path
+
+from pydantic_settings import BaseSettings
+
+FILE = Path(__file__)
+ROOT = FILE.parent.parent
+
+
+class Settings(BaseSettings):
+    # API SETTINGS
+    HOST: str
+    PORT: int
+    CORS_ORIGINS: list
+    CORS_HEADERS: list
+
+    # MODEL SETTINGS
+    MODEL_NAME: str = "ViT-B/32"
+    DEVICE: str = "cpu"
+
+    # FAISS DATABASE SETTINGS
+    INDEX_FILE_PATH: str
+    KEYFRAMES_GROUPS_JSON_PATH: str
+
+    class Config:
+        env_file = ROOT / ".env"
+
+
+settings = Settings()

src/itr/beit3/README.md

@@ -0,0 +1,28 @@
+# Using BEiT-3 to get text-vision embedding
+
+## For text embedding
+1. Create file ```test_model.py``` inside folder ```itr```.
+2. Using the code follow:
+```
+from beit3_model import Beit3Model
+
+if __name__ == '__main__':
+    vlm = Beit3Model(device='cpu')
+
+    print(vlm.get_embedding('A man who loves a girl.').shape)
+```
+
+## For image embedding
+1. Create file ```test_model.py``` inside folder ```itr```.
+2. Using the code follow:
+```
+from beit3_model import Beit3Model
+from torchvision.datasets.folder import default_loader
+
+if __name__ == '__main__':
+    loader = default_loader
+    image = loader('./path/to/your/image.jpg')
+
+    vlm = Beit3Model(device='cpu')
+    print(vlm.get_embedding(image).shape)
+```

src/itr/beit3_model.py

@@ -0,0 +1,109 @@
+import os
+import torch
+from functools import lru_cache
+from pathlib import Path
+from typing import Union
+
+from . import modeling_finetune, utils
+from PIL import Image
+from timm.data.constants import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
+from timm.models import create_model
+from torchvision import transforms
+from transformers import XLMRobertaTokenizer
+
+# Get current workdir of this file
+CWD = Path(__file__).parent
+print(CWD)
+
+
+class Preprocess:
+    def __init__(self, tokenizer):
+        self.max_len = 64
+        self.input_size = 384
+
+        self.tokenizer = tokenizer
+        self.transform = transforms.Compose(
+            [
+                transforms.Resize((self.input_size, self.input_size), interpolation=3),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD
+                ),
+            ]
+        )
+
+        self.bos_token_id = tokenizer.bos_token_id
+        self.eos_token_id = tokenizer.eos_token_id
+        self.pad_token_id = tokenizer.pad_token_id
+
+    def preprocess(self, input: Union[str, Image.Image]):
+        if isinstance(input, str):
+            tokens = self.tokenizer.tokenize(input)
+            tokens = self.tokenizer.convert_tokens_to_ids(tokens)
+
+            tokens = [self.bos_token_id] + tokens[:] + [self.eos_token_id]
+            num_tokens = len(tokens)
+            padding_mask = [0] * num_tokens + [1] * (self.max_len - num_tokens)
+
+            return (
+                torch.LongTensor(
+                    tokens + [self.pad_token_id] * (self.max_len - num_tokens)
+                ).unsqueeze(0),
+                torch.Tensor(padding_mask).unsqueeze(0),
+                num_tokens,
+            )
+        elif isinstance(input, Image.Image):
+            return self.transform(input).unsqueeze(0)
+        else:
+            raise Exception("Invalid input type")
+
+
+class Beit3Model:
+    def __init__(
+        self,
+        model_name: str = "beit3_base_patch16_384_retrieval",
+        model_path: str = os.path.join(
+            CWD,
+            "beit3_model/beit3_base_patch16_384_f30k_retrieval.pth",
+        ),
+        device: str = "cuda",
+    ):
+        self._load_model(model_name, model_path, device)
+        self.device = device
+
+    # @lru_cache(maxsize=1)
+    def _load_model(self, model_name, model_path, device: str = "cpu"):
+        self.model = create_model(
+            model_name,
+            pretrained=False,
+            drop_path_rate=0.1,
+            vocab_size=64010,
+            checkpoint_activations=False,
+        )
+
+        if model_name:
+            utils.load_model_and_may_interpolate(
+                model_path, self.model, "model|module", ""
+            )
+
+        self.preprocessor = Preprocess(
+            XLMRobertaTokenizer(os.path.join(CWD, "beit3_model/beit3.spm"))
+        )
+        self.model.to(device)
+
+    def get_embedding(self, input: Union[str, Image.Image]):
+        if isinstance(input, str):
+            token_ids, padding_mask, _ = self.preprocessor.preprocess(input)
+
+            _, vector = self.model(
+                text_description=token_ids, padding_mask=padding_mask, only_infer=True
+            )
+            vector = vector.cpu().detach().numpy().astype("float32")
+            return vector
+        elif isinstance(input, Image.Image):
+            image_input = self.preprocessor.preprocess(input)
+            image_input = image_input.to(self.device)
+            vector, _ = self.model(image=image_input, only_infer=True)
+            return vector
+        else:
+            raise Exception("Invalid input type")

src/itr/beit3_model/README.md

@@ -0,0 +1,6 @@
+# BEiT-3 Weight and Sentencepiece models
+
+1. Please download [beit3_base_patch16_384_retrieval.pth](https://conversationhub.blob.core.windows.net/beit-share-public/beit3/f30k_retrieval/beit3_base_patch16_384_f30k_retrieval.pth?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D) and [beit3.spm](https://conversationhub.blob.core.windows.net/beit-share-public/beit3/sentencepiece/beit3.spm?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D) model.
+
+2. Put those 2 model inside this folder
+

src/itr/dtb_cursor.py

@@ -0,0 +1,30 @@
+import json
+import faiss
+import os
+
+from functools import lru_cache
+from pathlib import Path
+
+
+class DatabaseCursor:
+    def __init__(self, index_file_path: str, keyframes_groups_json_path: str):
+        self._load_index(index_file_path)
+        self._load_keyframes_groups_info(keyframes_groups_json_path)
+
+    @lru_cache(maxsize=1)
+    def _load_index(self, index_file_path):
+        self.index = faiss.read_index(index_file_path)
+
+    @lru_cache(maxsize=1)
+    def _load_keyframes_groups_info(self, keyframes_groups_json_path: str):
+        with open(keyframes_groups_json_path) as file:
+            self.keyframes_group_info = json.loads(file.read())
+
+    def kNN_search(self, query_vector: str, topk: int = 10):
+        results = []
+        distances, ids = self.index.search(query_vector, topk)
+        for i in range(len(ids[0])):
+            frame_detail = self.keyframes_group_info[ids[0][i]]
+            frame_detail["distance"] = str(distances[0][i])
+            results.append(frame_detail)
+        return results

src/itr/modeling_finetune.py

@@ -0,0 +1,388 @@
+# --------------------------------------------------------
+# Image as a Foreign Language: BEiT Pretraining for Vision and Vision-Language Tasks (https://arxiv.org/abs/2208.10442)
+# Github source: https://github.com/microsoft/unilm/tree/master/beit3
+# Copyright (c) 2023 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------'
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from . import utils
+from .modeling_utils import BEiT3Wrapper, _get_base_config, _get_large_config
+from timm.models.registry import register_model
+
+
+class TwoLayerMLP(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features,
+        out_features,
+        norm_layer,
+        norm_input=True,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(in_features) if norm_input else nn.Identity()
+        self.dense1 = nn.Linear(in_features, hidden_features)
+        self.norm2 = norm_layer(hidden_features)
+        self.act = nn.GELU()
+        self.dense2 = nn.Linear(hidden_features, out_features)
+
+    def forward(self, x):
+        x = self.norm1(x)
+        x = self.dense1(x)
+        x = self.norm2(x)
+        x = self.act(x)
+        return self.dense2(x)
+
+
+class Pooler(nn.Module):
+    def __init__(self, input_features, output_features, norm_layer):
+        super().__init__()
+        self.norm = norm_layer(input_features)
+        self.dense = nn.Linear(input_features, output_features)
+        self.activation = nn.Tanh()
+
+    def forward(self, x):
+        cls_rep = x[:, 0, :]
+        cls_rep = self.norm(cls_rep)
+        pooled_output = self.dense(cls_rep)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class BEiT3ForVisualReasoning(BEiT3Wrapper):
+    def __init__(self, args, num_classes, norm_layer=nn.LayerNorm, **kwargs):
+        super().__init__(args=args)
+        embed_dim = args.encoder_embed_dim
+        self.head = TwoLayerMLP(
+            in_features=embed_dim * 4,
+            hidden_features=embed_dim * 2,
+            out_features=num_classes,
+            norm_layer=norm_layer,
+        )
+        init_scale = 0.001
+        self.head.apply(self._init_weights)
+        if isinstance(self.head.dense1, nn.Linear):
+            self.head.dense1.weight.data.mul_(init_scale)
+            self.head.dense1.bias.data.mul_(init_scale)
+
+        if isinstance(self.head.dense2, nn.Linear):
+            self.head.dense2.weight.data.mul_(init_scale)
+            self.head.dense2.bias.data.mul_(init_scale)
+
+    def forward(self, image_a, image_b, text_description, padding_mask, **kwargs):
+        bsz, _ = text_description.size()
+
+        vision_input = torch.cat((image_a, image_b), dim=0)
+        language_input = torch.cat((text_description, text_description), dim=0)
+        padding_mask = torch.cat((padding_mask, padding_mask), dim=0)
+
+        outputs = self.beit3(
+            textual_tokens=language_input,
+            visual_tokens=vision_input,
+            text_padding_position=padding_mask,
+        )
+        x = outputs["encoder_out"]
+        multiway_split_position = outputs["multiway_split_position"]
+
+        vision_cls = x[:, 0, :]
+        language_cls = x[:, multiway_split_position, :]
+        cls_rep = torch.cat((vision_cls, language_cls), dim=-1)
+        a, b = torch.split(cls_rep, split_size_or_sections=[bsz, bsz], dim=0)
+        cls_rep = torch.cat((a, b), dim=-1)
+        return self.head(cls_rep)
+
+
+class BEiT3ForImageClassification(BEiT3Wrapper):
+    def __init__(self, args, num_classes, norm_layer=nn.LayerNorm, **kwargs):
+        super().__init__(args=args)
+        embed_dim = args.encoder_embed_dim
+        self.fc_norm = norm_layer(embed_dim)
+        self.head = (
+            nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        )
+
+        self.fc_norm.apply(self._init_weights)
+        self.head.apply(self._init_weights)
+        init_scale = 0.001
+        if isinstance(self.head, nn.Linear):
+            self.head.weight.data.mul_(init_scale)
+            self.head.bias.data.mul_(init_scale)
+
+    def forward(self, image, **kwargs):
+        x = self.beit3(textual_tokens=None, visual_tokens=image)["encoder_out"]
+        t = x[:, 1:, :]
+        cls_x = self.fc_norm(t.mean(1))
+        return self.head(cls_x)
+
+
+class BEiT3ForCaptioning(BEiT3Wrapper):
+    def __init__(self, args, **kwargs):
+        super().__init__(args=args)
+        embed_dim = args.encoder_embed_dim
+        self.mlm_head = nn.Linear(embed_dim, args.vocab_size)
+        self.mlm_head.apply(self._init_weights)
+
+    def forward(
+        self,
+        image,
+        text_ids,
+        padding_mask,
+        language_masked_pos,
+        text_len=None,
+        incremental_state=None,
+        **kwargs
+    ):
+        text_len = text_len if text_len is not None else text_ids.size(1)
+        image_len = self.beit3.vision_embed.num_position_embeddings()
+        max_len = text_len + image_len
+        uni_mask = torch.zeros(
+            (max_len, max_len), dtype=torch.long, device=text_ids.device
+        )
+        i_start, i_end = 0, image_len
+        t_start, t_end = image_len, max_len
+        # triangle mask for caption to caption
+        uni_mask[t_start:t_end, t_start:t_end] = torch.tril(
+            torch.ones(text_len, text_len, dtype=torch.long, device=text_ids.device)
+        )
+        # full attention for caption to image
+        uni_mask[t_start:t_end, i_start:i_end] = 1
+        # full attention for image to image
+        uni_mask[i_start:i_end, i_start:i_end] = 1
+        uni_mask = 1 - uni_mask
+
+        if incremental_state is not None:
+            for idx in range(self.get_num_layers()):
+                if idx not in incremental_state:
+                    incremental_state[idx] = {}
+
+        # for incremental decoding
+        positions = None
+        if image is None:
+            uni_mask = uni_mask[-2:]
+            padding_mask = None
+            # start position (2 (fairseq starts at 2) + cur_position) is equal to text_len
+            positions = (
+                torch.arange(
+                    text_len, text_ids.size(1) + text_len, device=text_ids.device
+                )
+                .long()
+                .unsqueeze(0)
+            )
+
+        outputs = self.beit3(
+            textual_tokens=text_ids,
+            visual_tokens=image,
+            text_padding_position=padding_mask,
+            attn_mask=uni_mask,
+            incremental_state=incremental_state,
+            positions=positions,
+        )
+        if image is not None:
+            text_feats = outputs["encoder_out"][:, image_len:]
+        else:
+            text_feats = outputs["encoder_out"]
+
+        if language_masked_pos is not None:
+            text_feats = text_feats[language_masked_pos.bool()]
+
+        return self.mlm_head(text_feats), incremental_state
+
+
+class BEiT3ForVisualQuestionAnswering(BEiT3Wrapper):
+    def __init__(self, args, num_classes, norm_layer=nn.LayerNorm, **kwargs):
+        super().__init__(args=args)
+        embed_dim = args.encoder_embed_dim
+        self.pooler = Pooler(
+            input_features=embed_dim,
+            output_features=embed_dim,
+            norm_layer=norm_layer,
+        )
+        self.pooler.apply(self._init_weights)
+        self.head = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim * 2),
+            norm_layer(embed_dim * 2),
+            nn.GELU(),
+            nn.Linear(embed_dim * 2, num_classes),
+        )
+        self.head.apply(self._init_weights)
+
+    def forward(self, image, question, padding_mask, **kwargs):
+        outputs = self.beit3(
+            textual_tokens=question,
+            visual_tokens=image,
+            text_padding_position=padding_mask,
+        )
+        x = outputs["encoder_out"]
+        cls_rep = self.pooler(x)
+        return self.head(cls_rep)
+
+
+class BEiT3ForRetrieval(BEiT3Wrapper):
+    def __init__(self, args, **kwargs):
+        super().__init__(args=args)
+        embed_dim = args.encoder_embed_dim
+        self.language_head = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.vision_head = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.language_head.apply(self._init_weights)
+        self.vision_head.apply(self._init_weights)
+        self.criterion = utils.ClipLoss(
+            rank=utils.get_rank(),
+            world_size=utils.get_world_size(),
+        )
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def forward(
+        self,
+        image=None,
+        text_description=None,
+        padding_mask=None,
+        only_infer=False,
+        **kwargs
+    ):
+        if image is not None:
+            outputs = self.beit3(
+                textual_tokens=None,
+                visual_tokens=image,
+                text_padding_position=None,
+            )
+            x = outputs["encoder_out"]
+            vision_cls = self.vision_head(x[:, 0, :])
+            vision_cls = F.normalize(vision_cls, dim=-1)
+        else:
+            vision_cls = None
+
+        if text_description is not None:
+            outputs = self.beit3(
+                textual_tokens=text_description,
+                visual_tokens=None,
+                text_padding_position=padding_mask,
+            )
+            x = outputs["encoder_out"]
+            language_cls = self.language_head(x[:, 0, :])
+            language_cls = F.normalize(language_cls, dim=-1)
+        else:
+            language_cls = None
+
+        if only_infer:
+            return vision_cls, language_cls
+        else:
+            loss, logits_per_image, logits_per_text = self.criterion(
+                vision_cls, language_cls, self.logit_scale.exp()
+            )
+            return loss, vision_cls, language_cls
+
+
+@register_model
+def beit3_base_patch16_224_imageclassification(pretrained=False, **kwargs):
+    args = _get_base_config(**kwargs)
+    args.normalize_output = False
+    model = BEiT3ForImageClassification(args, num_classes=1000, **kwargs)
+    return model
+
+
+@register_model
+def beit3_large_patch16_224_imageclassification(pretrained=False, **kwargs):
+    args = _get_large_config(**kwargs)
+    args.normalize_output = False
+    model = BEiT3ForImageClassification(args, num_classes=1000, **kwargs)
+    return model
+
+
+@register_model
+def beit3_base_patch16_224_nlvr2(pretrained=False, **kwargs):
+    args = _get_base_config(**kwargs)
+    model = BEiT3ForVisualReasoning(args, num_classes=2, **kwargs)
+    return model
+
+
+@register_model
+def beit3_large_patch16_224_nlvr2(pretrained=False, **kwargs):
+    args = _get_large_config(**kwargs)
+    model = BEiT3ForVisualReasoning(args, num_classes=2, **kwargs)
+    return model
+
+
+@register_model
+def beit3_base_patch16_384_vqav2(pretrained=False, **kwargs):
+    args = _get_base_config(img_size=384, **kwargs)
+    args.normalize_output = False
+    model = BEiT3ForVisualQuestionAnswering(args, num_classes=3129, **kwargs)
+    return model
+
+
+@register_model
+def beit3_base_patch16_480_vqav2(pretrained=False, **kwargs):
+    args = _get_base_config(img_size=480, **kwargs)
+    args.normalize_output = False
+    model = BEiT3ForVisualQuestionAnswering(args, num_classes=3129, **kwargs)
+    return model
+
+
+@register_model
+def beit3_large_patch16_384_vqav2(pretrained=False, **kwargs):
+    args = _get_large_config(img_size=384, **kwargs)
+    args.normalize_output = False
+    model = BEiT3ForVisualQuestionAnswering(args, num_classes=3129, **kwargs)
+    return model
+
+
+@register_model
+def beit3_large_patch16_480_vqav2(pretrained=False, **kwargs):
+    args = _get_large_config(img_size=480, **kwargs)
+    args.normalize_output = False
+    model = BEiT3ForVisualQuestionAnswering(args, num_classes=3129, **kwargs)
+    return model
+
+
+@register_model
+def beit3_large_patch16_768_vqav2(pretrained=False, **kwargs):
+    args = _get_large_config(img_size=768, **kwargs)
+    args.normalize_output = False
+    model = BEiT3ForVisualQuestionAnswering(args, num_classes=3129, **kwargs)
+    return model
+
+
+@register_model
+def beit3_base_patch16_224_captioning(pretrained=False, **kwargs):
+    args = _get_base_config(**kwargs)
+    model = BEiT3ForCaptioning(args, **kwargs)
+    return model
+
+
+@register_model
+def beit3_base_patch16_480_captioning(pretrained=False, **kwargs):
+    args = _get_base_config(img_size=480, **kwargs)
+    model = BEiT3ForCaptioning(args, **kwargs)
+    return model
+
+
+@register_model
+def beit3_large_patch16_480_captioning(pretrained=False, **kwargs):
+    args = _get_large_config(img_size=480, **kwargs)
+    model = BEiT3ForCaptioning(args, **kwargs)
+    return model
+
+
+@register_model
+def beit3_base_patch16_224_retrieval(pretrained=False, **kwargs):
+    args = _get_base_config(**kwargs)
+    model = BEiT3ForRetrieval(args, **kwargs)
+    return model
+
+
+@register_model
+def beit3_base_patch16_384_retrieval(pretrained=False, **kwargs):
+    args = _get_base_config(img_size=384, **kwargs)
+    model = BEiT3ForRetrieval(args, **kwargs)
+    return model
+
+
+@register_model
+def beit3_large_patch16_384_retrieval(pretrained=False, **kwargs):
+    args = _get_large_config(img_size=384, **kwargs)
+    model = BEiT3ForRetrieval(args, **kwargs)
+    return model

src/itr/modeling_utils.py

@@ -0,0 +1,108 @@
+# --------------------------------------------------------
+# Image as a Foreign Language: BEiT Pretraining for Vision and Vision-Language Tasks (https://arxiv.org/abs/2208.10442)
+# Github source: https://github.com/microsoft/unilm/tree/master/beit3
+# Copyright (c) 2023 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------'
+
+import math
+
+import torch
+import torch.nn as nn
+from timm.models.layers import trunc_normal_ as __call_trunc_normal_
+from torchscale.architecture.config import EncoderConfig
+from torchscale.model.BEiT3 import BEiT3
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0):
+    __call_trunc_normal_(tensor, mean=mean, std=std, a=-std, b=std)
+
+
+def _get_base_config(
+    img_size=224,
+    patch_size=16,
+    drop_path_rate=0,
+    checkpoint_activations=None,
+    mlp_ratio=4,
+    vocab_size=64010,
+    **kwargs
+):
+    return EncoderConfig(
+        img_size=img_size,
+        patch_size=patch_size,
+        vocab_size=vocab_size,
+        multiway=True,
+        layernorm_embedding=False,
+        normalize_output=True,
+        no_output_layer=True,
+        drop_path_rate=drop_path_rate,
+        encoder_embed_dim=768,
+        encoder_attention_heads=12,
+        encoder_ffn_embed_dim=int(768 * mlp_ratio),
+        encoder_layers=12,
+        checkpoint_activations=checkpoint_activations,
+    )
+
+
+def _get_large_config(
+    img_size=224,
+    patch_size=16,
+    drop_path_rate=0,
+    checkpoint_activations=None,
+    mlp_ratio=4,
+    vocab_size=64010,
+    **kwargs
+):
+    return EncoderConfig(
+        img_size=img_size,
+        patch_size=patch_size,
+        vocab_size=vocab_size,
+        multiway=True,
+        layernorm_embedding=False,
+        normalize_output=True,
+        no_output_layer=True,
+        drop_path_rate=drop_path_rate,
+        encoder_embed_dim=1024,
+        encoder_attention_heads=16,
+        encoder_ffn_embed_dim=int(1024 * mlp_ratio),
+        encoder_layers=24,
+        checkpoint_activations=checkpoint_activations,
+    )
+
+
+class BEiT3Wrapper(nn.Module):
+    def __init__(self, args, **kwargs):
+        super().__init__()
+        self.args = args
+        self.beit3 = BEiT3(args)
+        self.apply(self._init_weights)
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def get_num_layers(self):
+        return self.beit3.encoder.num_layers
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {
+            'pos_embed',
+            'cls_token',
+            'beit3.encoder.embed_positions.A.weight',
+            'beit3.vision_embed.cls_token',
+            'logit_scale',
+        }
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)

src/itr/router.py

@@ -0,0 +1,49 @@
+# from .vlm_model import VisionLanguageModel
+from .beit3_model import Beit3Model
+from fastapi import APIRouter, File, status
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+
+from .dtb_cursor import DatabaseCursor
+
+
+class Item(BaseModel):
+    query_text: str
+    topk: int
+
+
+router = APIRouter()
+
+vectordb_cursor = None
+vlm_model = None
+
+
+def init_vectordb(**kargs):
+    # Singleton pattern
+    global vectordb_cursor
+    if vectordb_cursor is None:
+        vectordb_cursor = DatabaseCursor(**kargs)
+
+
+def init_model(device: str):
+    # Singleton
+    global vlm_model
+    if vlm_model is None:
+        vlm_model = Beit3Model(device=device)
+
+
+@router.post("/retrieval/image-text")
+async def retrieve(item: Item) -> JSONResponse:
+    try:
+        query_vector = vlm_model.get_embedding(input=item.query_text)
+        search_results = vectordb_cursor.kNN_search(query_vector, item.topk)
+    except Exception:
+        return JSONResponse(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            content={"message": "Search error"},
+        )
+
+    return JSONResponse(
+        status_code=status.HTTP_200_OK,
+        content={"message": "success", "details": search_results},
+    )

src/itr/utils.py

@@ -0,0 +1,891 @@
+# --------------------------------------------------------
+# Image as a Foreign Language: BEiT Pretraining for Vision and Vision-Language Tasks (https://arxiv.org/abs/2208.10442)
+# Github source: https://github.com/microsoft/unilm/tree/master/beit3
+# Copyright (c) 2023 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------'
+
+import argparse
+import datetime
+import io
+import json
+import math
+import os
+import time
+from collections import defaultdict, deque
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.utils import get_state_dict
+from torch import inf
+
+
+def bool_flag(s):
+    """
+    Parse boolean arguments from the command line.
+    """
+    FALSY_STRINGS = {"off", "false", "0"}
+    TRUTHY_STRINGS = {"on", "true", "1"}
+    if s.lower() in FALSY_STRINGS:
+        return False
+    elif s.lower() in TRUTHY_STRINGS:
+        return True
+    else:
+        raise argparse.ArgumentTypeError("invalid value for a boolean flag")
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value,
+        )
+
+
+class MetricLogger:
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{attr}'")
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(f"{name}: {str(meter)}")
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}',
+        ]
+        if torch.cuda.is_available():
+            log_msg.append('max mem: {memory:.0f}')
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(
+            '{} Total time: {} ({:.4f} s / it)'.format(
+                header, total_time_str, total_time / len(iterable)
+            )
+        )
+
+
+def _load_checkpoint_for_ema(model_ema, checkpoint):
+    """
+    Workaround for ModelEma._load_checkpoint to accept an already-loaded object
+    """
+    mem_file = io.BytesIO()
+    torch.save(checkpoint, mem_file)
+    mem_file.seek(0)
+    model_ema._load_checkpoint(mem_file)
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def _get_rank_env():
+    if "RANK" in os.environ:
+        return int(os.environ["RANK"])
+    else:
+        return int(os.environ['OMPI_COMM_WORLD_RANK'])
+
+
+def _get_local_rank_env():
+    if "LOCAL_RANK" in os.environ:
+        return int(os.environ["LOCAL_RANK"])
+    else:
+        return int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+
+
+def _get_world_size_env():
+    if "WORLD_SIZE" in os.environ:
+        return int(os.environ["WORLD_SIZE"])
+    else:
+        return int(os.environ['OMPI_COMM_WORLD_SIZE'])
+
+
+# The implementation code is modified from DeiT (https://github.com/facebookresearch/deit.git)
+def init_distributed_mode(args):
+    if args.dist_on_itp:
+        args.rank = _get_rank_env()
+        args.world_size = _get_world_size_env()  # int(os.environ['OMPI_COMM_WORLD_SIZE'])
+        args.gpu = _get_local_rank_env()
+        args.dist_url = "tcp://{}:{}".format(os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
+        os.environ['LOCAL_RANK'] = str(args.gpu)
+        os.environ['RANK'] = str(args.rank)
+        os.environ['WORLD_SIZE'] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print(
+        f'| distributed init (rank {args.rank}): {args.dist_url}, gpu {args.gpu}',
+        flush=True,
+    )
+    torch.distributed.init_process_group(
+        backend=args.dist_backend,
+        init_method=args.dist_url,
+        world_size=args.world_size,
+        rank=args.rank,
+        timeout=datetime.timedelta(0, 7200),
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+def load_state_dict(model, state_dict, prefix='', ignore_missing="relative_position_index"):
+    missing_keys = []
+    unexpected_keys = []
+    error_msgs = []
+    # copy state_dict so _load_from_state_dict can modify it
+    metadata = getattr(state_dict, '_metadata', None)
+    state_dict = state_dict.copy()
+    if metadata is not None:
+        state_dict._metadata = metadata
+
+    def load(module, prefix=''):
+        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+        module._load_from_state_dict(
+            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs
+        )
+        for name, child in module._modules.items():
+            if child is not None:
+                load(child, prefix + name + '.')
+
+    load(model, prefix=prefix)
+
+    warn_missing_keys = []
+    ignore_missing_keys = []
+    for key in missing_keys:
+        keep_flag = True
+        for ignore_key in ignore_missing.split('|'):
+            if ignore_key in key:
+                keep_flag = False
+                break
+        if keep_flag:
+            warn_missing_keys.append(key)
+        else:
+            ignore_missing_keys.append(key)
+
+    missing_keys = warn_missing_keys
+
+    if len(missing_keys) > 0:
+        print(
+            "Weights of {} not initialized from pretrained model: {}".format(
+                model.__class__.__name__, missing_keys
+            )
+        )
+    if len(unexpected_keys) > 0:
+        print(
+            "Weights from pretrained model not used in {}: {}".format(
+                model.__class__.__name__, unexpected_keys
+            )
+        )
+    if len(ignore_missing_keys) > 0:
+        print(
+            "Ignored weights of {} not initialized from pretrained model: {}".format(
+                model.__class__.__name__, ignore_missing_keys
+            )
+        )
+    if len(error_msgs) > 0:
+        print('\n'.join(error_msgs))
+
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(
+        self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True
+    ):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(
+                    optimizer
+                )  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.0)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(
+            torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]),
+            norm_type,
+        )
+    return total_norm
+
+
+def cosine_scheduler(
+    base_value,
+    final_value,
+    epochs,
+    niter_per_ep,
+    warmup_epochs=0,
+    start_warmup_value=0,
+    warmup_steps=-1,
+    sched_type="cos",
+):
+    warmup_schedule = np.array([])
+    warmup_iters = warmup_epochs * niter_per_ep
+    if warmup_steps > 0:
+        warmup_iters = warmup_steps
+    print("Set warmup steps = %d" % warmup_iters)
+    if warmup_epochs > 0:
+        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+    if sched_type == "cos":
+        iters = np.arange(epochs * niter_per_ep - warmup_iters)
+        schedule = np.array(
+            [
+                final_value
+                + 0.5 * (base_value - final_value) * (1 + math.cos(math.pi * i / (len(iters))))
+                for i in iters
+            ]
+        )
+    elif sched_type == "linear":
+        schedule = np.linspace(base_value, final_value, epochs * niter_per_ep - warmup_iters)
+    else:
+        raise NotImplementedError()
+
+    schedule = np.concatenate((warmup_schedule, schedule))
+
+    assert len(schedule) == epochs * niter_per_ep
+    return schedule
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler, model_ema=None):
+    output_dir = Path(args.output_dir)
+    if loss_scaler is not None:
+        checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            if model_ema is not None:
+                to_save['model_ema'] = get_state_dict(model_ema)
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch, "args": args}
+        if model_ema is not None:
+            client_state['model_ema'] = get_state_dict(model_ema)
+        model.save_checkpoint(
+            save_dir=args.output_dir, tag="checkpoint-%s" % epoch, client_state=client_state
+        )
+
+
+def auto_load_model(args, model, model_without_ddp, optimizer, loss_scaler, model_ema=None):
+    output_dir = Path(args.output_dir)
+    if loss_scaler is not None:
+        # torch.amp
+        if args.auto_resume and len(args.resume) == 0:
+            import glob
+
+            all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*.pth'))
+            latest_ckpt = -1
+            for ckpt in all_checkpoints:
+                t = ckpt.split('-')[-1].split('.')[0]
+                if t.isdigit():
+                    latest_ckpt = max(int(t), latest_ckpt)
+            if latest_ckpt >= 0:
+                args.resume = os.path.join(output_dir, 'checkpoint-%d.pth' % latest_ckpt)
+            print("Auto resume checkpoint: %s" % args.resume)
+
+        if args.resume:
+            if args.resume.startswith('https'):
+                checkpoint = torch.hub.load_state_dict_from_url(
+                    args.resume, map_location='cpu', check_hash=True
+                )
+            else:
+                checkpoint = torch.load(args.resume, map_location='cpu')
+            model_without_ddp.load_state_dict(checkpoint['model'])
+            print("Resume checkpoint %s" % args.resume)
+            if 'optimizer' in checkpoint and 'epoch' in checkpoint:
+                optimizer.load_state_dict(checkpoint['optimizer'])
+                args.start_epoch = checkpoint['epoch'] + 1
+                if hasattr(args, 'model_ema') and args.model_ema:
+                    _load_checkpoint_for_ema(model_ema, checkpoint['model_ema'])
+                if 'scaler' in checkpoint:
+                    loss_scaler.load_state_dict(checkpoint['scaler'])
+                print("With optim & sched!")
+    else:
+        # deepspeed, only support '--auto_resume'.
+        if args.auto_resume:
+            import glob
+
+            all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*'))
+            latest_ckpt = -1
+            for ckpt in all_checkpoints:
+                t = ckpt.split('-')[-1].split('.')[0]
+                if t.isdigit():
+                    latest_ckpt = max(int(t), latest_ckpt)
+            if latest_ckpt >= 0:
+                args.resume = os.path.join(output_dir, 'checkpoint-%d' % latest_ckpt)
+                print("Auto resume checkpoint: %d" % latest_ckpt)
+                _, client_states = model.load_checkpoint(
+                    args.output_dir, tag='checkpoint-%d' % latest_ckpt
+                )
+                args.start_epoch = client_states['epoch'] + 1
+                if model_ema is not None:
+                    if args.model_ema:
+                        _load_checkpoint_for_ema(model_ema, client_states['model_ema'])
+
+
+# The implementation code is modified from DeiT (https://github.com/facebookresearch/deit.git)
+def load_model_and_may_interpolate(ckpt_path, model, model_key, model_prefix):
+    if ckpt_path.startswith('https'):
+        checkpoint = torch.hub.load_state_dict_from_url(
+            ckpt_path, map_location='cpu', check_hash=True
+        )
+    else:
+        checkpoint = torch.load(ckpt_path, map_location='cpu')
+
+    print("Load ckpt from %s" % ckpt_path)
+    checkpoint_model = None
+    for model_key in model_key.split('|'):
+        if model_key in checkpoint:
+            checkpoint_model = checkpoint[model_key]
+            print("Load state_dict by model_key = %s" % model_key)
+            break
+
+    if checkpoint_model is None:
+        checkpoint_model = checkpoint
+
+    state_dict = model.state_dict()
+    for k in ['head.weight', 'head.bias']:
+        if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
+            print(f"Removing key {k} from pretrained checkpoint")
+            del checkpoint_model[k]
+
+    # interpolate position embedding
+    for pos_embed_key in (
+        "vision_pos_embed",
+        "pos_embed",
+        "beit3.encoder.embed_positions.A.weight",
+    ):
+        if pos_embed_key in checkpoint_model:
+            pos_embed_checkpoint = checkpoint_model[pos_embed_key]
+            embedding_size = pos_embed_checkpoint.shape[-1]
+            if pos_embed_key == "beit3.encoder.embed_positions.A.weight":
+                # being consistent with Fairseq, which starts from 2 for position embedding
+                torchscale_model = True
+                num_patches = model.beit3.vision_embed.num_patches
+                num_extra_tokens = (
+                    model.beit3.vision_embed.num_position_embeddings() + 2 - num_patches
+                )
+            else:
+                torchscale_model = False
+                num_patches = model.patch_embed.num_patches
+                num_extra_tokens = getattr(model, pos_embed_key).shape[-2] - num_patches
+            # height (== width) for the checkpoint position embedding
+            orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+            # height (== width) for the new position embedding
+            new_size = int(num_patches**0.5)
+            # class_token and dist_token are kept unchanged
+            if orig_size != new_size:
+                print(
+                    "Position interpolate from %dx%d to %dx%d"
+                    % (orig_size, orig_size, new_size, new_size)
+                )
+                if torchscale_model:
+                    extra_tokens = pos_embed_checkpoint[:num_extra_tokens].unsqueeze(0)
+                    # only the position tokens are interpolated
+                    pos_tokens = pos_embed_checkpoint[num_extra_tokens:]
+                else:
+                    extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+                    # only the position tokens are interpolated
+                    pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+                pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(
+                    0, 3, 1, 2
+                )
+                pos_tokens = torch.nn.functional.interpolate(
+                    pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False
+                )
+                pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+                new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+                if torchscale_model:
+                    new_pos_embed = new_pos_embed.squeeze(0)
+                checkpoint_model[pos_embed_key] = new_pos_embed
+
+    load_state_dict(model, checkpoint_model, prefix=model_prefix)
+
+
+def create_ds_config(args):
+    args.deepspeed_config = os.path.join(args.output_dir, "deepspeed_config.json")
+    with open(args.deepspeed_config, mode="w") as writer:
+        ds_config = {
+            "train_batch_size": args.batch_size * args.update_freq * get_world_size(),
+            "train_micro_batch_size_per_gpu": args.batch_size,
+            "steps_per_print": 1000,
+            "optimizer": {
+                "type": "Adam",
+                "adam_w_mode": True,
+                "params": {
+                    "lr": args.lr,
+                    "weight_decay": args.weight_decay,
+                    "bias_correction": True,
+                    "betas": [args.opt_betas[0], args.opt_betas[1]],
+                    "eps": args.opt_eps,
+                },
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 0,
+                "initial_scale_power": getattr(args, "initial_scale_power", 12),
+                "loss_scale_window": 1000,
+                "hysteresis": 2,
+                "min_loss_scale": 1,
+            },
+            "amp": {"enabled": False, "opt_level": "O2"},
+        }
+
+        if args.clip_grad is not None:
+            ds_config.update({'gradient_clipping': args.clip_grad})
+
+        if args.zero_stage == 1:
+            ds_config.update(
+                {"zero_optimization": {"stage": args.zero_stage, "reduce_bucket_size": 5e8}}
+            )
+        elif args.zero_stage > 1:
+            raise NotImplementedError()
+
+        writer.write(json.dumps(ds_config, indent=2))
+
+
+def merge_batch_tensors_by_dict_key(batch):
+    batch_tensors = {}
+    for tensor_key in batch[0]:
+        if isinstance(batch[0][tensor_key], torch.Tensor):
+            batch_tensors[tensor_key] = torch.stack([d[tensor_key] for d in batch])
+        else:
+            batch_tensors[tensor_key] = torch.tensor(
+                [d[tensor_key] for d in batch], dtype=torch.long
+            )
+    return batch_tensors
+
+
+def get_loss_scale_for_deepspeed(model):
+    optimizer = model.optimizer
+    loss_scale = None
+    if hasattr(optimizer, 'loss_scale'):
+        loss_scale = optimizer.loss_scale
+    elif hasattr(optimizer, 'cur_scale'):
+        loss_scale = optimizer.cur_scale
+    return loss_scale
+
+
+class GatherLayer(torch.autograd.Function):
+    """
+    Gather tensors from all workers with support for backward propagation:
+    This implementation does not cut the gradients as torch.distributed.all_gather does.
+    """
+
+    @staticmethod
+    def forward(ctx, x):
+        output = [torch.zeros_like(x) for _ in range(dist.get_world_size())]
+        dist.all_gather(output, x)
+        return tuple(output)
+
+    @staticmethod
+    def backward(ctx, *grads):
+        all_gradients = torch.stack(grads)
+        dist.all_reduce(all_gradients)
+        return all_gradients[dist.get_rank()]
+
+
+def gather_features(
+    image_features,
+    text_features,
+):
+    gathered_image_features = GatherLayer.apply(image_features)
+    gathered_text_features = GatherLayer.apply(text_features)
+    all_image_features = torch.cat(gathered_image_features)
+    all_text_features = torch.cat(gathered_text_features)
+
+    return all_image_features, all_text_features
+
+
+# The implementation code is modified from open_clip (https://github.com/mlfoundations/open_clip.git)
+class ClipLoss(nn.Module):
+    def __init__(
+        self,
+        cache_labels=False,
+        rank=0,
+        world_size=1,
+    ):
+        super().__init__()
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def forward(self, image_features, text_features, logit_scale):
+        device = image_features.device
+        if self.world_size > 1:
+            all_image_features, all_text_features = gather_features(image_features, text_features)
+
+            logits_per_image = logit_scale * image_features @ all_text_features.T
+            logits_per_text = logit_scale * text_features @ all_image_features.T
+        else:
+            logits_per_image = logit_scale * image_features @ text_features.T
+            logits_per_text = logit_scale * text_features @ image_features.T
+
+        # calculated ground-truth and cache if enabled
+        num_logits = logits_per_image.shape[0]
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if self.world_size > 1:
+                labels = labels + num_logits * self.rank
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+
+        total_loss = (
+            F.cross_entropy(logits_per_image, labels) + F.cross_entropy(logits_per_text, labels)
+        ) / 2
+        return total_loss, logits_per_image, logits_per_text
+
+
+def write_result_to_jsonl(test_stats, result_file):
+    with open(result_file, mode="w", encoding="utf-8") as writer:
+        writer.write(json.dumps(test_stats, indent=None))
+
+
+def read_result_from_jsonl(result_file):
+    with open(result_file, encoding="utf-8") as reader:
+        return json.load(reader)
+
+
+class BertCaptioningLoss(nn.Module):
+    def __init__(self, label_smoothing, drop_worst_ratio, drop_worst_after):
+        super().__init__()
+        self.label_smoothing = label_smoothing
+        self.drop_worst_ratio = drop_worst_ratio
+        self.drop_worst_after = drop_worst_after
+        self.log_soft = nn.LogSoftmax(dim=1)
+        self.kl = nn.KLDivLoss(reduction='none')
+        self.iter = 0
+
+    def forward(self, logits, target, iter):
+        eps = self.label_smoothing
+        n_class = logits.size(1)
+        one_hot = torch.zeros_like(logits).scatter(1, target.view(-1, 1), 1)
+        one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
+        log_prb = self.log_soft(logits)
+        loss = self.kl(log_prb, one_hot).sum(1)
+
+        if self.drop_worst_ratio > 0 and iter > self.drop_worst_after:
+            loss, _ = torch.topk(
+                loss, k=int(loss.shape[0] * (1 - self.drop_worst_ratio)), largest=False
+            )
+        loss = loss.mean()
+
+        return loss
+
+
+class BeamHypotheses:
+    def __init__(self, n_hyp, max_length, length_penalty, early_stopping):
+        """
+        Initialize n-best list of hypotheses.
+        """
+        self.max_length = max_length - 1  # ignoring bos_token
+        self.length_penalty = length_penalty
+        self.early_stopping = early_stopping
+        self.n_hyp = n_hyp
+        self.hyp = []
+        self.worst_score = 1e9
+
+    def __len__(self):
+        """
+        Number of hypotheses in the list.
+        """
+        return len(self.hyp)
+
+    def add(self, hyp, sum_logprobs):
+        """
+        Add a new hypothesis to the list.
+        """
+        score = sum_logprobs / len(hyp) ** self.length_penalty
+        if len(self) < self.n_hyp or score > self.worst_score:
+            self.hyp.append((score, hyp))
+            if len(self) > self.n_hyp:
+                sorted_scores = sorted([(s, idx) for idx, (s, _) in enumerate(self.hyp)])
+                del self.hyp[sorted_scores[0][1]]
+                self.worst_score = sorted_scores[1][0]
+            else:
+                self.worst_score = min(score, self.worst_score)
+
+    def is_done(self, best_sum_logprobs):
+        """
+        If there are enough hypotheses and that none of the hypotheses being generated
+        can become better than the worst one in the heap, then we are done with this sentence.
+        """
+        if len(self) < self.n_hyp:
+            return False
+        elif self.early_stopping:
+            return True
+        else:
+            return self.worst_score >= best_sum_logprobs / self.max_length**self.length_penalty
+
+
+def dump_predictions(args, result, file_suffix):
+    global_rank = get_rank()
+    jsons = None
+    if global_rank >= 0:
+        output_file = os.path.join(args.task_cache_path, f"submit_{global_rank}_{file_suffix}.json")
+        with open(output_file, "w") as fp:
+            json.dump(result, fp, indent=2)
+        torch.distributed.barrier()
+
+        if global_rank == 0:
+            world_size = get_world_size()
+            jsons = []
+            for i in range(world_size):
+                each_file = os.path.join(args.task_cache_path, f"submit_{i}_{file_suffix}.json")
+                with open(each_file) as fp:
+                    jsons += json.load(fp)
+
+            new_jsons = []
+            res_dict = dict()
+            if args.task in ["coco_captioning", "nocaps"]:
+                qid_key = "image_id"
+            else:
+                # for VQAv2
+                qid_key = "question_id"
+            for item in jsons:
+                if item[qid_key] in res_dict:
+                    continue
+                new_jsons.append(item)
+                res_dict[item[qid_key]] = item
+            jsons = new_jsons
+
+        torch.distributed.barrier()
+        os.remove(output_file)
+    else:
+        jsons = result
+
+    result_file = os.path.join(args.output_dir, f"submit_{file_suffix}.json")
+    if jsons is not None:
+        with open(result_file, "w") as fp:
+            json.dump(jsons, fp, indent=2)
+        print("Infer %d examples into %s" % (len(jsons), result_file))
+    return result_file

src/itr/vlm_model.py

@@ -0,0 +1,31 @@
+from functools import lru_cache
+from typing import Union
+
+import clip
+from PIL import Image
+
+
+class VisionLanguageModel:
+    def __init__(self, model_name: str = "ViT-B/32", device: str = "cuda"):
+        self._load_model(model_name, device)
+        self.device = device
+
+    @lru_cache(maxsize=1)
+    def _load_model(self, model_name, device: str = "cpu"):
+        self.model, self.processor = clip.load(model_name, device=device)
+
+    def get_embedding(self, input: Union[str, Image.Image]):
+        if isinstance(input, str):
+            tokens = clip.tokenize(input).to(self.device)
+            vector = self.model.encode_text(tokens)
+            vector /= vector.norm(dim=-1, keepdim=True)
+            vector = vector.cpu().detach().numpy().astype("float32")
+            return vector
+        elif isinstance(input, Image.Image):
+            image_input = self.preprocess(input).unsqueeze(0).to(self.device)
+            vector = self.model.encode_image(image_input)
+            vector /= vector.norm(dim=-1, keepdim=True)
+            vector = vector.cpu().detach().numpy().astype("float32")
+            return vector
+        else:
+            raise Exception("Invalid input type")

src/main.py

@@ -0,0 +1,67 @@
+import torch
+import os
+from config import settings
+from fastapi import FastAPI, Request, status
+from fastapi.exceptions import RequestValidationError
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, RedirectResponse
+from itr.router import init_model, init_vectordb
+from itr.router import router as router
+from pathlib import Path
+
+app = FastAPI(title="[BeiT-3] Text-to-image Retrieval API")
+
+SERVICE_ROOT = Path(__file__).parent.parent
+
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=settings.CORS_ORIGINS,
+    allow_headers=settings.CORS_HEADERS,
+    allow_credentials=True,
+    allow_methods=["*"],
+)
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    # Get the original 'detail' list of errors
+    details = exc.errors()
+    error_details = []
+
+    for error in details:
+        error_details.append({"error": f"{error['msg']} {str(error['loc'])}"})
+    return JSONResponse(content={"message": error_details})
+
+
+@app.on_event("startup")
+async def startup_event():
+    init_vectordb(
+        index_file_path=os.path.join(SERVICE_ROOT, settings.INDEX_FILE_PATH),
+        keyframes_groups_json_path=settings.KEYFRAMES_GROUPS_JSON_PATH,
+    )
+    device = (
+        "cuda" if settings.DEVICE == "cuda" and torch.cuda.is_available() else "cpu"
+    )
+    init_model(device=device)
+
+
+@app.get("/", include_in_schema=False)
+async def root() -> None:
+    return RedirectResponse("/docs")
+
+
+@app.get("/health", status_code=status.HTTP_200_OK, tags=["health"])
+async def perform_healthcheck() -> None:
+    return JSONResponse(content={"message": "success"})
+
+
+app.include_router(router)
+
+
+# Start API
+if __name__ == "__main__":
+    print(os.listdir(os.path.join(SERVICE_ROOT, "data/faiss-index/")))
+    import uvicorn
+
+    uvicorn.run("main:app", host=settings.HOST, port=settings.PORT, reload=True)