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