Upload ClipMDModel.py

ClipMDModel.py

@@ -0,0 +1,138 @@
+from transformers import CLIPModel
+import torch
+from typing import Optional, Tuple
+
+
+def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
+    return torch.nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))
+
+
+def clip_loss(logits_per_text: torch.Tensor) -> torch.Tensor:
+    caption_loss = contrastive_loss(logits_per_text)
+    image_loss = contrastive_loss(logits_per_text.T)
+    return (caption_loss + image_loss) / 2.0
+
+
+class ClipMDModel(CLIPModel):
+
+    def embed_text(self,
+                   input_ids:torch.LongTensor,
+                   attention_mask:torch.LongTensor,
+                   output_attentions: Optional[bool] = None,
+                   output_hidden_states: Optional[bool] = None,
+                   position_ids: Optional[torch.LongTensor] = None,
+                  ):
+        """
+        :param input_ids: tokenized text from CLIPProcessor.
+        :param attention_mask: attention mask from CLIPProcessor.
+        :return: text embeddings of input_ids (tokens longer then 77 tokens
+        is embeded using a sliding window and pooling).
+        """
+        tokens = []
+        masks = []
+        pos = []
+
+        for i in range(input_ids.size()[0]):
+            ten = input_ids[i]
+            mask = attention_mask[i]
+            mask = mask[mask.nonzero().flatten()]
+            ten = ten[:mask.size()[0]]
+
+            if not pos:
+                pos.append([0, 0])
+            else:
+                pos.append([pos[-1][1], pos[-1][1]])
+
+            #spliting tokenized text into input sized chunks with an overlapping window.
+            if ten.size()[0]>77:
+                tokens.append(ten.unfold(dimension = 0,size = 77, step = 70))
+                masks.append(mask.unfold(dimension = 0,size = 77, step = 70))
+
+                pos[-1][1]+=tokens[-1].size()[0]
+
+                ten=ten[tokens[-1].size()[0]*70:]
+                mask=mask[tokens[-1].size()[0]*70:]
+
+            if ten.size()[0] > 0:
+                new_mask = torch.zeros((1, 77)).to(self.device)
+                new_mask[:, 0:mask.size()[0]] = mask
+
+                new_ten = torch.full((1, 77), 49407).to(self.device)
+                new_ten[:, 0:ten.size()[0]] = ten
+
+                tokens.append(new_ten)
+                masks.append(new_mask)
+                pos[-1][1] += 1
+        #encoding the tokenized text
+        embedded = self.get_text_features(input_ids=torch.cat(tokens, 0),
+                                          attention_mask=torch.cat(masks, 0),
+                                         output_attentions=output_attentions,
+                                         output_hidden_states=output_hidden_states,
+                                         position_ids=position_ids,
+                                         )
+        
+        #pooling the embeddings of segments that came from the same original text
+        embeddings = []
+        for p in pos:
+            if p[1] - p[0] == 1:
+                embeddings.append(embedded[p[0]].unsqueeze(0))
+            else:
+                embeddings.append(torch.mean(embedded[p[0]:p[1]], dim=0).unsqueeze(0))
+
+        return torch.cat(embeddings, 0)
+
+    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,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Tuple:
+        """
+        :param input_ids: tokenized text from CLIPProcessor.
+        :param attention_mask: attention mask from CLIPProcessor.
+        :param pixel_values: pixel values from CLIPProcessor.
+        :param return_loss: boolean that indicates if loss should be returned
+        :return: image-caption cosine similarity as logits per image and per caption (also loss if return_loss is true)
+        """
+        # Use CLIP model's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = self.config.use_return_dict
+        
+        #encoding the images 
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        image_embeds = vision_outputs[1]
+        image_embeds = self.visual_projection(image_embeds)
+
+        #encoding the text captions
+        text_embeds =self.embed_text(input_ids=input_ids,
+                                     attention_mask=attention_mask,
+                                     output_attentions=output_attentions,
+                                     output_hidden_states=output_hidden_states,
+                                     position_ids=position_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
+
+        if return_loss:
+            loss = clip_loss(logits_per_text)
+            return logits_per_image,logits_per_text,loss
+        return logits_per_image,logits_per_text