Use remote implementation instead

config.json

@@ -6,8 +6,8 @@
     "JinaCLIPModel"
   ],
   "auto_map": {
-    "AutoConfig": "configuration_clip.JinaCLIPConfig",
-    "AutoModel": "modeling_clip.JinaCLIPModel"
+    "AutoConfig": "tomaarsen/jina-clip-implementation-st--configuration_clip.JinaCLIPConfig",
+    "AutoModel": "tomaarsen/jina-clip-implementation-st--modeling_clip.JinaCLIPModel"
   },
   "initializer_factor": 1.0,
   "logit_scale_init_value": 2.6592,

configuration_clip.py

@@ -1,304 +0,0 @@
-# coding=utf-8
-#
-# Code mainly copied from:
-# https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/configuration_clip.py
-# and adjusted for Jina CLIP
-
-import os
-from copy import deepcopy
-from typing import Any, Dict, Optional, Union
-
-from transformers import PretrainedConfig, logging
-
-logger = logging.get_logger(__name__)
-
-
-""" Jina CLIP model configuration """
-
-
-class JinaCLIPTextConfig(PretrainedConfig):
-    model_type = 'jina_clip_text'
-
-    def __init__(
-        self,
-        embed_dim: int = 768,
-        hf_model_name_or_path: str = 'jinaai/jina-bert-flash-implementation',
-        hf_model_config_kwargs: Optional[Dict[str, Any]] = None,
-        pooler_type: Optional[str] = None,
-        proj_type: Optional[str] = None,
-        proj_bias: bool = False,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        self.embed_dim = embed_dim
-        self.hf_model_name_or_path = hf_model_name_or_path
-        self.hf_model_config_kwargs = hf_model_config_kwargs or {}
-        self.pooler_type = pooler_type
-        self.proj_type = proj_type
-        self.proj_bias = proj_bias
-
-    @classmethod
-    def from_pretrained(
-        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
-    ) -> 'PretrainedConfig':
-        cls._set_token_in_kwargs(kwargs)
-
-        configdict, kwargs = cls.get_config_dict(
-            pretrained_model_name_or_path, **kwargs
-        )
-
-        # get the text config dict if we are loading from JinaCLIPConfig
-        if configdict.get('model_type') == 'jina_clip':
-            configdict = configdict['text_config']
-
-        if (
-            'model_type' in configdict
-            and hasattr(cls, 'model_type')
-            and configdict['model_type'] != cls.model_type
-        ):
-            logger.warning(
-                f'You are using a model of type {configdict["model_type"]} to '
-                f'instantiate a model of type {cls.model_type}. This is not supported '
-                'for all configurations of models and can yield errors.'
-            )
-
-        return cls.from_dict(configdict, **kwargs)
-
-
-class JinaCLIPVisionConfig(PretrainedConfig):
-    model_type = 'jina_clip_vision'
-
-    def __init__(
-        self,
-        embed_dim: int = 768,
-        width: int = 768,
-        image_size: int = 224,
-        patch_size: int = 16,
-        layers: int = 12,
-        head_width: int = 64,
-        mlp_ratio: float = 4.0,
-        ls_init_value: Optional[float] = None,
-        patch_dropout: float = 0.0,
-        qkv_bias: bool = True,
-        fused_layer_norm: bool = False,
-        x_attention: bool = False,
-        post_norm: bool = False,
-        rope_embeddings: bool = False,
-        pt_hw_seq_len: int = 16,
-        intp_freq: bool = False,
-        naive_swiglu: bool = False,
-        subln: bool = False,
-        drop_path_rate: float = 0.0,
-        proj_type: Optional[str] = None,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        self.layers = layers
-        self.embed_dim = embed_dim
-        self.width = width
-        self.head_width = head_width
-        self.mlp_ratio = mlp_ratio
-        self.image_size = image_size
-        self.patch_size = patch_size
-        self.ls_init_value = ls_init_value
-        self.patch_dropout = patch_dropout
-        self.qkv_bias = qkv_bias
-        self.fused_layer_norm = fused_layer_norm
-        self.x_attention = x_attention
-        self.post_norm = post_norm
-        self.rope_embeddings = rope_embeddings
-        self.pt_hw_seq_len = pt_hw_seq_len
-        self.intp_freq = intp_freq
-        self.naive_swiglu = naive_swiglu
-        self.subln = subln
-        self.drop_path_rate = drop_path_rate
-        self.proj_type = proj_type
-
-    @classmethod
-    def from_pretrained(
-        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
-    ) -> 'PretrainedConfig':
-        cls._set_token_in_kwargs(kwargs)
-
-        configdict, kwargs = cls.get_config_dict(
-            pretrained_model_name_or_path, **kwargs
-        )
-
-        # get the vision config dict if we are loading from JinaCLIPConfig
-        if configdict.get('model_type') == 'jina_clip':
-            configdict = configdict['vision_config']
-
-        if (
-            'model_type' in configdict
-            and hasattr(cls, 'model_type')
-            and configdict['model_type'] != cls.model_type
-        ):
-            logger.warning(
-                f'You are using a model of type {configdict["model_type"]} to '
-                f'instantiate a model of type {cls.model_type}. This is not supported '
-                'for all configurations of models and can yield errors.'
-            )
-
-        return cls.from_dict(configdict, **kwargs)
-
-
-class JinaCLIPConfig(PretrainedConfig):
-    model_type = 'jina_clip'
-    is_composition = True
-
-    def __init__(
-        self,
-        text_config: Optional[Dict] = None,
-        vision_config: Optional[Dict] = None,
-        add_projections: bool = False,
-        projection_dim: int = 768,
-        logit_scale_init_value: float = 2.6592,
-        use_text_flash_attn: Optional[bool] = None,
-        use_vision_xformers: Optional[bool] = None,
-        **kwargs,
-    ):
-        # If `_config_dict` exist, we use them for the backward compatibility.
-        # We pop out these 2 attributes before calling `super().__init__` to avoid
-        # them being saved (which causes a lot of confusion!).
-
-        text_config_dict: Optional[Dict] = kwargs.pop('text_config_dict', None)
-        vision_config_dict: Optional[Dict] = kwargs.pop('vision_config_dict', None)
-        self.use_text_flash_attn = use_text_flash_attn
-        self.use_vision_xformers = use_vision_xformers
-
-        super().__init__(**kwargs)
-
-        if text_config_dict is not None:
-            if text_config is None:
-                text_config = {}
-
-            # This is the complete result when using `text_config_dict`.
-            _text_config_dict = JinaCLIPTextConfig(**text_config_dict).to_dict()
-
-            # Give a warning if the values exist in both `_text_config_dict` and
-            # `text_config` but being different.
-            for key, value in _text_config_dict.items():
-                if (
-                    key in text_config
-                    and value != text_config[key]
-                    and key not in ['transformers_version']
-                ):
-                    # If specified in `text_config_dict`
-                    if key in text_config_dict:
-                        message = (
-                            f'`{key}` is found in both `text_config_dict` and '
-                            f'`text_config` but with different values. '
-                            f'The value `text_config_dict["{key}"]` will be used '
-                            f'instead.'
-                        )
-                    # If inferred from default argument values (
-                    # just to be super careful)
-                    else:
-                        message = (
-                            f'`text_config_dict` is provided which will be used to '
-                            f'initialize `JinaCLIPTextConfig`. The '
-                            f'value `text_config["{key}"]` will be overriden.'
-                        )
-                    logger.info(message)
-
-            # Update all values in `text_config` with the ones in `_text_config_dict`.
-            text_config.update(_text_config_dict)
-
-        if vision_config_dict is not None:
-            if vision_config is None:
-                vision_config = {}
-
-            # This is the complete result when using `vision_config_dict`.
-            _vision_config_dict = JinaCLIPVisionConfig(**vision_config_dict).to_dict()
-            # convert keys to string instead of integer
-            if 'id2label' in _vision_config_dict:
-                _vision_config_dict['id2label'] = {
-                    str(key): value
-                    for key, value in _vision_config_dict['id2label'].items()
-                }
-
-            # Give a warning if the values exist in both `_vision_config_dict`
-            # and `vision_config` but being different.
-            for key, value in _vision_config_dict.items():
-                if (
-                    key in vision_config
-                    and value != vision_config[key]
-                    and key not in ['transformers_version']
-                ):
-                    # If specified in `vision_config_dict`
-                    if key in vision_config_dict:
-                        message = (
-                            f'`{key}` is found in both `vision_config_dict` and '
-                            f'`vision_config` but with different '
-                            f'values. The value `vision_config_dict["{key}"]` will '
-                            f'be used instead.'
-                        )
-                    # If inferred from default argument values
-                    # (just to be super careful)
-                    else:
-                        message = (
-                            f'`vision_config_dict` is provided which will be used to '
-                            f'initialize `JinaCLIPVisionConfig`. '
-                            f'The value `vision_config["{key}"]` will be overriden.'
-                        )
-                    logger.info(message)
-
-            # Update all values in `vision_config` with the ones in
-            # `_vision_config_dict`.
-            vision_config.update(_vision_config_dict)
-
-        if text_config is None:
-            text_config = {}
-            logger.info(
-                '`text_config` is `None`. Initializing the `JinaCLIPTextConfig` with '
-                'default values.'
-            )
-
-        if vision_config is None:
-            vision_config = {}
-            logger.info(
-                '`vision_config` is `None`. initializing the `JinaCLIPVisionConfig` '
-                'with default values.'
-            )
-
-        self.text_config = JinaCLIPTextConfig(**text_config)
-        self.vision_config = JinaCLIPVisionConfig(**vision_config)
-
-        self.add_projections = add_projections
-        self.projection_dim = projection_dim
-        self.logit_scale_init_value = logit_scale_init_value
-        self.initializer_factor = 1.0
-
-        if not self.add_projections:
-            if self.text_config.embed_dim != self.vision_config.embed_dim:
-                raise ValueError(
-                    'When projections are disabled (`add_projections=False`), text '
-                    'and vision towers need to have the same embedding dimensionality. '
-                    f'Currently text embedding dim is {self.text_config.embed_dim} != '
-                    f'{self.vision_config.embed_dim} of the vision tower. '
-                    'Either set the same output dim for both towers, or enable '
-                    'projections with `add_projections=True`.'
-                )
-
-    @classmethod
-    def from_text_vision_configs(
-        cls,
-        text_config: JinaCLIPTextConfig,
-        vision_config: JinaCLIPVisionConfig,
-        **kwargs,
-    ):
-        return cls(
-            text_config=text_config.to_dict(),
-            vision_config=vision_config.to_dict(),
-            projection_dim=text_config.projection_dim,
-            **kwargs,
-        )
-
-    def to_dict(self):
-        output = deepcopy(self.__dict__)
-        output['text_config'] = self.text_config.to_dict()
-        output['vision_config'] = self.vision_config.to_dict()
-        output['model_type'] = self.__class__.model_type
-        return output

custom_st.py

@@ -1,174 +0,0 @@
-import base64
-from io import BytesIO
-import json
-import os
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-from .custom_st_2 import OtherClass
-import requests
-import torch
-from torch import nn
-from transformers import AutoConfig, AutoModel, AutoTokenizer, AutoImageProcessor
-from PIL import Image
-
-OtherClass()
-
-class Transformer(nn.Module):
-    """Huggingface AutoModel to generate token embeddings.
-    Loads the correct class, e.g. BERT / RoBERTa etc.
-
-    Args:
-        model_name_or_path: Huggingface models name
-            (https://huggingface.co/models)
-        max_seq_length: Truncate any inputs longer than max_seq_length
-        model_args: Keyword arguments passed to the Huggingface
-            Transformers model
-        tokenizer_args: Keyword arguments passed to the Huggingface
-            Transformers tokenizer
-        config_args: Keyword arguments passed to the Huggingface
-            Transformers config
-        cache_dir: Cache dir for Huggingface Transformers to store/load
-            models
-        do_lower_case: If true, lowercases the input (independent if the
-            model is cased or not)
-        tokenizer_name_or_path: Name or path of the tokenizer. When
-            None, then model_name_or_path is used
-    """
-
-    def __init__(
-        self,
-        model_name_or_path: str,
-        max_seq_length: Optional[int] = None,
-        model_args: Optional[Dict[str, Any]] = None,
-        tokenizer_args: Optional[Dict[str, Any]] = None,
-        config_args: Optional[Dict[str, Any]] = None,
-        cache_dir: Optional[str] = None,
-        do_lower_case: bool = False,
-        tokenizer_name_or_path: str = None,
-    ) -> None:
-        super(Transformer, self).__init__()
-        self.config_keys = ["max_seq_length", "do_lower_case"]
-        self.do_lower_case = do_lower_case
-        if model_args is None:
-            model_args = {}
-        if tokenizer_args is None:
-            tokenizer_args = {}
-        if config_args is None:
-            config_args = {}
-
-        config = AutoConfig.from_pretrained(model_name_or_path, **config_args, cache_dir=cache_dir)
-        self.jina_clip = AutoModel.from_pretrained(
-            model_name_or_path, config=config, cache_dir=cache_dir, **model_args
-        )
-
-        if max_seq_length is not None and "model_max_length" not in tokenizer_args:
-            tokenizer_args["model_max_length"] = max_seq_length
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            tokenizer_name_or_path if tokenizer_name_or_path is not None else model_name_or_path,
-            cache_dir=cache_dir,
-            **tokenizer_args,
-        )
-        self.preprocessor = AutoImageProcessor.from_pretrained(
-            tokenizer_name_or_path if tokenizer_name_or_path is not None else model_name_or_path,
-            cache_dir=cache_dir,
-            **tokenizer_args,
-        )
-
-        # No max_seq_length set. Try to infer from model
-        if max_seq_length is None:
-            if (
-                hasattr(self.jina_clip, "config")
-                and hasattr(self.jina_clip.config, "max_position_embeddings")
-                and hasattr(self.tokenizer, "model_max_length")
-            ):
-                max_seq_length = min(self.jina_clip.config.max_position_embeddings, self.tokenizer.model_max_length)
-
-        self.max_seq_length = max_seq_length
-
-        if tokenizer_name_or_path is not None:
-            self.jina_clip.config.tokenizer_class = self.tokenizer.__class__.__name__
-
-    def forward(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        """Returns token_embeddings, cls_token"""
-        if "input_ids" in features:
-            embedding = self.jina_clip.get_text_features(input_ids=features["input_ids"])
-        else:
-            embedding = self.jina_clip.get_image_features(pixel_values=features["pixel_values"])
-        return {"sentence_embedding": embedding}
-
-    def get_word_embedding_dimension(self) -> int:
-        return self.config.text_config.embed_dim
-
-    def decode_data_image(data_image_str):
-        header, data = data_image_str.split(',', 1)
-        image_data = base64.b64decode(data)
-        return Image.open(BytesIO(image_data))
-
-    def tokenize(
-        self, batch: Union[List[str]], padding: Union[str, bool] = True
-    ) -> Dict[str, torch.Tensor]:
-        """Tokenizes a text and maps tokens to token-ids"""
-        images = []
-        texts = []
-        for sample in batch:
-            if isinstance(sample, str):
-                if sample.startswith('http'):
-                    response = requests.get(sample)
-                    images.append(Image.open(BytesIO(response.content)).convert('RGB'))
-                elif sample.startswith('data:image/'):
-                    images.append(self.decode_data_image(sample).convert('RGB'))
-                else:
-                    # TODO: Make sure that Image.open fails for non-image files
-                    try:
-                        images.append(Image.open(sample).convert('RGB'))
-                    except:
-                        texts.append(sample)
-            elif isinstance(sample, Image.Image):
-                images.append(sample.convert('RGB'))
-
-        if images and texts:
-            raise ValueError('Batch must contain either images or texts, not both')
-
-        if texts:
-            return self.tokenizer(
-                texts,
-                padding=padding,
-                truncation="longest_first",
-                return_tensors="pt",
-                max_length=self.max_seq_length,
-            )
-        elif images:
-            return self.preprocessor(images)
-        return {}
-
-    def save(self, output_path: str, safe_serialization: bool = True) -> None:
-        self.jina_clip.save_pretrained(output_path, safe_serialization=safe_serialization)
-        self.tokenizer.save_pretrained(output_path)
-        self.preprocessor.save_pretrained(output_path)
-
-    @staticmethod
-    def load(input_path: str) -> "Transformer":
-        # Old classes used other config names than 'sentence_bert_config.json'
-        for config_name in [
-            "sentence_bert_config.json",
-            "sentence_roberta_config.json",
-            "sentence_distilbert_config.json",
-            "sentence_camembert_config.json",
-            "sentence_albert_config.json",
-            "sentence_xlm-roberta_config.json",
-            "sentence_xlnet_config.json",
-        ]:
-            sbert_config_path = os.path.join(input_path, config_name)
-            if os.path.exists(sbert_config_path):
-                break
-
-        with open(sbert_config_path) as fIn:
-            config = json.load(fIn)
-        # Don't allow configs to set trust_remote_code
-        if "model_args" in config and "trust_remote_code" in config["model_args"]:
-            config["model_args"].pop("trust_remote_code")
-        if "tokenizer_args" in config and "trust_remote_code" in config["tokenizer_args"]:
-            config["tokenizer_args"].pop("trust_remote_code")
-        if "config_args" in config and "trust_remote_code" in config["config_args"]:
-            config["config_args"].pop("trust_remote_code")
-        return Transformer(model_name_or_path=input_path, **config)

custom_st_2.py

@@ -1,3 +0,0 @@
-
-class OtherClass:
-    pass

eva_model.py

@@ -1,764 +0,0 @@
-# --------------------------------------------------------
-# Adapted from EVA CLIP
-# https://github.com/baaivision/EVA/tree/master/EVA-CLIP/rei/eva_clip
-# --------------------------------------------------------
-
-import math
-import os
-from functools import partial
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-try:
-    from timm.models.layers import drop_path, to_2tuple, trunc_normal_
-except ImportError or ModuleNotFoundError:
-    from timm.layers import drop_path, to_2tuple, trunc_normal_
-
-from .rope_embeddings import VisionRotaryEmbeddingFast
-
-if os.getenv('ENV_TYPE') == 'deepspeed':
-    try:
-        from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
-    except ImportError or ModuleNotFoundError:
-        from torch.utils.checkpoint import checkpoint
-else:
-    from torch.utils.checkpoint import checkpoint
-
-try:
-    import xformers.ops as xops
-except ImportError:
-    xops = None
-
-
-class PatchDropout(nn.Module):
-    """
-    https://arxiv.org/abs/2212.00794
-    """
-
-    def __init__(self, prob, exclude_first_token=True):
-        super().__init__()
-        assert 0 <= prob < 1.0
-        self.prob = prob
-        self.exclude_first_token = exclude_first_token  # exclude CLS token
-
-    def forward(self, x):
-        if not self.training or self.prob == 0.0:
-            return x
-
-        if self.exclude_first_token:
-            cls_tokens, x = x[:, :1], x[:, 1:]
-        else:
-            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
-
-        batch = x.size()[0]
-        num_tokens = x.size()[1]
-
-        batch_indices = torch.arange(batch)
-        batch_indices = batch_indices[..., None]
-
-        keep_prob = 1 - self.prob
-        num_patches_keep = max(1, int(num_tokens * keep_prob))
-
-        rand = torch.randn(batch, num_tokens)
-        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
-
-        x = x[batch_indices, patch_indices_keep]
-
-        if self.exclude_first_token:
-            x = torch.cat((cls_tokens, x), dim=1)
-
-        return x, patch_indices_keep
-
-
-class DropPath(nn.Module):
-    """Drop paths (Stochastic Depth) per sample  (when applied in main path of
-    residual blocks)."""
-
-    def __init__(self, drop_prob=None):
-        super(DropPath, self).__init__()
-        self.drop_prob = drop_prob
-
-    def forward(self, x):
-        return drop_path(x, self.drop_prob, self.training)
-
-    def extra_repr(self) -> str:
-        return 'p={}'.format(self.drop_prob)
-
-
-class Mlp(nn.Module):
-    def __init__(
-        self,
-        in_features,
-        hidden_features=None,
-        out_features=None,
-        act_layer=nn.GELU,
-        norm_layer=nn.LayerNorm,
-        drop=0.0,
-        subln=False,
-    ):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
-
-        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
-
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.drop = nn.Dropout(drop)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        # x = self.drop(x)
-        # commit this for the orignal BERT implement
-        x = self.ffn_ln(x)
-
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
-
-
-class SwiGLU(nn.Module):
-    def __init__(
-        self,
-        in_features,
-        hidden_features=None,
-        out_features=None,
-        act_layer=nn.SiLU,
-        drop=0.0,
-        norm_layer=nn.LayerNorm,
-        subln=False,
-    ):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-
-        self.w1 = nn.Linear(in_features, hidden_features)
-        self.w2 = nn.Linear(in_features, hidden_features)
-
-        self.act = act_layer()
-        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
-        self.w3 = nn.Linear(hidden_features, out_features)
-
-        self.drop = nn.Dropout(drop)
-
-    def forward(self, x):
-        x1 = self.w1(x)
-        x2 = self.w2(x)
-        hidden = self.act(x1) * x2
-        x = self.ffn_ln(hidden)
-        x = self.w3(x)
-        x = self.drop(x)
-        return x
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim,
-        num_heads=8,
-        qkv_bias=False,
-        qk_scale=None,
-        attn_drop=0.0,
-        proj_drop=0.0,
-        window_size=None,
-        attn_head_dim=None,
-        xattn=False,
-        rope=None,
-        subln=False,
-        norm_layer=nn.LayerNorm,
-    ):
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        if attn_head_dim is not None:
-            head_dim = attn_head_dim
-        all_head_dim = head_dim * self.num_heads
-        self.scale = qk_scale or head_dim**-0.5
-
-        self.subln = subln
-        if self.subln:
-            self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
-            self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
-            self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
-        else:
-            self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
-
-        if qkv_bias:
-            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
-            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
-        else:
-            self.q_bias = None
-            self.v_bias = None
-
-        if window_size:
-            self.window_size = window_size
-            self.num_relative_distance = (2 * window_size[0] - 1) * (
-                2 * window_size[1] - 1
-            ) + 3
-            self.relative_position_bias_table = nn.Parameter(
-                torch.zeros(self.num_relative_distance, num_heads)
-            )  # 2*Wh-1 * 2*Ww-1, nH
-            # cls to token & token 2 cls & cls to cls
-
-            # get pair-wise relative position index for each token inside the window
-            coords_h = torch.arange(window_size[0])
-            coords_w = torch.arange(window_size[1])
-            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
-            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-            relative_coords = (
-                coords_flatten[:, :, None] - coords_flatten[:, None, :]
-            )  # 2, Wh*Ww, Wh*Ww
-            relative_coords = relative_coords.permute(
-                1, 2, 0
-            ).contiguous()  # Wh*Ww, Wh*Ww, 2
-            relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
-            relative_coords[:, :, 1] += window_size[1] - 1
-            relative_coords[:, :, 0] *= 2 * window_size[1] - 1
-            relative_position_index = torch.zeros(
-                size=(window_size[0] * window_size[1] + 1,) * 2,
-                dtype=relative_coords.dtype,
-            )
-            relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
-            relative_position_index[0, 0:] = self.num_relative_distance - 3
-            relative_position_index[0:, 0] = self.num_relative_distance - 2
-            relative_position_index[0, 0] = self.num_relative_distance - 1
-
-            self.register_buffer('relative_position_index', relative_position_index)
-        else:
-            self.window_size = None
-            self.relative_position_bias_table = None
-            self.relative_position_index = None
-
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
-        # self.proj = nn.Linear(all_head_dim, all_head_dim)
-        self.proj = nn.Linear(all_head_dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-        self.xattn = xattn
-        self.xattn_drop = attn_drop
-
-        self.rope = rope
-
-    def forward(self, x, rel_pos_bias=None, attn_mask=None):
-        B, N, C = x.shape
-        if self.subln:
-            q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
-            k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
-            v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
-
-            q = q.reshape(B, N, self.num_heads, -1).permute(
-                0, 2, 1, 3
-            )  # B, num_heads, N, C
-            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
-            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
-        else:
-            qkv_bias = None
-            if self.q_bias is not None:
-                qkv_bias = torch.cat(
-                    (
-                        self.q_bias,
-                        torch.zeros_like(self.v_bias, requires_grad=False),
-                        self.v_bias,
-                    )
-                )
-
-            qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
-            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(
-                2, 0, 3, 1, 4
-            )  # 3, B, num_heads, N, C
-            q, k, v = qkv[0], qkv[1], qkv[2]
-
-        if self.rope:
-            # slightly fast impl
-            q_t = q[:, :, 1:, :]
-            ro_q_t = self.rope(q_t)
-            q = torch.cat((q[:, :, :1, :], ro_q_t), -2).type_as(v)
-
-            k_t = k[:, :, 1:, :]
-            ro_k_t = self.rope(k_t)
-            k = torch.cat((k[:, :, :1, :], ro_k_t), -2).type_as(v)
-
-        if self.xattn:
-            if xops is None:
-                raise ValueError(
-                    "Can't use xattn without xformers. Please 'pip install xformers'"
-                )
-            q = q.permute(0, 2, 1, 3)  # B, num_heads, N, C -> B, N, num_heads, C
-            k = k.permute(0, 2, 1, 3)
-            v = v.permute(0, 2, 1, 3)
-
-            x = xops.memory_efficient_attention(
-                q,
-                k,
-                v,
-                p=self.xattn_drop,
-                scale=self.scale,
-            )
-            x = x.reshape(B, N, -1)
-            x = self.inner_attn_ln(x)
-            x = self.proj(x)
-            x = self.proj_drop(x)
-        else:
-            q = q * self.scale
-            attn = q @ k.transpose(-2, -1)
-
-            if self.relative_position_bias_table is not None:
-                relative_position_bias = self.relative_position_bias_table[
-                    self.relative_position_index.view(-1)
-                ].view(
-                    self.window_size[0] * self.window_size[1] + 1,
-                    self.window_size[0] * self.window_size[1] + 1,
-                    -1,
-                )  # Wh*Ww,Wh*Ww,nH
-                relative_position_bias = relative_position_bias.permute(
-                    2, 0, 1
-                ).contiguous()  # nH, Wh*Ww, Wh*Ww
-                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
-
-            if rel_pos_bias is not None:
-                attn = attn + rel_pos_bias.type_as(attn)
-
-            if attn_mask is not None:
-                attn_mask = attn_mask.bool()
-                attn = attn.masked_fill(~attn_mask[:, None, None, :], float('-inf'))
-
-            attn = attn.softmax(dim=-1)
-            attn = self.attn_drop(attn)
-
-            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
-            x = self.inner_attn_ln(x)
-            x = self.proj(x)
-            x = self.proj_drop(x)
-        return x
-
-
-class Block(nn.Module):
-    def __init__(
-        self,
-        dim,
-        num_heads,
-        mlp_ratio=4.0,
-        qkv_bias=False,
-        qk_scale=None,
-        drop=0.0,
-        attn_drop=0.0,
-        drop_path=0.0,
-        init_values=None,
-        act_layer=nn.GELU,
-        norm_layer=nn.LayerNorm,
-        window_size=None,
-        attn_head_dim=None,
-        xattn=False,
-        rope=None,
-        postnorm=False,
-        subln=False,
-        naiveswiglu=False,
-    ):
-        super().__init__()
-        self.norm1 = norm_layer(dim)
-        self.attn = Attention(
-            dim,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            qk_scale=qk_scale,
-            attn_drop=attn_drop,
-            proj_drop=drop,
-            window_size=window_size,
-            attn_head_dim=attn_head_dim,
-            xattn=xattn,
-            rope=rope,
-            subln=subln,
-            norm_layer=norm_layer,
-        )
-        # NOTE: drop path for stochastic depth, we shall see if this is better
-        # than dropout here
-        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-
-        if naiveswiglu:
-            self.mlp = SwiGLU(
-                in_features=dim,
-                hidden_features=mlp_hidden_dim,
-                subln=subln,
-                norm_layer=norm_layer,
-            )
-        else:
-            self.mlp = Mlp(
-                in_features=dim,
-                hidden_features=mlp_hidden_dim,
-                act_layer=act_layer,
-                subln=subln,
-                drop=drop,
-            )
-
-        if init_values is not None and init_values > 0:
-            self.gamma_1 = nn.Parameter(
-                init_values * torch.ones((dim,)), requires_grad=True
-            )
-            self.gamma_2 = nn.Parameter(
-                init_values * torch.ones((dim,)), requires_grad=True
-            )
-        else:
-            self.gamma_1, self.gamma_2 = None, None
-
-        self.postnorm = postnorm
-
-    def forward(self, x, rel_pos_bias=None, attn_mask=None):
-        if self.gamma_1 is None:
-            if self.postnorm:
-                x = x + self.drop_path(
-                    self.norm1(
-                        self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)
-                    )
-                )
-                x = x + self.drop_path(self.norm2(self.mlp(x)))
-            else:
-                x = x + self.drop_path(
-                    self.attn(
-                        self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask
-                    )
-                )
-                x = x + self.drop_path(self.mlp(self.norm2(x)))
-        else:
-            if self.postnorm:
-                x = x + self.drop_path(
-                    self.gamma_1
-                    * self.norm1(
-                        self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)
-                    )
-                )
-                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
-            else:
-                x = x + self.drop_path(
-                    self.gamma_1
-                    * self.attn(
-                        self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask
-                    )
-                )
-                x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
-        return x
-
-
-class PatchEmbed(nn.Module):
-    """Image to Patch Embedding"""
-
-    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
-        super().__init__()
-        img_size = to_2tuple(img_size)
-        patch_size = to_2tuple(patch_size)
-        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
-        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.num_patches = num_patches
-
-        self.proj = nn.Conv2d(
-            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
-        )
-
-    def forward(self, x, **kwargs):
-        target_dtype = self.proj.weight.dtype
-        B, C, H, W = x.shape
-        # FIXME look at relaxing size constraints
-        assert H == self.img_size[0] and W == self.img_size[1], (
-            f"Input image size ({H}*{W}) doesn't match model "
-            f'({self.img_size[0]}*{self.img_size[1]}).'
-        )
-        x = self.proj(x.to(dtype=target_dtype)).flatten(2).transpose(1, 2)
-        return x
-
-
-class RelativePositionBias(nn.Module):
-    def __init__(self, window_size, num_heads):
-        super().__init__()
-        self.window_size = window_size
-        self.num_relative_distance = (2 * window_size[0] - 1) * (
-            2 * window_size[1] - 1
-        ) + 3
-        self.relative_position_bias_table = nn.Parameter(
-            torch.zeros(self.num_relative_distance, num_heads)
-        )  # 2*Wh-1 * 2*Ww-1, nH
-        # cls to token & token 2 cls & cls to cls
-
-        # get pair-wise relative position index for each token inside the window
-        coords_h = torch.arange(window_size[0])
-        coords_w = torch.arange(window_size[1])
-        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
-        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-        relative_coords = (
-            coords_flatten[:, :, None] - coords_flatten[:, None, :]
-        )  # 2, Wh*Ww, Wh*Ww
-        relative_coords = relative_coords.permute(
-            1, 2, 0
-        ).contiguous()  # Wh*Ww, Wh*Ww, 2
-        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
-        relative_coords[:, :, 1] += window_size[1] - 1
-        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
-        relative_position_index = torch.zeros(
-            size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype
-        )
-        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
-        relative_position_index[0, 0:] = self.num_relative_distance - 3
-        relative_position_index[0:, 0] = self.num_relative_distance - 2
-        relative_position_index[0, 0] = self.num_relative_distance - 1
-
-        self.register_buffer('relative_position_index', relative_position_index)
-
-    def forward(self):
-        relative_position_bias = self.relative_position_bias_table[
-            self.relative_position_index.view(-1)
-        ].view(
-            self.window_size[0] * self.window_size[1] + 1,
-            self.window_size[0] * self.window_size[1] + 1,
-            -1,
-        )  # Wh*Ww,Wh*Ww,nH
-        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
-
-
-class EVAVisionTransformer(nn.Module):
-    """Vision Transformer with support for patch or hybrid CNN input stage"""
-
-    def __init__(
-        self,
-        img_size=224,
-        patch_size=16,
-        in_chans=3,
-        num_classes=0,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4.0,
-        qkv_bias=False,
-        qk_scale=None,
-        drop_rate=0.0,
-        attn_drop_rate=0.0,
-        drop_path_rate=0.0,
-        norm_layer=nn.LayerNorm,
-        init_values=None,
-        patch_dropout=0.0,
-        use_abs_pos_emb=True,
-        use_rel_pos_bias=False,
-        use_shared_rel_pos_bias=False,
-        rope=False,
-        use_mean_pooling=True,
-        init_scale=0.001,
-        grad_checkpointing=False,
-        xattn=False,
-        postnorm=False,
-        pt_hw_seq_len=16,
-        intp_freq=False,
-        naiveswiglu=False,
-        subln=False,
-        proj_type=None,
-    ):
-        super().__init__()
-        self.image_size = img_size
-        self.num_classes = num_classes
-        self.num_features = (
-            self.embed_dim
-        ) = embed_dim  # num_features for consistency with other models
-
-        self.patch_embed = PatchEmbed(
-            img_size=img_size,
-            patch_size=patch_size,
-            in_chans=in_chans,
-            embed_dim=embed_dim,
-        )
-        num_patches = self.patch_embed.num_patches
-
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        # self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        if use_abs_pos_emb:
-            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
-        else:
-            self.pos_embed = None
-        self.pos_drop = nn.Dropout(p=drop_rate)
-
-        if use_shared_rel_pos_bias:
-            self.rel_pos_bias = RelativePositionBias(
-                window_size=self.patch_embed.patch_shape, num_heads=num_heads
-            )
-        else:
-            self.rel_pos_bias = None
-
-        if rope:
-            half_head_dim = embed_dim // num_heads // 2
-            hw_seq_len = img_size // patch_size
-            self.rope = VisionRotaryEmbeddingFast(
-                dim=half_head_dim,
-                pt_seq_len=pt_hw_seq_len,
-                ft_seq_len=hw_seq_len if intp_freq else None,
-                patch_dropout=patch_dropout,
-            )
-        else:
-            self.rope = None
-
-        self.naiveswiglu = naiveswiglu
-
-        dpr = [
-            x.item() for x in torch.linspace(0, drop_path_rate, depth)
-        ]  # stochastic depth decay rule
-        self.use_rel_pos_bias = use_rel_pos_bias
-        self.blocks = nn.ModuleList(
-            [
-                Block(
-                    dim=embed_dim,
-                    num_heads=num_heads,
-                    mlp_ratio=mlp_ratio,
-                    qkv_bias=qkv_bias,
-                    qk_scale=qk_scale,
-                    drop=drop_rate,
-                    attn_drop=attn_drop_rate,
-                    drop_path=dpr[i],
-                    norm_layer=norm_layer,
-                    init_values=init_values,
-                    window_size=self.patch_embed.patch_shape
-                    if use_rel_pos_bias
-                    else None,
-                    xattn=xattn,
-                    rope=self.rope,
-                    postnorm=postnorm,
-                    subln=subln,
-                    naiveswiglu=naiveswiglu,
-                )
-                for i in range(depth)
-            ]
-        )
-        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
-        self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
-        if (num_classes == embed_dim) and (proj_type is None):
-            self.head = nn.Identity()
-        elif proj_type == 'linear':
-            self.head = nn.Linear(embed_dim, num_classes, bias=qkv_bias)
-        elif proj_type == 'mlp':
-            hidden_size = (embed_dim + num_classes) // 2
-            self.proj = nn.Sequential(
-                nn.Linear(embed_dim, hidden_size, bias=qkv_bias),
-                nn.GELU(),
-                nn.Linear(hidden_size, num_classes, bias=qkv_bias),
-            )
-
-        if self.pos_embed is not None:
-            trunc_normal_(self.pos_embed, std=0.02)
-
-        trunc_normal_(self.cls_token, std=0.02)
-
-        self.apply(self._init_weights)
-        self.fix_init_weight()
-
-        if isinstance(self.head, nn.Linear):
-            trunc_normal_(self.head.weight, std=0.02)
-            self.head.weight.data.mul_(init_scale)
-            if qkv_bias:
-                self.head.bias.data.mul_(init_scale)
-
-        # setting a patch_dropout of 0. would mean it is disabled and this function
-        # would be the identity fn
-        self.patch_dropout = (
-            PatchDropout(patch_dropout) if patch_dropout > 0.0 else nn.Identity()
-        )
-
-        self.grad_checkpointing = grad_checkpointing
-
-    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)
-            if self.naiveswiglu:
-                rescale(layer.mlp.w3.weight.data, layer_id + 1)
-            else:
-                rescale(layer.mlp.fc2.weight.data, layer_id + 1)
-
-    def get_cast_dtype(self) -> torch.dtype:
-        return self.blocks[0].mlp.fc2.weight.dtype
-
-    def _init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=0.02)
-            if 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)
-
-    def get_num_layers(self):
-        return len(self.blocks)
-
-    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
-        assert (
-            unlocked_groups == 0
-        ), 'partial locking not currently supported for this model'
-        for param in self.parameters():
-            param.requires_grad = False
-
-    @torch.jit.ignore
-    def set_grad_checkpointing(self, enable=True):
-        self.grad_checkpointing = enable
-
-    @torch.jit.ignore
-    def no_weight_decay(self):
-        return {'pos_embed', 'cls_token'}
-
-    def get_classifier(self):
-        return self.head
-
-    def reset_classifier(self, num_classes, global_pool=''):
-        self.num_classes = num_classes
-        self.head = (
-            nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
-        )
-
-    def forward_features(self, x, return_all_features=False):
-        x = self.patch_embed(x)
-        batch_size, seq_len, _ = x.size()
-
-        cls_tokens = self.cls_token.expand(
-            batch_size, -1, -1
-        )  # stole cls_tokens impl from Phil Wang, thanks
-        x = torch.cat((cls_tokens, x), dim=1)
-        if self.pos_embed is not None:
-            x = x + self.pos_embed
-        x = self.pos_drop(x)
-
-        # a patch_dropout of 0. would mean it is disabled and this function would do
-        # nothing but return what was passed in
-        if self.rope is not None:
-            if self.training and not isinstance(self.patch_dropout, nn.Identity):
-                x, patch_indices_keep = self.patch_dropout(x)
-                self.rope.forward = partial(
-                    self.rope.forward, patch_indices_keep=patch_indices_keep
-                )
-            else:
-                self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
-                x = self.patch_dropout(x)
-        else:
-            x = self.patch_dropout(x)
-
-        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
-        for blk in self.blocks:
-            if self.grad_checkpointing:
-                x = checkpoint(blk, x, (rel_pos_bias,))
-            else:
-                x = blk(x, rel_pos_bias=rel_pos_bias)
-
-        if not return_all_features:
-            x = self.norm(x)
-            if self.fc_norm is not None:
-                return self.fc_norm(x.mean(1))
-            else:
-                return x[:, 0]
-        return x
-
-    def forward(self, x, return_all_features=False):
-        if return_all_features:
-            return self.forward_features(x, return_all_features)
-        x = self.forward_features(x)
-        x = self.head(x)
-        return x

hf_model.py

@@ -1,297 +0,0 @@
-import re
-from typing import Dict, Optional, Tuple
-
-import torch
-import torch.nn as nn
-from transformers import AutoConfig, AutoModel, PretrainedConfig
-from transformers.modeling_outputs import (
-    BaseModelOutput,
-    BaseModelOutputWithPooling,
-    BaseModelOutputWithPoolingAndCrossAttentions,
-)
-
-"""
-HF architecture mapping
-"""
-
-_HF_ARCH_DICT = {
-    # https://huggingface.co/docs/transformers/model_doc/roberta#roberta
-    'roberta': {
-        'config_names': {
-            'context_length': 'max_position_embeddings',
-            'vocab_size': 'vocab_size',
-            'width': 'hidden_size',
-            'heads': 'num_attention_heads',
-            'layers': 'num_hidden_layers',
-            'layer_attr': 'layer',
-            'token_embeddings_attr': 'embeddings',
-        },
-        'pooler': 'mean_pooler',
-    },
-    # https://huggingface.co/docs/transformers/model_doc/xlm-roberta#transformers.XLMRobertaConfig
-    'xlm-roberta': {
-        'config_names': {
-            'context_length': 'max_position_embeddings',
-            'vocab_size': 'vocab_size',
-            'width': 'hidden_size',
-            'heads': 'num_attention_heads',
-            'layers': 'num_hidden_layers',
-            'layer_attr': 'layer',
-            'token_embeddings_attr': 'embeddings',
-        },
-        'pooler': 'mean_pooler',
-    },
-    # https://huggingface.co/docs/transformers/model_doc/mt5#mt5
-    'mt5': {
-        'config_names': {
-            # unlimited seqlen
-            # https://github.com/google-research/text-to-text-transfer-transformer/issues/273
-            # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/t5/modeling_t5.py#L374
-            'context_length': '',
-            'vocab_size': 'vocab_size',
-            'width': 'd_model',
-            'heads': 'num_heads',
-            'layers': 'num_layers',
-            'layer_attr': 'block',
-            'token_embeddings_attr': 'embed_tokens',
-        },
-        'pooler': 'mean_pooler',
-    },
-    # https://huggingface.co/docs/transformers/model_doc/bert
-    'bert': {
-        'config_names': {
-            'context_length': 'max_position_embeddings',
-            'vocab_size': 'vocab_size',
-            'width': 'hidden_size',
-            'heads': 'num_attention_heads',
-            'layers': 'num_hidden_layers',
-        },
-        'pooler': 'cls_pooler',
-    },
-    # https://huggingface.co/docs/transformers/model_doc/m2m_100
-    'm2m_100': {
-        'config_names': {
-            'context_length': 'max_position_embeddings',
-            'vocab_size': 'vocab_size',
-            'width': 'd_model',
-            'heads': 'encoder_attention_heads',
-            'layers': 'encoder_layers',
-        },
-        'pooler': 'cls_pooler',
-    },
-}
-
-
-"""
-Pooling functions
-"""
-
-_POOLERS = {}
-
-
-def _camel2snake(s):
-    return re.sub(r'(?<!^)(?=[A-Z])', '_', s).lower()
-
-
-def register_pooler(cls):
-    """Decorator registering pooler class"""
-    _POOLERS[_camel2snake(cls.__name__)] = cls
-    return cls
-
-
-@register_pooler
-class MeanPooler(nn.Module):
-    """Mean pooling"""
-
-    @staticmethod
-    def forward(x: BaseModelOutput, attention_mask: torch.Tensor):
-        masked_output = x.last_hidden_state * attention_mask.unsqueeze(-1)
-        return masked_output.sum(dim=1) / attention_mask.sum(-1, keepdim=True)
-
-
-@register_pooler
-class MaxPooler(nn.Module):
-    """
-    Max pooling
-    """
-
-    @staticmethod
-    def forward(x: BaseModelOutput, attention_mask: torch.Tensor):
-        masked_output = x.last_hidden_state.masked_fill(
-            attention_mask.unsqueeze(-1), -torch.inf
-        )
-        return masked_output.max(1).values
-
-
-@register_pooler
-class ClsPooler(nn.Module):
-    """
-    CLS token pooling
-    """
-
-    def __init__(self, use_pooler_output=True):
-        super().__init__()
-        self.cls_token_position = 0
-        self.use_pooler_output = use_pooler_output
-
-    def forward(self, x: BaseModelOutput, _: torch.Tensor):
-        if (
-            self.use_pooler_output
-            and isinstance(
-                x,
-                (
-                    BaseModelOutputWithPooling,
-                    BaseModelOutputWithPoolingAndCrossAttentions,
-                ),
-            )
-            and (x.pooler_output is not None)
-        ):
-            return x.pooler_output
-
-        return x.last_hidden_state[:, self.cls_token_position, :]
-
-
-"""
-HF text model
-"""
-
-
-class HFTextEncoder(nn.Module):
-    output_tokens: torch.jit.Final[bool]
-
-    def __init__(
-        self,
-        model_name_or_path: str,
-        output_dim: int,
-        config: PretrainedConfig = None,
-        pooler_type: str = None,
-        proj_type: str = None,
-        proj_bias: bool = False,
-        pretrained: bool = True,
-        output_tokens: bool = False,
-        trust_remote_code: bool = False,
-        revision: Optional[str] = None,
-        model_config_kwargs: Optional[Dict] = None,
-    ):
-        super().__init__()
-        self.output_tokens = output_tokens
-        self.output_dim = output_dim
-
-        # TODO: find better way to get this information
-        uses_transformer_pooler = pooler_type == 'cls_pooler'
-        model_config_kwargs = model_config_kwargs or {}
-
-        if config is None:
-            self.config = AutoConfig.from_pretrained(
-                model_name_or_path,
-                trust_remote_code=trust_remote_code,
-                code_revision=revision,
-            )
-            self.config.update(model_config_kwargs)
-            create_func, model_args = (
-                (AutoModel.from_pretrained, model_name_or_path)
-                if pretrained
-                else (AutoModel.from_config, self.config)
-            )
-            # TODO: do all model configs have this attribute?
-            #  PretrainedConfig does so yes??
-            if (
-                hasattr(self.config, 'is_encoder_decoder')
-                and self.config.is_encoder_decoder
-            ):
-                self.transformer = create_func(model_args)
-                self.transformer = self.transformer.encoder
-            else:
-                self.transformer = create_func(
-                    model_args,
-                    trust_remote_code=trust_remote_code,
-                    add_pooling_layer=uses_transformer_pooler,
-                    code_revision=revision,
-                )
-        else:
-            self.config = config
-            self.config.update(model_config_kwargs)
-            self.transformer = AutoModel.from_config(self.config)
-
-        if pooler_type is None:  # get default arch pooler
-            pooler_type = _HF_ARCH_DICT[self.config.model_type]['pooler']
-
-        # FIXME downstream users of OpenCLIP models use these attr,
-        #  need to verify valid across all models
-        self.vocab_size = getattr(self.config, 'vocab_size', 0)
-        self.context_length = getattr(self.config, 'max_position_embeddings', 0)
-
-        self.pooler = _POOLERS[pooler_type]()
-
-        d_model = getattr(
-            self.config, _HF_ARCH_DICT[self.config.model_type]['config_names']['width']
-        )
-        if (d_model == output_dim) and (proj_type is None):  # do we always need a proj?
-            self.proj = nn.Identity()
-        elif proj_type == 'linear':
-            self.proj = nn.Linear(d_model, output_dim, bias=proj_bias)
-        elif proj_type == 'mlp':
-            hidden_size = (d_model + output_dim) // 2
-            self.proj = nn.Sequential(
-                nn.Linear(d_model, hidden_size, bias=proj_bias),
-                nn.GELU(),
-                nn.Linear(hidden_size, output_dim, bias=proj_bias),
-            )
-
-    def forward(self, x: torch.Tensor):
-        attn_mask = (x != self.config.pad_token_id).long()
-        out = self.transformer(input_ids=x, attention_mask=attn_mask)
-        pooled_out = self.pooler(out, attn_mask)
-        projected = self.proj(pooled_out)
-
-        seq_len = out.last_hidden_state.shape[1]
-        tokens = (
-            out.last_hidden_state[
-                :, torch.arange(seq_len) != self.pooler.cls_token_position, :
-            ]
-            if isinstance(self.pooler, ClsPooler)
-            else out.last_hidden_state
-        )
-
-        if self.output_tokens:
-            return projected, tokens
-        return projected
-
-    def lock(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
-        if not unlocked_layers:  # full freezing
-            for n, p in self.transformer.named_parameters():
-                p.requires_grad = (
-                    (not freeze_layer_norm) if 'LayerNorm' in n.split('.') else False
-                )
-            return
-
-        encoder = (
-            self.transformer.encoder
-            if hasattr(self.transformer, 'encoder')
-            else self.transformer
-        )
-        layer_list = getattr(
-            encoder, _HF_ARCH_DICT[self.config.model_type]['config_names']['layer_attr']
-        )
-        print(f'Unlocking {unlocked_layers}/{len(layer_list) + 1} layers of hf model')
-        embeddings = getattr(
-            self.transformer,
-            _HF_ARCH_DICT[self.config.model_type]['config_names'][
-                'token_embeddings_attr'
-            ],
-        )
-        modules = [embeddings, *layer_list][:-unlocked_layers]
-        # freeze layers
-        for module in modules:
-            for n, p in module.named_parameters():
-                p.requires_grad = (
-                    (not freeze_layer_norm) if 'LayerNorm' in n.split('.') else False
-                )
-
-    @torch.jit.ignore
-    def set_grad_checkpointing(self, _=True):
-        self.transformer.gradient_checkpointing_enable()
-
-    def init_parameters(self):
-        pass
-

modeling_clip.py

@@ -1,570 +0,0 @@
-# coding=utf-8
-#
-# Code mainly copied from:
-# https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/modeling_clip.py
-# and adjusted for Jina CLIP
-
-from functools import partial
-from typing import List, Optional, Tuple, Union
-from io import BytesIO
-import requests
-import base64
-import numpy as np
-import torch
-import torch.nn.functional as f
-import torch.utils.checkpoint
-from torch import nn
-from transformers import (
-    AutoImageProcessor,
-    AutoTokenizer,
-    BatchEncoding,
-    BatchFeature,
-    PreTrainedModel,
-    logging,
-)
-from transformers.models.clip.modeling_clip import (
-    CLIPOutput,
-    CLIPTextModelOutput,
-    CLIPVisionModelOutput,
-    clip_loss,
-)
-
-try:
-    from tqdm.autonotebook import trange
-
-    has_tqdm = True
-except ImportError:
-    has_tqdm = False
-
-from .configuration_clip import JinaCLIPConfig, JinaCLIPTextConfig, JinaCLIPVisionConfig
-from .eva_model import EVAVisionTransformer
-from .hf_model import HFTextEncoder
-# needed for HF to correctly import in cache
-from .rope_embeddings import VisionRotaryEmbeddingFast  # noqa: F401
-from .transform import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD, image_transform  # noqa: F401
-
-logger = logging.get_logger(__name__)
-
-
-""" Jina CLIP model implementation """
-
-
-class LayerNorm(nn.LayerNorm):
-    """Subclass torch's LayerNorm (with cast back to input dtype)."""
-
-    def forward(self, x: torch.Tensor):
-        origtype = x.dtype
-        x = f.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
-        return x.to(origtype)
-
-
-def _build_text_tower(config: JinaCLIPTextConfig) -> HFTextEncoder:
-    return HFTextEncoder(
-        model_name_or_path=config.hf_model_name_or_path,
-        output_dim=config.embed_dim,
-        pooler_type=config.pooler_type,
-        proj_type=config.proj_type,
-        proj_bias=config.proj_bias,
-        pretrained=False,
-        output_tokens=False,
-        trust_remote_code=True,
-        revision=None,
-        model_config_kwargs=config.hf_model_config_kwargs,
-    )
-
-
-def _build_vision_tower(config: JinaCLIPVisionConfig) -> EVAVisionTransformer:
-    norm_layer = partial(LayerNorm, eps=1e-6)
-
-    if config.fused_layer_norm:
-        try:
-            from apex.normalization import FusedLayerNorm
-
-            norm_layer = partial(FusedLayerNorm, eps=1e-6)
-        except (ModuleNotFoundError, ImportError):
-            logger.warning('Please install apex to use fused layer norm, ignoring')
-
-    return EVAVisionTransformer(
-        img_size=config.image_size,
-        patch_size=config.patch_size,
-        num_classes=config.embed_dim,
-        use_mean_pooling=False,
-        init_values=config.ls_init_value,
-        patch_dropout=config.patch_dropout,
-        embed_dim=config.width,
-        depth=config.layers,
-        num_heads=config.width // config.head_width,
-        mlp_ratio=config.mlp_ratio,
-        qkv_bias=config.qkv_bias,
-        drop_path_rate=config.drop_path_rate,
-        norm_layer=norm_layer,
-        xattn=config.x_attention,
-        rope=config.rope_embeddings,
-        postnorm=config.post_norm,
-        pt_hw_seq_len=config.pt_hw_seq_len,
-        intp_freq=config.intp_freq,
-        naiveswiglu=config.naive_swiglu,
-        subln=config.subln,
-        proj_type=config.proj_type,
-    )
-
-
-class JinaCLIPPreTrainedModel(PreTrainedModel):
-    """
-    An abstract class to handle weights initialization and a simple interface for
-    downloading and loading pretrained models.
-    """
-
-    config_class = JinaCLIPConfig
-    base_model_prefix = 'clip'
-    supports_gradient_checkpointing = True
-
-    def _init_weights(self, module):
-        """Initialize the weights"""
-        if isinstance(module, JinaCLIPModel):
-            if isinstance(module.text_projection, nn.Linear):
-                nn.init.normal_(
-                    module.text_projection.weight,
-                    std=module.text_embed_dim**-0.5 * self.config.initializer_factor,
-                )
-            if isinstance(module.text_projection, nn.Linear):
-                nn.init.normal_(
-                    module.visual_projection.weight,
-                    std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
-                )
-        if isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-        if isinstance(module, nn.Linear) and module.bias is not None:
-            module.bias.data.zero_()
-
-
-class JinaCLIPTextModel(JinaCLIPPreTrainedModel):
-    config_class = JinaCLIPTextConfig
-
-    def __init__(self, config: JinaCLIPTextConfig):
-        super().__init__(config)
-        self.text_model = _build_text_tower(config)
-        self.post_init()
-
-    def forward(
-        self,
-        input_ids: Union[None, torch.Tensor, BatchEncoding] = None,
-        return_dict: Optional[bool] = None,
-        *_,
-        **__,
-    ) -> Union[Tuple[Optional[torch.FloatTensor], ...], CLIPTextModelOutput]:
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-        x = input_ids.input_ids if isinstance(input_ids, BatchEncoding) else input_ids
-        feats = self.text_model(x=x)
-        out = CLIPTextModelOutput(text_embeds=feats)
-        return out if return_dict else out.to_tuple()
-
-
-class JinaCLIPVisionModel(JinaCLIPPreTrainedModel):
-    config_class = JinaCLIPVisionConfig
-    main_input_name = 'pixel_values'
-
-    def __init__(self, config: JinaCLIPVisionConfig):
-        super().__init__(config)
-        self.vision_model = _build_vision_tower(config)
-        self.post_init()
-
-    def forward(
-        self,
-        pixel_values: Union[None, torch.FloatTensor, BatchFeature] = None,
-        return_dict: Optional[bool] = None,
-        *_,
-        **__,
-    ) -> Union[Tuple[Optional[torch.FloatTensor], ...], CLIPVisionModelOutput]:
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-        x = (
-            pixel_values.pixel_values
-            if isinstance(pixel_values, BatchFeature)
-            else pixel_values
-        )
-        feats = self.vision_model(x=x)
-        out = CLIPVisionModelOutput(image_embeds=feats)
-        return out if return_dict else out.to_tuple()
-
-
-class JinaCLIPModel(JinaCLIPPreTrainedModel):
-    config_class = JinaCLIPConfig
-
-    def __init__(self, config: JinaCLIPConfig):
-        super().__init__(config)
-
-        if not isinstance(config.text_config, JinaCLIPTextConfig):
-            raise ValueError(
-                'Attribute config.text_config is expected to be of type '
-                f'JinaCLIPTextConfig but is of type {type(config.text_config)}.'
-            )
-
-        if not isinstance(config.vision_config, JinaCLIPVisionConfig):
-            raise ValueError(
-                'Attribute config.vision_config is expected to be of type '
-                f'JinaCLIPVisionConfig but is of type {type(config.vision_config)}.'
-            )
-
-        text_config = config.text_config
-        vision_config = config.vision_config
-
-        if config.use_text_flash_attn is not None:
-            text_config.hf_model_config_kwargs['use_flash_attn'] = config.use_text_flash_attn
-        if config.use_vision_xformers is not None:
-            vision_config.x_attention = config.use_vision_xformers
-
-        self.add_projections = config.add_projections
-        self.projection_dim = config.projection_dim
-        self.text_embed_dim = text_config.embed_dim
-        self.vision_embed_dim = vision_config.embed_dim
-
-        self.text_model = _build_text_tower(text_config)
-        self.vision_model = _build_vision_tower(vision_config)
-        self.logit_scale = nn.Parameter(
-            torch.tensor(self.config.logit_scale_init_value)
-        )
-
-        if self.add_projections:
-            self.visual_projection = nn.Linear(
-                self.vision_embed_dim, self.projection_dim, bias=False
-            )
-            self.text_projection = nn.Linear(
-                self.text_embed_dim, self.projection_dim, bias=False
-            )
-        else:
-            self.visual_projection = nn.Identity()
-            self.text_projection = nn.Identity()
-
-        self.tokenizer = None
-        self.preprocess = None
-        self.post_init()
-
-    def get_tokenizer(self):
-        if not self.tokenizer:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.config._name_or_path, trust_remote_code=True
-            )
-        return self.tokenizer
-
-    def get_preprocess(self):
-        if not self.preprocess:
-            self.preprocess = AutoImageProcessor.from_pretrained(
-                self.config._name_or_path, trust_remote_code=True
-            )
-        return self.preprocess
-
-    def get_text_features(
-        self,
-        input_ids: Union[None, torch.Tensor, BatchEncoding] = None,
-        *_,
-        **__,
-    ) -> torch.FloatTensor:
-        x = input_ids.input_ids if isinstance(input_ids, BatchEncoding) else input_ids
-        return self.text_projection(self.text_model(x=x))
-
-    def get_image_features(
-        self,
-        pixel_values: Union[None, torch.FloatTensor, BatchFeature] = None,
-        *_,
-        **__,
-    ) -> torch.FloatTensor:
-        x = (
-            pixel_values.pixel_values
-            if isinstance(pixel_values, BatchFeature)
-            else pixel_values
-        )
-        return self.visual_projection(self.vision_model(x=x))
-
-    @torch.inference_mode()
-    def encode_text(
-        self,
-        sentences: Union[str, List[str]],
-        batch_size: int = 32,
-        show_progress_bar: Optional[bool] = None,
-        convert_to_numpy: bool = True,
-        convert_to_tensor: bool = False,
-        device: Optional[torch.device] = None,
-        normalize_embeddings: bool = True,
-        **tokenizer_kwargs,
-    ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
-        """
-        Computes sentence embeddings
-         Args:
-             sentences(`str` or `List[str]`):
-                 Sentence or sentences to be encoded
-             batch_size(`int`, *optional*, defaults to 32):
-                 Batch size for the computation
-             show_progress_bar(`bool`, *optional*, defaults to None):
-                 Show a progress bar when encoding sentences.
-                 If set to None, progress bar is only shown when
-                 `logger.level == logging.INFO` or `logger.level == logging.DEBUG`.
-             convert_to_numpy(`bool`, *optional*, defaults to True):
-                 If true, the output is a list of numpy vectors.
-                 Else, it is a list of pytorch tensors.
-             convert_to_tensor(`bool`, *optional*, defaults to False):
-                 If true, you get one large tensor as return.
-                 Overwrites any setting from convert_to_numpy
-             device(`torch.device`, *optional*, defaults to None):
-                 Which torch.device to use for the computation
-             normalize_embeddings(`bool`, *optional*, defaults to False):
-                 If set to true, returned vectors will have length 1. In that case,
-                 the faster dot-product (util.dot_score) instead of cosine similarity
-                 can be used.
-             tokenizer_kwargs(`Dict[str, Any]`, *optional*, defaults to {}):
-                 Keyword arguments for the tokenizer
-         Returns:
-             By default, a list of tensors is returned.
-             If convert_to_tensor, a stacked tensor is returned.
-             If convert_to_numpy, a numpy matrix is returned.
-        """
-        is_training = self.training
-        self.eval()
-        all_embeddings = []
-
-        self.tokenizer = self.get_tokenizer()
-
-        if show_progress_bar is None:
-            show_progress_bar = (
-                logger.getEffectiveLevel() == logging.INFO
-                or logger.getEffectiveLevel() == logging.DEBUG
-            )
-
-        if convert_to_tensor:
-            convert_to_numpy = False
-
-        input_was_string = False
-        if isinstance(sentences, str) or not hasattr(sentences, '__len__'):
-            sentences = [sentences]
-            input_was_string = True
-
-        if device is not None:
-            self.to(device)
-
-        permutation = np.argsort([-len(i) for i in sentences])
-        inverse_permutation = np.argsort(permutation)
-        sentences = [sentences[idx] for idx in permutation]
-
-        tokenizer_kwargs['padding'] = tokenizer_kwargs.get('padding', True)
-        tokenizer_kwargs['max_length'] = tokenizer_kwargs.get('max_length', 512)
-        tokenizer_kwargs['truncation'] = tokenizer_kwargs.get('truncation', True)
-
-        if has_tqdm:
-            range_iter = trange(
-                0,
-                len(sentences),
-                batch_size,
-                desc='Encoding',
-                disable=not show_progress_bar,
-            )
-        else:
-            range_iter = range(0, len(sentences), batch_size)
-
-        for i in range_iter:
-            encoded_input = self.tokenizer(
-                sentences[i : i + batch_size],
-                return_tensors='pt',
-                **tokenizer_kwargs,
-            ).to(self.device)
-
-            embeddings = self.get_text_features(input_ids=encoded_input)
-            if normalize_embeddings:
-                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
-            if convert_to_numpy:
-                embeddings = embeddings.cpu()
-            all_embeddings.extend(embeddings)
-
-        all_embeddings = [all_embeddings[idx] for idx in inverse_permutation]
-
-        if convert_to_tensor:
-            all_embeddings = torch.stack(all_embeddings)
-        elif convert_to_numpy:
-            all_embeddings = np.asarray([emb.to(torch.float32).numpy() for emb in all_embeddings])
-
-        if input_was_string:
-            all_embeddings = all_embeddings[0]
-
-        self.train(is_training)
-        return all_embeddings
-
-    def decode_data_image(data_image_str):
-        header, data = data_image_str.split(',', 1)
-        image_data = base64.b64decode(data)
-        return Image.open(BytesIO(image_data))
-
-    @torch.inference_mode()
-    def encode_image(
-        self,
-        images: Union[str, List[Union[str, "Image.Image"]]],
-        batch_size: int = 32,
-        show_progress_bar: Optional[bool] = None,
-        convert_to_numpy: bool = True,
-        convert_to_tensor: bool = False,
-        device: Optional[torch.device] = None,
-        normalize_embeddings: bool = True,
-    ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
-        """
-        Computes image embeddings.
-    
-        Args:
-            images(`str` or `List[Union[str, Image.Image]]`):
-                image paths, URLs, PIL images, or data:image/ strings to be encoded
-            batch_size(`int`, *optional*, defaults to 32):
-                Batch size for the computation
-            show_progress_bar(`bool`, *optional*, defaults to None):
-                Show a progress bar when encoding images.
-                If set to None, progress bar is only shown when
-                `logger.level == logging.INFO` or `logger.level == logging.DEBUG`.
-            convert_to_numpy(`bool`, *optional*, defaults to True):
-                If true, the output is a list of numpy vectors.
-                Else, it is a list of pytorch tensors.
-            convert_to_tensor(`bool`, *optional*, defaults to False):
-                If true, you get one large tensor as return.
-                Overwrites any setting from convert_to_numpy
-            device(`torch.device`, *optional*, defaults to None):
-                Which torch.device to use for the computation
-            normalize_embeddings(`bool`, *optional*, defaults to False):
-                If set to true, returned vectors will have length 1. In that case,
-                the faster dot-product (util.dot_score) instead of cosine similarity
-                can be used.
-        Returns:
-            By default, a list of tensors is returned.
-            If convert_to_tensor, a stacked tensor is returned.
-            If convert_to_numpy, a numpy matrix is returned.
-        """
-        
-        is_training = self.training
-        self.eval()
-    
-        self.preprocess = self.get_preprocess()
-        all_embeddings = []
-    
-        if show_progress_bar is None:
-            show_progress_bar = (
-                logger.getEffectiveLevel() == logging.INFO
-                or logger.getEffectiveLevel() == logging.DEBUG
-            )
-    
-        if convert_to_tensor:
-            convert_to_numpy = False
-    
-        input_was_single_img = False
-        if isinstance(images, str) or not hasattr(images, '__len__'):
-            images = [images]
-            input_was_single_img = True
-    
-        if device is not None:
-            self.to(device)
-    
-        permutation = np.argsort([-len(str(i)) for i in images])
-        inverse_permutation = np.argsort(permutation)
-        images = [images[idx] for idx in permutation]
-    
-        if has_tqdm:
-            range_iter = trange(
-                0,
-                len(images),
-                batch_size,
-                desc='Encoding',
-                disable=not show_progress_bar,
-            )
-        else:
-            range_iter = range(0, len(images), batch_size)
-
-        from PIL import Image
-    
-        for i in range_iter:
-            batch_images = images[i:i+batch_size]
-            processed_inputs = []
-    
-            for img in batch_images:
-                if isinstance(img, str):
-                    if img.startswith('http'):
-                        response = requests.get(img)
-                        image = Image.open(BytesIO(response.content)).convert('RGB')
-                    elif img.startswith('data:image/'):
-                        image = decode_data_image(img).convert('RGB')
-                    else:
-                        image = Image.open(img).convert('RGB')
-                elif isinstance(img, Image.Image):
-                    image = img.convert('RGB')
-                else:
-                    raise ValueError("Unsupported image format")
-    
-                processed_inputs.append(image)
-    
-            processed_inputs = self.preprocess(processed_inputs)
-            processed_inputs = processed_inputs.to(self.device)
-            embeddings = self.get_image_features(processed_inputs)
-            
-            if normalize_embeddings:
-                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
-            if convert_to_numpy:
-                embeddings = embeddings.cpu()
-            all_embeddings.extend(embeddings)
-    
-        all_embeddings = [all_embeddings[idx] for idx in inverse_permutation]
-    
-        if convert_to_tensor:
-            all_embeddings = torch.stack(all_embeddings)
-        elif convert_to_numpy:
-            all_embeddings = np.asarray([emb.to(torch.float32).numpy() for emb in all_embeddings])
-    
-        if input_was_single_img:
-            all_embeddings = all_embeddings[0]
-    
-        self.train(is_training)
-        return all_embeddings
-
-    def forward(
-        self,
-        input_ids: Union[None, torch.Tensor, BatchEncoding] = None,
-        pixel_values: Union[None, torch.FloatTensor, BatchFeature] = None,
-        return_dict: Optional[bool] = None,
-        return_loss: Optional[bool] = None,
-        *_,
-        **__,
-    ) -> Union[Tuple[Optional[torch.FloatTensor], ...], CLIPOutput]:
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-        image_embeds = self.get_image_features(pixel_values=pixel_values)
-        text_embeds = self.get_text_features(input_ids=input_ids)
-
-        # normalized features
-        image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
-        text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
-
-        # cosine similarity as logits
-        logit_scale = self.logit_scale.exp()
-        logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
-        logits_per_image = logits_per_text.t()
-
-        loss = None
-        if return_loss:
-            loss = clip_loss(logits_per_text)
-
-        if not return_dict:
-            output = (
-                logits_per_image,
-                logits_per_text,
-                text_embeds,
-                image_embeds,
-                None,
-                None,
-            )
-            return ((loss,) + output) if loss is not None else output
-
-        return CLIPOutput(
-            loss=loss,
-            logits_per_image=logits_per_image,
-            logits_per_text=logits_per_text,
-            text_embeds=text_embeds,
-            image_embeds=image_embeds,
-            text_model_output=None,
-            vision_model_output=None,
-        )

modules.json

@@ -3,6 +3,6 @@
     "idx": 0,
     "name": "0",
     "path": "",
-    "type": "jina-clip-implementation-st.6d5aa7d8b428eaba8d7908d86f43f5dd5ad6ad93.custom_st.Transformer"
+    "type": "tomaarsen/jina-clip-implementation-st--custom_st.Transformer"
   }
 ]

processing_clip.py

@@ -1,88 +0,0 @@
-# coding=utf-8
-#
-# Code mainly copied from:
-# https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/image_processing_clip.py
-# and adjusted for Jina CLIP
-
-from typing import Tuple, Union
-
-import torch
-from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
-from transformers.image_utils import ImageInput, make_list_of_images
-from transformers.models.clip import CLIPProcessor
-
-from .transform import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD, image_transform
-
-""" Jina CLIP processor implementation """
-
-
-class JinaCLIPProcessor(CLIPProcessor):
-    image_processor_class = 'AutoImageProcessor'
-    tokenizer_class = 'AutoTokenizer'
-
-
-""" Jina CLIP image processor implementation """
-
-
-class JinaCLIPImageProcessor(BaseImageProcessor):
-    model_input_names = ['pixel_values']
-    _valid_processor_keys = [
-        'size',
-        'mean',
-        'std',
-        'resize_mode',
-        'interpolation',
-        'fill_color',
-    ]
-
-    def __init__(
-        self,
-        size: Union[int, Tuple[int, int]] = 224,
-        mean: Union[float, Tuple[float]] = OPENAI_DATASET_MEAN,
-        std: Union[float, Tuple[float]] = OPENAI_DATASET_STD,
-        resize_mode: str = 'shortest',
-        interpolation: str = 'bicubic',
-        fill_color: int = 0,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-        self.size = size
-        self.mean = mean
-        self.std = std
-        self.resize_mode = resize_mode
-        self.interpolation = interpolation
-        self.fill_color = fill_color
-        self.transform = self._build_transform()
-
-    def _build_transform(self):
-        return image_transform(
-            image_size=self.size,
-            is_train=False,
-            mean=self.mean,
-            std=self.std,
-            resize_mode=self.resize_mode,
-            interpolation=self.interpolation,
-            fill_color=self.fill_color,
-            aug_cfg=None,
-        )
-
-    def to_dict(self):
-        output = super().to_dict()
-        output.pop('transform')
-        return output
-
-    def preprocess(self, images: ImageInput, **kwargs) -> BatchFeature:
-
-        _transform_needs_rebuild = False
-        for k, v in kwargs.items():
-            if k in self._valid_processor_keys:
-                if v != getattr(self, k):
-                    setattr(self, k, v)
-                    _transform_needs_rebuild = True
-
-        if _transform_needs_rebuild:
-            self.transform = self._build_transform()
-
-        images = make_list_of_images(images)
-        out = torch.stack([self.transform(image) for image in images], dim=0)
-        return BatchFeature(data={'pixel_values': out})

rope_embeddings.py

@@ -1,165 +0,0 @@
-# --------------------------------------------------------
-# Adapted from EVA CLIP
-# https://github.com/baaivision/EVA/tree/master/EVA-CLIP/rei/eva_clip
-# --------------------------------------------------------
-
-import logging
-from math import pi
-
-import torch
-from einops import rearrange, repeat
-from torch import nn
-
-
-def broadcast(tensors, dim=-1):
-    num_tensors = len(tensors)
-    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
-    assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
-    shape_len = list(shape_lens)[0]
-    dim = (dim + shape_len) if dim < 0 else dim
-    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
-    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
-    assert all(
-        [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]
-    ), 'invalid dimensions for broadcastable concatentation'
-    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
-    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
-    expanded_dims.insert(dim, (dim, dims[dim]))
-    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
-    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
-    return torch.cat(tensors, dim=dim)
-
-
-def rotate_half(x):
-    x = rearrange(x, '... (d r) -> ... d r', r=2)
-    x1, x2 = x.unbind(dim=-1)
-    x = torch.stack((-x2, x1), dim=-1)
-    return rearrange(x, '... d r -> ... (d r)')
-
-
-class VisionRotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim,
-        pt_seq_len,
-        ft_seq_len=None,
-        custom_freqs=None,
-        freqs_for='lang',
-        theta=10000,
-        max_freq=10,
-        num_freqs=1,
-    ):
-        super().__init__()
-        if custom_freqs:
-            freqs = custom_freqs
-        elif freqs_for == 'lang':
-            freqs = 1.0 / (
-                theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
-            )
-        elif freqs_for == 'pixel':
-            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
-        elif freqs_for == 'constant':
-            freqs = torch.ones(num_freqs).float()
-        else:
-            raise ValueError(f'unknown modality {freqs_for}')
-
-        if ft_seq_len is None:
-            ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs_h = torch.einsum('..., f -> ... f', t, freqs)
-        freqs_h = repeat(freqs_h, '... n -> ... (n r)', r=2)
-
-        freqs_w = torch.einsum('..., f -> ... f', t, freqs)
-        freqs_w = repeat(freqs_w, '... n -> ... (n r)', r=2)
-
-        freqs = broadcast((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
-
-        self.register_buffer('freqs_cos', freqs.cos())
-        self.register_buffer('freqs_sin', freqs.sin())
-
-        logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
-
-    def forward(self, t, start_index=0):
-        rot_dim = self.freqs_cos.shape[-1]
-        end_index = start_index + rot_dim
-        assert rot_dim <= t.shape[-1], (
-            f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in '
-            f'all the positions {rot_dim}'
-        )
-        t_left, t, t_right = (
-            t[..., :start_index],
-            t[..., start_index:end_index],
-            t[..., end_index:],
-        )
-        t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)
-
-        return torch.cat((t_left, t, t_right), dim=-1)
-
-
-class VisionRotaryEmbeddingFast(nn.Module):
-    def __init__(
-        self,
-        dim,
-        pt_seq_len,
-        ft_seq_len=None,
-        custom_freqs=None,
-        freqs_for='lang',
-        theta=10000,
-        max_freq=10,
-        num_freqs=1,
-        patch_dropout=0.0,
-    ):
-        super().__init__()
-        if custom_freqs:
-            freqs = custom_freqs
-        elif freqs_for == 'lang':
-            freqs = 1.0 / (
-                theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
-            )
-        elif freqs_for == 'pixel':
-            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
-        elif freqs_for == 'constant':
-            freqs = torch.ones(num_freqs).float()
-        else:
-            raise ValueError(f'unknown modality {freqs_for}')
-
-        if ft_seq_len is None:
-            ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum('..., f -> ... f', t, freqs)
-        freqs = repeat(freqs, '... n -> ... (n r)', r=2)
-        freqs = broadcast((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
-        freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
-
-        self.patch_dropout = patch_dropout
-
-        self.register_buffer('freqs_cos', freqs_cos)
-        self.register_buffer('freqs_sin', freqs_sin)
-
-        logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
-
-    def forward(self, t, patch_indices_keep=None):
-        if patch_indices_keep is not None:
-            batch = t.size()[0]
-            batch_indices = torch.arange(batch)
-            batch_indices = batch_indices[..., None]
-
-            freqs_cos = repeat(
-                self.freqs_cos, 'i j -> n i m j', n=t.shape[0], m=t.shape[1]
-            )
-            freqs_sin = repeat(
-                self.freqs_sin, 'i j -> n i m j', n=t.shape[0], m=t.shape[1]
-            )
-
-            freqs_cos = freqs_cos[batch_indices, patch_indices_keep]
-            freqs_cos = rearrange(freqs_cos, 'n i m j -> n m i j')
-            freqs_sin = freqs_sin[batch_indices, patch_indices_keep]
-            freqs_sin = rearrange(freqs_sin, 'n i m j -> n m i j')
-
-            return t * freqs_cos + rotate_half(t) * freqs_sin
-
-        return t * self.freqs_cos + rotate_half(t) * self.freqs_sin

transform.py

@@ -1,458 +0,0 @@
-import numbers
-import random
-import warnings
-from dataclasses import asdict, dataclass
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
-
-import torch
-import torchvision.transforms.functional as F
-from torchvision.transforms import (
-    CenterCrop,
-    ColorJitter,
-    Compose,
-    Grayscale,
-    InterpolationMode,
-    Normalize,
-    RandomResizedCrop,
-    Resize,
-    ToTensor,
-)
-from transformers.image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
-
-OPENAI_DATASET_MEAN = tuple(OPENAI_CLIP_MEAN)
-OPENAI_DATASET_STD = tuple(OPENAI_CLIP_STD)
-
-
-@dataclass
-class PreprocessCfg:
-    size: Union[int, Tuple[int, int]] = 224
-    mode: str = 'RGB'
-    mean: Tuple[float, ...] = OPENAI_DATASET_MEAN
-    std: Tuple[float, ...] = OPENAI_DATASET_STD
-    interpolation: str = 'bicubic'
-    resize_mode: str = 'shortest'
-    fill_color: int = 0
-
-    def __post_init__(self):
-        assert self.mode in ('RGB',)
-
-    @property
-    def num_channels(self):
-        return 3
-
-    @property
-    def input_size(self):
-        return (self.num_channels,) + (self.size, self.size)
-
-
-_PREPROCESS_KEYS = set(asdict(PreprocessCfg()).keys())
-
-
-def merge_preprocess_dict(
-    base: Union[PreprocessCfg, Dict],
-    overlay: Dict,
-):
-    """Merge overlay key-value pairs on top of base preprocess cfg or dict.
-    Input dicts are filtered based on PreprocessCfg fields.
-    """
-    if isinstance(base, PreprocessCfg):
-        base_clean = asdict(base)
-    else:
-        base_clean = {k: v for k, v in base.items() if k in _PREPROCESS_KEYS}
-    if overlay:
-        overlay_clean = {
-            k: v for k, v in overlay.items() if k in _PREPROCESS_KEYS and v is not None
-        }
-        base_clean.update(overlay_clean)
-    return base_clean
-
-
-def merge_preprocess_kwargs(base: Union[PreprocessCfg, Dict], **kwargs):
-    return merge_preprocess_dict(base, kwargs)
-
-
-@dataclass
-class AugmentationCfg:
-    scale: Tuple[float, float] = (0.9, 1.0)
-    ratio: Optional[Tuple[float, float]] = None
-    color_jitter: Optional[
-        Union[float, Tuple[float, float, float], Tuple[float, float, float, float]]
-    ] = None
-    re_prob: Optional[float] = None
-    re_count: Optional[int] = None
-    use_timm: bool = False
-
-    # params for simclr_jitter_gray
-    color_jitter_prob: float = None
-    gray_scale_prob: float = None
-
-
-def _setup_size(size, error_msg):
-    if isinstance(size, numbers.Number):
-        return int(size), int(size)
-
-    if isinstance(size, Sequence) and len(size) == 1:
-        return size[0], size[0]
-
-    if len(size) != 2:
-        raise ValueError(error_msg)
-
-    return size
-
-
-class ResizeKeepRatio:
-    """Resize and Keep Ratio
-
-    Copy & paste from `timm`
-    """
-
-    def __init__(
-        self,
-        size,
-        longest=0.0,
-        interpolation=InterpolationMode.BICUBIC,
-        random_scale_prob=0.0,
-        random_scale_range=(0.85, 1.05),
-        random_aspect_prob=0.0,
-        random_aspect_range=(0.9, 1.11),
-    ):
-        if isinstance(size, (list, tuple)):
-            self.size = tuple(size)
-        else:
-            self.size = (size, size)
-        self.interpolation = interpolation
-        self.longest = float(longest)  # [0, 1] where 0 == shortest edge, 1 == longest
-        self.random_scale_prob = random_scale_prob
-        self.random_scale_range = random_scale_range
-        self.random_aspect_prob = random_aspect_prob
-        self.random_aspect_range = random_aspect_range
-
-    @staticmethod
-    def get_params(
-        img,
-        target_size,
-        longest,
-        random_scale_prob=0.0,
-        random_scale_range=(0.85, 1.05),
-        random_aspect_prob=0.0,
-        random_aspect_range=(0.9, 1.11),
-    ):
-        """Get parameters"""
-        source_size = img.size[::-1]  # h, w
-        h, w = source_size
-        target_h, target_w = target_size
-        ratio_h = h / target_h
-        ratio_w = w / target_w
-        ratio = max(ratio_h, ratio_w) * longest + min(ratio_h, ratio_w) * (
-            1.0 - longest
-        )
-        if random_scale_prob > 0 and random.random() < random_scale_prob:
-            ratio_factor = random.uniform(random_scale_range[0], random_scale_range[1])
-            ratio_factor = (ratio_factor, ratio_factor)
-        else:
-            ratio_factor = (1.0, 1.0)
-        if random_aspect_prob > 0 and random.random() < random_aspect_prob:
-            aspect_factor = random.uniform(
-                random_aspect_range[0], random_aspect_range[1]
-            )
-            ratio_factor = (
-                ratio_factor[0] / aspect_factor,
-                ratio_factor[1] * aspect_factor,
-            )
-        size = [round(x * f / ratio) for x, f in zip(source_size, ratio_factor)]
-        return size
-
-    def __call__(self, img):
-        """
-        Args:
-            img (PIL Image): Image to be cropped and resized.
-
-        Returns:
-            PIL Image: Resized, padded to at least target size, possibly
-            cropped to exactly target size
-        """
-        size = self.get_params(
-            img,
-            self.size,
-            self.longest,
-            self.random_scale_prob,
-            self.random_scale_range,
-            self.random_aspect_prob,
-            self.random_aspect_range,
-        )
-        img = F.resize(img, size, self.interpolation)
-        return img
-
-    def __repr__(self):
-        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
-        format_string += f', interpolation={self.interpolation})'
-        format_string += f', longest={self.longest:.3f})'
-        return format_string
-
-
-def center_crop_or_pad(
-    img: torch.Tensor, output_size: List[int], fill=0
-) -> torch.Tensor:
-    """Center crops and/or pads the given image.
-    If the image is torch Tensor, it is expected
-    to have [..., H, W] shape, where ... means an arbitrary number of leading
-    dimensions. If image size is smaller than output size along any edge, image is
-    padded with 0 and then center cropped.
-
-    Args:
-        img (PIL Image or Tensor): Image to be cropped.
-        output_size (sequence or int): (height, width) of the crop box. If int or
-        sequence with single int, it is used for both directions.
-        fill (int, Tuple[int]): Padding color
-
-    Returns:
-        PIL Image or Tensor: Cropped image.
-    """
-    if isinstance(output_size, numbers.Number):
-        output_size = (int(output_size), int(output_size))
-    elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
-        output_size = (output_size[0], output_size[0])
-
-    _, image_height, image_width = F.get_dimensions(img)
-    crop_height, crop_width = output_size
-
-    if crop_width > image_width or crop_height > image_height:
-        padding_ltrb = [
-            (crop_width - image_width) // 2 if crop_width > image_width else 0,
-            (crop_height - image_height) // 2 if crop_height > image_height else 0,
-            (crop_width - image_width + 1) // 2 if crop_width > image_width else 0,
-            (crop_height - image_height + 1) // 2 if crop_height > image_height else 0,
-        ]
-        img = F.pad(img, padding_ltrb, fill=fill)
-        _, image_height, image_width = F.get_dimensions(img)
-        if crop_width == image_width and crop_height == image_height:
-            return img
-
-    crop_top = int(round((image_height - crop_height) / 2.0))
-    crop_left = int(round((image_width - crop_width) / 2.0))
-    return F.crop(img, crop_top, crop_left, crop_height, crop_width)
-
-
-class CenterCropOrPad(torch.nn.Module):
-    """Crops the given image at the center.
-    If the image is torch Tensor, it is expected
-    to have [..., H, W] shape, where ... means an arbitrary number of leading
-    dimensions. If image size is smaller than output size along any edge, image is
-    padded with 0 and then center cropped.
-
-    Args:
-        size (sequence or int): Desired output size of the crop. If size is an
-            int instead of sequence like (h, w), a square crop (size, size) is
-            made. If provided a sequence of length 1, it will be interpreted as
-            (size[0], size[0]).
-    """
-
-    def __init__(self, size, fill=0):
-        super().__init__()
-        self.size = _setup_size(
-            size, error_msg='Please provide only two dimensions (h, w) for size.'
-        )
-        self.fill = fill
-
-    def forward(self, img):
-        """
-        Args:
-            img (PIL Image or Tensor): Image to be cropped.
-
-        Returns:
-            PIL Image or Tensor: Cropped image.
-        """
-        return center_crop_or_pad(img, self.size, fill=self.fill)
-
-    def __repr__(self) -> str:
-        return f'{self.__class__.__name__}(size={self.size})'
-
-
-def _convert_to_rgb(image):
-    return image.convert('RGB')
-
-
-class _ColorJitter(object):
-    """
-    Apply Color Jitter to the PIL image with a specified probability.
-    """
-
-    def __init__(self, brightness=0.0, contrast=0.0, saturation=0.0, hue=0.0, p=0.8):
-        assert 0.0 <= p <= 1.0
-        self.p = p
-        self.transf = ColorJitter(
-            brightness=brightness, contrast=contrast, saturation=saturation, hue=hue
-        )
-
-    def __call__(self, img):
-        if random.random() < self.p:
-            return self.transf(img)
-        else:
-            return img
-
-
-class _GrayScale(object):
-    """
-    Apply Gray Scale to the PIL image with a specified probability.
-    """
-
-    def __init__(self, p=0.2):
-        assert 0.0 <= p <= 1.0
-        self.p = p
-        self.transf = Grayscale(num_output_channels=3)
-
-    def __call__(self, img):
-        if random.random() < self.p:
-            return self.transf(img)
-        else:
-            return img
-
-
-def image_transform(
-    image_size: Union[int, Tuple[int, int]],
-    is_train: bool,
-    mean: Optional[Tuple[float, ...]] = None,
-    std: Optional[Tuple[float, ...]] = None,
-    resize_mode: Optional[str] = None,
-    interpolation: Optional[str] = None,
-    fill_color: int = 0,
-    aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
-):
-    mean = mean or OPENAI_DATASET_MEAN
-    if not isinstance(mean, (list, tuple)):
-        mean = (mean,) * 3
-
-    std = std or OPENAI_DATASET_STD
-    if not isinstance(std, (list, tuple)):
-        std = (std,) * 3
-
-    interpolation = interpolation or 'bicubic'
-    assert interpolation in ['bicubic', 'bilinear', 'random']
-    # NOTE random is ignored for interpolation_mode, so defaults to BICUBIC for
-    # inference if set
-    interpolation_mode = (
-        InterpolationMode.BILINEAR
-        if interpolation == 'bilinear'
-        else InterpolationMode.BICUBIC
-    )
-
-    resize_mode = resize_mode or 'shortest'
-    assert resize_mode in ('shortest', 'longest', 'squash')
-
-    if isinstance(aug_cfg, dict):
-        aug_cfg = AugmentationCfg(**aug_cfg)
-    else:
-        aug_cfg = aug_cfg or AugmentationCfg()
-
-    normalize = Normalize(mean=mean, std=std)
-
-    if is_train:
-        aug_cfg_dict = {k: v for k, v in asdict(aug_cfg).items() if v is not None}
-        use_timm = aug_cfg_dict.pop('use_timm', False)
-        if use_timm:
-            from timm.data import create_transform  # timm can still be optional
-
-            if isinstance(image_size, (tuple, list)):
-                assert len(image_size) >= 2
-                input_size = (3,) + image_size[-2:]
-            else:
-                input_size = (3, image_size, image_size)
-
-            aug_cfg_dict.setdefault('color_jitter', None)  # disable by default
-            # drop extra non-timm items
-            aug_cfg_dict.pop('color_jitter_prob', None)
-            aug_cfg_dict.pop('gray_scale_prob', None)
-
-            train_transform = create_transform(
-                input_size=input_size,
-                is_training=True,
-                hflip=0.0,
-                mean=mean,
-                std=std,
-                re_mode='pixel',
-                interpolation=interpolation,
-                **aug_cfg_dict,
-            )
-        else:
-            train_transform = [
-                RandomResizedCrop(
-                    image_size,
-                    scale=aug_cfg_dict.pop('scale'),
-                    interpolation=InterpolationMode.BICUBIC,
-                ),
-                _convert_to_rgb,
-            ]
-            if aug_cfg.color_jitter_prob:
-                assert (
-                    aug_cfg.color_jitter is not None and len(aug_cfg.color_jitter) == 4
-                )
-                train_transform.extend(
-                    [_ColorJitter(*aug_cfg.color_jitter, p=aug_cfg.color_jitter_prob)]
-                )
-            if aug_cfg.gray_scale_prob:
-                train_transform.extend([_GrayScale(aug_cfg.gray_scale_prob)])
-            train_transform.extend(
-                [
-                    ToTensor(),
-                    normalize,
-                ]
-            )
-            train_transform = Compose(train_transform)
-            if aug_cfg_dict:
-                warnings.warn(
-                    f'Unused augmentation cfg items, specify `use_timm` to use '
-                    f'({list(aug_cfg_dict.keys())}).'
-                )
-        return train_transform
-    else:
-        if resize_mode == 'longest':
-            transforms = [
-                ResizeKeepRatio(
-                    image_size, interpolation=interpolation_mode, longest=1
-                ),
-                CenterCropOrPad(image_size, fill=fill_color),
-            ]
-        elif resize_mode == 'squash':
-            if isinstance(image_size, int):
-                image_size = (image_size, image_size)
-            transforms = [
-                Resize(image_size, interpolation=interpolation_mode),
-            ]
-        else:
-            assert resize_mode == 'shortest'
-            if not isinstance(image_size, (tuple, list)):
-                image_size = (image_size, image_size)
-            if image_size[0] == image_size[1]:
-                # simple case, use torchvision built-in Resize w/ shortest edge mode
-                # (scalar size arg)
-                transforms = [Resize(image_size[0], interpolation=interpolation_mode)]
-            else:
-                # resize shortest edge to matching target dim for non-square target
-                transforms = [ResizeKeepRatio(image_size)]
-            transforms += [CenterCrop(image_size)]
-
-        transforms.extend(
-            [
-                _convert_to_rgb,
-                ToTensor(),
-                normalize,
-            ]
-        )
-        return Compose(transforms)
-
-
-def image_transform_v2(
-    cfg: PreprocessCfg,
-    is_train: bool,
-    aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
-):
-    return image_transform(
-        image_size=cfg.size,
-        is_train=is_train,
-        mean=cfg.mean,
-        std=cfg.std,
-        interpolation=cfg.interpolation,
-        resize_mode=cfg.resize_mode,
-        fill_color=cfg.fill_color,
-        aug_cfg=aug_cfg,
-    )