Upload modeling_esm_plusplus.py with huggingface_hub

modeling_esm_plusplus.py

@@ -1949,7 +1949,7 @@ class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin):
             s_max=transformer_output.s_max,
         )
 
-class ESMplusplusForMaskedLM(PreTrainedESMplusplusModel, EmbeddingMixin):
+class ESMplusplusForMaskedLM(FastPLMTestTimeTrainingMixin, PreTrainedESMplusplusModel, EmbeddingMixin):
     """
     ESM++ model for masked language modeling.
     Implements the base ESM++ architecture with a masked language modeling head.
@@ -1971,6 +1971,7 @@ class ESMplusplusForMaskedLM(PreTrainedESMplusplusModel, EmbeddingMixin):
         self.ce_loss = nn.CrossEntropyLoss()
         self.tokenizer = EsmSequenceTokenizer()
         self.init_weights()
+        self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"})
 
     def get_input_embeddings(self):
         return self.embed
@@ -2006,6 +2007,9 @@ class ESMplusplusForMaskedLM(PreTrainedESMplusplusModel, EmbeddingMixin):
             store_all_hidden_states=store_all_hidden_states,
         )
 
+    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
+        return [self.transformer]
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,