Update modeling_clip_masked_lm.py

modeling_clip_masked_lm.py

@@ -2,33 +2,38 @@ from typing import Optional, Tuple, Union
 
 import torch
 from torch import nn
-from transformers import CLIPTextConfig, CLIPTextModel
+from transformers import CLIPTextConfig
 from transformers.modeling_outputs import MaskedLMOutput
-from transformers.models.clip.modeling_clip import CLIPPreTrainedModel
+from transformers.models.clip.modeling_clip import (
+    CLIPPreTrainedModel,
+    CLIPTextTransformer,
+)
 from transformers.models.roberta.modeling_roberta import RobertaLMHead
 
 
 class CLIPTextModelForMaskedLM(CLIPPreTrainedModel):
     config_class = CLIPTextConfig
 
+    _no_split_modules = ["CLIPEncoderLayer"]
+
     def __init__(self, config: CLIPTextConfig):
         super().__init__(config)
-        self.clip_text_model = CLIPTextModel(config)
+        self.text_model = CLIPTextTransformer(config)
         self.lm_head = RobertaLMHead(config)
 
         self.post_init()
 
-    def get_input_embeddings(self):
-        return self.clip_text_model.text_model.embeddings.token_embedding
+    def get_input_embeddings(self) -> nn.Module:
+        return self.text_model.embeddings.token_embedding
 
-    def set_input_embeddings(self, value):
-        self.clip_text_model.text_model.embeddings.token_embedding = value
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.text_model.embeddings.token_embedding = value
 
-    def get_output_embeddings(self):
+    def get_output_embeddings(self) -> nn.Module:
         return self.lm_head.decoder
 
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head.decoder = new_embeddings
+    def set_output_embeddings(self, value: nn.Module) -> None:
+        self.lm_head.decoder = value
 
     def forward(
         self,
@@ -44,7 +49,7 @@ class CLIPTextModelForMaskedLM(CLIPPreTrainedModel):
             return_dict if return_dict is not None else self.config.use_return_dict
         )
 
-        outputs = self.clip_text_model(
+        outputs = self.text_model(
             input_ids=input_ids,
             attention_mask=attention_mask,
             position_ids=position_ids,