Upload 2 files

config.json

@@ -4,6 +4,9 @@
   "architectures": [
     "OpenCLIPVisionTextDualEncoderModel"
   ],
+  "auto_map": {
+    "AutoModel": "modeling_clip.OpenCLIPVisionTextDualEncoderModel"
+  },
   "logit_scale_init_value": 2.6592,
   "model_type": "vision-text-dual-encoder",
   "projection_dim": 1024,
@@ -82,7 +85,7 @@
     "top_p": 1.0,
     "torch_dtype": null,
     "torchscript": false,
-    "transformers_version": "4.27.0.dev0",
+    "transformers_version": "4.24.0",
     "type_vocab_size": 1,
     "typical_p": 1.0,
     "use_bfloat16": false,
@@ -104,6 +107,7 @@
     "decoder_start_token_id": null,
     "diversity_penalty": 0.0,
     "do_sample": false,
+    "dropout": 0.0,
     "early_stopping": false,
     "encoder_no_repeat_ngram_size": 0,
     "eos_token_id": null,
@@ -146,7 +150,6 @@
     "patch_size": 14,
     "prefix": null,
     "problem_type": null,
-    "projection_dim": 512,
     "pruned_heads": {},
     "remove_invalid_values": false,
     "repetition_penalty": 1.0,
@@ -164,7 +167,7 @@
     "top_p": 1.0,
     "torch_dtype": null,
     "torchscript": false,
-    "transformers_version": "4.27.0.dev0",
+    "transformers_version": "4.24.0",
     "typical_p": 1.0,
     "use_bfloat16": false
   }

modeling_clip.py

@@ -0,0 +1,129 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, VisionTextDualEncoderConfig, VisionTextDualEncoderModel
+from transformers.models.vision_text_dual_encoder.modeling_vision_text_dual_encoder import clip_loss, CLIPOutput
+
+
+class MeanPooler(nn.Module):
+    """Mean pooling"""
+
+    def forward(self, x, attention_mask):
+        masked_output = x.last_hidden_state * attention_mask.unsqueeze(-1)
+        return masked_output.sum(dim=1) / attention_mask.sum(-1, keepdim=True)
+
+
+class OpenCLIPVisionTextDualEncoderModel(VisionTextDualEncoderModel):
+    def __init__(
+        self,
+        config: Optional[VisionTextDualEncoderConfig] = None,
+        vision_model: Optional[PreTrainedModel] = None,
+        text_model: Optional[PreTrainedModel] = None,
+        add_text_model_pooling_layer: bool = False,
+    ):
+        super().__init__(config, vision_model, text_model)
+
+        # Remove text pooling layer
+        if not add_text_model_pooling_layer:
+            self.text_model.pooler = None
+
+        # Add mean pooling
+        self.pooler = MeanPooler()
+        # Overwrite text projection
+        hidden_size = (self.text_embed_dim + self.projection_dim) // 2
+        self.text_projection = nn.Sequential(
+            nn.Linear(self.text_embed_dim, hidden_size, bias=False),
+            nn.GELU(),
+            nn.Linear(hidden_size, self.projection_dim, bias=False),
+        )
+
+    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,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = self.pooler(text_outputs, attention_mask)
+        text_features = self.text_projection(pooled_output)
+
+        return text_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
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        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,
+        )
+
+        image_embeds = vision_outputs[1]  # pooler_output
+        image_embeds = self.visual_projection(image_embeds)
+
+        pooled_output = self.pooler(text_outputs, attention_mask)
+        text_embeds = self.text_projection(pooled_output)
+
+        # 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_outputs,
+            vision_model_output=vision_outputs,
+        )