Upload modeling_dplm2.py with huggingface_hub

modeling_dplm2.py

@@ -12,7 +12,7 @@ from torch.nn import functional as F
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple, Union
 
-from transformers import EsmTokenizer
+from transformers import AutoTokenizer, EsmTokenizer
 from transformers.modeling_outputs import (
     BaseModelOutputWithPastAndCrossAttentions,
     BaseModelOutputWithPoolingAndCrossAttentions,
@@ -694,6 +694,9 @@ class DPLM2ForMaskedLM(DPLM2PreTrainedModel, EmbeddingMixin):
         self.loss_fct = nn.CrossEntropyLoss()
         self.post_init()
         self.pad_id = config.pad_token_id
+        self.tokenizer = self.__class__.tokenizer
+        if isinstance(config._name_or_path, str) and len(config._name_or_path) > 0:
+            self.tokenizer = AutoTokenizer.from_pretrained(config._name_or_path)
 
     def get_input_embeddings(self) -> nn.Module:
         return self.esm.embeddings.word_embeddings