Upload modeling_dplm2.py with huggingface_hub

modeling_dplm2.py

@@ -1819,7 +1819,7 @@ class DPLM2Model(DPLM2PreTrainedModel, EmbeddingMixin):
         )
 
 
-class DPLM2ForMaskedLM(DPLM2PreTrainedModel, EmbeddingMixin):
+class DPLM2ForMaskedLM(FastPLMTestTimeTrainingMixin, DPLM2PreTrainedModel, EmbeddingMixin):
     config_class = DPLM2Config
     def __init__(self, config, dropout: float = 0.1, vocab_size: Optional[int] = None):
         config.hidden_dropout_prob = dropout
@@ -1835,6 +1835,7 @@ class DPLM2ForMaskedLM(DPLM2PreTrainedModel, EmbeddingMixin):
         self.tokenizer = self.__class__.tokenizer
         if isinstance(config._name_or_path, str) and len(config._name_or_path) > 0:
             self.tokenizer = EsmTokenizer.from_pretrained(config._name_or_path)
+        self.init_ttt({"lora_target_replace_module": "ModifiedEsmAttention"})
 
     def get_input_embeddings(self) -> nn.Module:
         return self.esm.get_input_embeddings()
@@ -1874,6 +1875,9 @@ class DPLM2ForMaskedLM(DPLM2PreTrainedModel, EmbeddingMixin):
             store_all_hidden_states=store_all_hidden_states,
         )
 
+    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
+        return [self.esm]
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,