add

modeling_clip.py

@@ -39,7 +39,6 @@ from transformers.utils import (
 )
 from .configuration_clip import CLIPConfig, CLIPTextConfig, CLIPVisionConfig
 
-
 if is_flash_attn_2_available():
     from transformers.modeling_flash_attention_utils import _flash_attention_forward
 
@@ -603,16 +602,15 @@ class CLIPPreTrainedModel(PreTrainedModel):
             fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
             nn.init.normal_(module.fc1.weight, std=fc_std)
             nn.init.normal_(module.fc2.weight, std=in_proj_std)
-        elif isinstance(module, CLIPModel):
-            pass
+        elif isinstance(module, LLM2CLIPModel):
             # nn.init.normal_(
             #     module.text_projection.weight,
             #     std=module.text_embed_dim**-0.5 * self.config.initializer_factor,
             # )
-            # nn.init.normal_(
-            #     module.visual_projection.weight,
-            #     std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
-            # )
+            nn.init.normal_(
+                module.visual_projection.weight,
+                std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
+            )
         elif isinstance(module, CLIPVisionModelWithProjection):
             nn.init.normal_(
                 module.visual_projection.weight,
@@ -1112,80 +1110,97 @@ class CLIPVisionModel(CLIPPreTrainedModel):
 
 
 @add_start_docstrings(CLIP_START_DOCSTRING)
-class CLIPModel(CLIPPreTrainedModel):
+class LLM2CLIPModel(CLIPPreTrainedModel):
     config_class = CLIPConfig
     _no_split_modules = ["CLIPTextEmbeddings", "CLIPEncoderLayer", "CLIPVisionEmbeddings"]
 
     def __init__(self, config: CLIPConfig):
         super().__init__(config)
+        # if not isinstance(config.text_config, CLIPTextConfig):
+        #     raise TypeError(
+        #         "config.text_config is expected to be of type CLIPTextConfig but is of type"
+        #         f" {type(config.text_config)}."
+        #     )
+
         if not isinstance(config.vision_config, CLIPVisionConfig):
             raise TypeError(
                 "config.vision_config is expected to be of type CLIPVisionConfig but is of type"
                 f" {type(config.vision_config)}."
             )
 
+        # text_config = config.text_config
         vision_config = config.vision_config
 
         self.projection_dim = config.projection_dim
+        # self.text_embed_dim = text_config.hidden_size
         self.vision_embed_dim = vision_config.hidden_size
+        
+        adapter = LLM2CLIP_Adapter()
+        self.text_adapter = adapter
+
+        # text_model = CLIPTextModel._from_config(text_config, attn_implementation=config._attn_implementation)
+        # self.text_model = text_model.text_model
 
         vision_model = CLIPVisionModel._from_config(vision_config, attn_implementation=config._attn_implementation)
         self.vision_model = vision_model.vision_model
 
-        # self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)
-        scale = self.vision_embed_dim ** -0.5
-        self.visual_projection =  nn.Parameter(scale * torch.randn(self.vision_embed_dim, self.projection_dim))
+        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)
         self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
 
         # Initialize weights and apply final processing
         self.post_init()
-
-    @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING)
-    def get_text_features(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> torch.FloatTensor:
-        r"""
-        Returns:
-            text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
-            applying the projection layer to the pooled output of [`CLIPTextModel`].
-
-        Examples:
-
-        ```python
-        >>> from transformers import AutoTokenizer, CLIPModel
-
-        >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-        >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
-
-        >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
-        >>> text_features = model.get_text_features(**inputs)
-        ```"""
-        # 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 = return_dict if return_dict is not None else self.config.use_return_dict
-
-        text_outputs = self.text_model(
-            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,
-        )
-
-        pooled_output = text_outputs[1]
-        text_features = self.text_projection(pooled_output)
-
-        return text_features
+        
+    def get_text_features(self, inputs):
+        #TODO: make this more flexible and configurable
+        return self.text_adapter(inputs)  
+    
+    # @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING)
+    # def get_text_features(
+    #     self,
+    #     input_ids: Optional[torch.Tensor] = None,
+    #     attention_mask: Optional[torch.Tensor] = None,
+    #     position_ids: Optional[torch.Tensor] = None,
+    #     output_attentions: Optional[bool] = None,
+    #     output_hidden_states: Optional[bool] = None,
+    #     return_dict: Optional[bool] = None,
+    # ) -> torch.FloatTensor:
+    #     r"""
+    #     Returns:
+    #         text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
+    #         applying the projection layer to the pooled output of [`CLIPTextModel`].
+
+    #     Examples:
+
+    #     ```python
+    #     >>> from transformers import AutoTokenizer, CLIPModel
+
+    #     >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+    #     >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+
+    #     >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
+    #     >>> text_features = model.get_text_features(**inputs)
+    #     ```"""
+    #     # 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 = return_dict if return_dict is not None else self.config.use_return_dict
+
+    #     text_outputs = self.text_model(
+    #         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,
+    #     )
+
+    #     pooled_output = text_outputs[1]
+    #     text_features = self.text_projection(pooled_output)
+
+    #     return text_features
 
     @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING)
     def get_image_features(
@@ -1232,7 +1247,7 @@ class CLIPModel(CLIPPreTrainedModel):
         )
 
         pooled_output = vision_outputs[1]  # pooled_output
-        image_features = pooled_output @ self.visual_projection
+        image_features = self.visual_projection(pooled_output)
 
         return image_features
 
@@ -1413,7 +1428,40 @@ class CLIPTextModelWithProjection(CLIPPreTrainedModel):
             attentions=text_outputs.attentions,
         )
 
-
+class LinearBlock(nn.Module):
+    def __init__(self, dim, expansion_factor=4, dropout=0.,norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.fn = nn.Sequential(
+            nn.Linear(dim, int(expansion_factor * dim)),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(int(expansion_factor * dim), dim),
+        )
+        self.ln = norm_layer(dim)
+        
+    def forward(self, x):
+        return x + self.fn(self.ln(x))
+       
+class LLM2CLIP_Adapter(nn.Module):
+    def __init__(self):
+        super().__init__()
+        #TODO: make this more flexible and configurable
+        # hard-coded values from the  LLM2CLIP model
+        text_embedding_dim = 4096
+        expansion_factor = 2
+        adaptor_num_layers = 4
+        proj_bias = True
+        output_dim = 1280
+        self.adaptor = nn.Sequential(
+            *[LinearBlock(text_embedding_dim, expansion_factor) for _ in range(adaptor_num_layers)],
+            nn.LayerNorm(text_embedding_dim),
+            nn.Linear(text_embedding_dim, output_dim, bias=proj_bias),
+        )
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = torch.nn.functional.normalize(hidden_states, p=2, dim=1)
+        hidden_states = self.adaptor(hidden_states)
+        return hidden_states
+    
 @add_start_docstrings(
     """
     CLIP Vision Model with a projection layer on top (a linear layer on top of the pooled output).