Upload model

config.json

@@ -0,0 +1,179 @@
+{
+  "_commit_hash": "e9754cb57664705bacb62145ea8a977f269a456b",
+  "_name_or_path": "flavour/vtde-dinov2-small-bge-small-en",
+  "architectures": [
+    "VTDEModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modelling.VTDEConfig",
+    "AutoModel": "modelling.VTDEModel"
+  },
+  "logit_scale_init_value": 2.6592,
+  "model_type": "vtde",
+  "projection_dim": 384,
+  "text_config": {
+    "_name_or_path": "BAAI/bge-small-en",
+    "add_cross_attention": false,
+    "architectures": [
+      "BertModel"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "classifier_dropout": null,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 384,
+    "id2label": {
+      "0": "LABEL_0"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 1536,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0
+    },
+    "layer_norm_eps": 1e-12,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 512,
+    "min_length": 0,
+    "model_type": "bert",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 12,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 12,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 0,
+    "position_embedding_type": "absolute",
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float32",
+    "torchscript": false,
+    "transformers_version": "4.32.0.dev0",
+    "type_vocab_size": 2,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 30522
+  },
+  "text_pooling_mode": "mean",
+  "torch_dtype": "float32",
+  "transformers_version": null,
+  "vision_config": {
+    "_name_or_path": "facebook/dinov2-small",
+    "add_cross_attention": false,
+    "architectures": [
+      "Dinov2Model"
+    ],
+    "attention_probs_dropout_prob": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "drop_path_rate": 0.0,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 384,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 518,
+    "initializer_range": 0.02,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "layerscale_value": 1.0,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "mlp_ratio": 4,
+    "model_type": "dinov2",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 6,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 12,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "qkv_bias": true,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float32",
+    "torchscript": false,
+    "transformers_version": "4.32.0.dev0",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_swiglu_ffn": false
+  },
+  "vision_pooling_mode": "max"
+}

model.safetensors

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

modelling.py

@@ -0,0 +1,174 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/12_modelling.ipynb.
+
+# %% auto 0
+__all__ = ['VTDEConfig', 'VTDEModel']
+
+# %% ../notebooks/12_modelling.ipynb 1
+from transformers.models.clip.modeling_clip import CLIPOutput, clip_loss
+from typing import Optional, Tuple, Union
+from transformers import VisionTextDualEncoderConfig, AutoModel, PreTrainedModel, VisionTextDualEncoderModel
+import torch
+
+class VTDEConfig(VisionTextDualEncoderConfig):
+    model_type = "vtde"
+    
+    def __init__(self, projection_dim=512, logit_scale_init_value=2.6592, 
+    text_pooling_mode='mean',
+    vision_pooling_mode='max',
+    **kwargs):
+        """
+        pooling_mode in ['mean', 'max', 'cls']
+        https://arxiv.org/pdf/2210.09996.pdf
+        https://github.com/kahnchana/clippy/blob/3c102c29c32f7c66c6e52e09b795fe9c061bbb03/src/open_clip/hf_model.py#L56
+        """
+        self.text_pooling_mode = text_pooling_mode
+        self.vision_pooling_mode = vision_pooling_mode
+        super().__init__(projection_dim, logit_scale_init_value, **kwargs)
+
+class VTDEModel(VisionTextDualEncoderModel):
+    config_class = VTDEConfig
+    base_model_prefix = "vtde"
+
+    def __init__(
+        self,
+        config: Optional[VTDEConfig] = None,
+        vision_model: Optional[PreTrainedModel] = None,
+        text_model: Optional[PreTrainedModel] = None,
+    ):
+        # You can customize the constructor if needed
+        super().__init__(config, vision_model, text_model)
+        self.text_pooling_mode = config.text_pooling_mode
+        self.vision_pooling_mode = config.vision_pooling_mode
+
+    def get_text_features(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        token_type_ids=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        if self.text_pooling_mode == 'cls':
+            pooled_output = text_outputs[1]
+        elif self.text_pooling_mode == 'mean':
+            pooled_output = torch.mean(text_outputs[0], dim=1)
+        elif self.text_pooling_mode == 'max':
+            pooled_output = torch.max(text_outputs[0], dim=1)[0]
+        elif self.text_pooling_mode == 'norm':
+            """we select the patch with the largest norm"""
+            last_hidden_states = text_outputs[0]
+            patch_norms = torch.norm(last_hidden_states[:, 1:, :], dim=-1)
+            max_norm_idx = torch.argmax(patch_norms, dim=1)
+            pooled_output = last_hidden_states[:, max_norm_idx, :][:, 0, :]
+        else:
+            "We want to raise the name of the pooling mode"
+            raise NotImplementedError
+
+        text_features = self.text_projection(pooled_output)
+
+        return text_features
+
+    def get_image_features(
+        self,
+        pixel_values=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        if self.vision_pooling_mode == 'cls':
+            pooled_output = vision_outputs[1]
+        elif self.vision_pooling_mode == 'mean':
+            pooled_output = torch.mean(vision_outputs[0], dim=1)
+        elif self.vision_pooling_mode == 'max':
+            pooled_output = torch.max(vision_outputs[0], dim=1)[0]
+        elif self.vision_pooling_mode == 'norm':
+            """we select the patch with the largest norm"""
+            last_hidden_states = vision_outputs[0]
+            patch_norms = torch.norm(last_hidden_states[:, 1:, :], dim=-1)
+            max_norm_idx = torch.argmax(patch_norms, dim=1)
+            pooled_output = last_hidden_states[:, max_norm_idx, :][:, 0, :]
+        else:
+            raise NotImplementedError
+
+        image_features = self.visual_projection(pooled_output)
+
+        return image_features
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        return_loss: Optional[bool] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], CLIPOutput]:
+
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        image_embeds = self.get_image_features(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        text_embeds = self.get_text_features(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        # normalized features
+        image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
+        text_embeds = text_embeds / text_embeds.norm(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, text_outputs, vision_outputs)
+            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=text_embeds,
+            vision_model_output=image_embeds,
+        )
+    
+VTDEConfig.register_for_auto_class()
+VTDEModel.register_for_auto_class("AutoModel")