Upload ProSSTXForMaskedLM

config.json

@@ -1,12 +1,12 @@
 {
   "_name_or_path": "/rgzn/limc/ProSST/oracle_checkpoint3/ss_2051_0_aa2pos_pos2aa_aa2ss_ss2aa_False/ProSSTX-2048",
   "architectures": [
-    "ProSSTForMaskedLM"
+    "ProSSTXForMaskedLM"
   ],
   "attention_probs_dropout_prob": 0.1,
   "auto_map": {
-    "AutoConfig": "configuration_prosst.ProSSTConfig",
-    "AutoModelForMaskedLM": "modeling_prosst.ProSSTForMaskedLM"
+    "AutoConfig": "configuration_prosst.ProSSTXConfig",
+    "AutoModelForMaskedLM": "modeling_prosst.ProSSTXForMaskedLM"
   },
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
@@ -18,7 +18,7 @@
   "max_position_embeddings": -1,
   "max_relative_positions": 1024,
   "mlm_probability": 0.15,
-  "model_type": "ProSST",
+  "model_type": "ProSSTX",
   "num_attention_heads": 16,
   "num_hidden_layers": 24,
   "pad_token_id": 0,

configuration_prosst.py

@@ -1,7 +1,7 @@
 from transformers import PretrainedConfig
 
-class ProSSTConfig(PretrainedConfig):
-    model_type = "ProSST"
+class ProSSTXConfig(PretrainedConfig):
+    model_type = "ProSSTX"
 
     def __init__(
         self,
@@ -68,4 +68,4 @@ class ProSSTConfig(PretrainedConfig):
         self.pooler_dropout = pooler_dropout
         self.pooler_hidden_act = pooler_hidden_act
         
-ProSSTConfig.register_for_auto_class()
+ProSSTXConfig.register_for_auto_class()

modeling_prosst.py

@@ -12,7 +12,7 @@ from transformers.modeling_outputs import (
     TokenClassifierOutput,
 )
 from transformers.modeling_utils import PreTrainedModel
-from .configuration_prosst import ProSSTConfig
+from .configuration_prosst import ProSSTXConfig
 import torch.nn.functional as F
 from functools import partial
 
@@ -262,7 +262,7 @@ class ContextPooler(nn.Module):
         return self.config.hidden_size
 
 
-class ProSSTLayerNorm(nn.Module):
+class ProSSTXLayerNorm(nn.Module):
     """LayerNorm module in the TF style (epsilon inside the square root)."""
 
     def __init__(self, size, eps=1e-12):
@@ -286,7 +286,7 @@ class ProSSTLayerNorm(nn.Module):
 
 class DisentangledSelfAttention(nn.Module):
 
-    def __init__(self, config: ProSSTConfig):
+    def __init__(self, config: ProSSTXConfig):
         super().__init__()
         self.config = config
         self.num_attention_heads = config.num_attention_heads
@@ -526,11 +526,11 @@ class DisentangledSelfAttention(nn.Module):
             return score, disentangled_attentions
 
 
-class ProSSTSelfOutput(nn.Module):
+class ProSSTXSelfOutput(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        self.LayerNorm = ProSSTLayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.LayerNorm = ProSSTXLayerNorm(config.hidden_size, config.layer_norm_eps)
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
 
     def forward(self, hidden_states, input_tensor):
@@ -540,12 +540,12 @@ class ProSSTSelfOutput(nn.Module):
         return hidden_states
 
 
-class ProSSTAttention(nn.Module):
+class ProSSTXAttention(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
         self.self = DisentangledSelfAttention(config)
-        self.output = ProSSTSelfOutput(config)
+        self.output = ProSSTXSelfOutput(config)
 
     def forward(
         self,
@@ -573,7 +573,7 @@ class ProSSTAttention(nn.Module):
             return attention_output
 
 
-class ProSSTIntermediate(nn.Module):
+class ProSSTXIntermediate(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
@@ -588,11 +588,11 @@ class ProSSTIntermediate(nn.Module):
         return hidden_states
 
 
-class ProSSTOutput(nn.Module):
+class ProSSTXOutput(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.LayerNorm = ProSSTLayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.LayerNorm = ProSSTXLayerNorm(config.hidden_size, config.layer_norm_eps)
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.config = config
 
@@ -603,13 +603,13 @@ class ProSSTOutput(nn.Module):
         return hidden_states
 
 
-class ProSSTLayer(nn.Module):
+class ProSSTXLayer(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
-        self.attention = ProSSTAttention(config)
-        self.intermediate = ProSSTIntermediate(config)
-        self.output = ProSSTOutput(config)
+        self.attention = ProSSTXAttention(config)
+        self.intermediate = ProSSTXIntermediate(config)
+        self.output = ProSSTXOutput(config)
 
     def forward(
         self,
@@ -638,13 +638,13 @@ class ProSSTLayer(nn.Module):
             return layer_output
 
 
-class ProSSTEncoder(nn.Module):
+class ProSSTXEncoder(nn.Module):
     """Modified BertEncoder with relative position bias support"""
 
     def __init__(self, config):
         super().__init__()
         self.layer = nn.ModuleList(
-            [ProSSTLayer(config) for _ in range(config.num_hidden_layers)]
+            [ProSSTXLayer(config) for _ in range(config.num_hidden_layers)]
         )
         self.relative_attention = config.relative_attention
         if self.relative_attention:
@@ -709,7 +709,7 @@ class ProSSTEncoder(nn.Module):
         )
 
 
-class ProSSTEmbeddings(nn.Module):
+class ProSSTXEmbeddings(nn.Module):
     """Construct the embeddings from word, position and token_type embeddings."""
 
     def __init__(self, config):
@@ -720,7 +720,7 @@ class ProSSTEmbeddings(nn.Module):
         self.word_embeddings = nn.Embedding(
             config.vocab_size, self.embedding_size, padding_idx=self.pad_token_id
         )
-        self.LayerNorm = ProSSTLayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.LayerNorm = ProSSTXLayerNorm(config.hidden_size, config.layer_norm_eps)
 
         # 绝对位置编码
         self.position_biased_input = config.position_biased_input
@@ -742,7 +742,7 @@ class ProSSTEmbeddings(nn.Module):
         # SS embeddings
         if config.ss_vocab_size > 0:
             self.ss_embeddings = nn.Embedding(config.ss_vocab_size, self.embedding_size)
-            self.ss_layer_norm = ProSSTLayerNorm(
+            self.ss_layer_norm = ProSSTXLayerNorm(
                 config.hidden_size, config.layer_norm_eps
             )
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
@@ -812,14 +812,14 @@ class ProSSTEmbeddings(nn.Module):
         return embeddings, None
 
 
-class ProSSTPreTrainedModel(PreTrainedModel):
+class ProSSTXPreTrainedModel(PreTrainedModel):
     """
     An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
     models.
     """
 
-    config_class = ProSSTConfig
-    base_model_prefix = "ProSST"
+    config_class = ProSSTXConfig
+    base_model_prefix = "ProSSTX"
     _keys_to_ignore_on_load_unexpected = ["position_embeddings"]
     supports_gradient_checkpointing = True
 
@@ -837,16 +837,16 @@ class ProSSTPreTrainedModel(PreTrainedModel):
                 module.weight.data[module.padding_idx].zero_()
 
     def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, ProSSTEncoder):
+        if isinstance(module, ProSSTXEncoder):
             module.gradient_checkpointing = value
 
 
-class ProSSTModel(ProSSTPreTrainedModel):
+class ProSSTXModel(ProSSTXPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.config = config
-        self.embeddings = ProSSTEmbeddings(config)
-        self.encoder = ProSSTEncoder(config)
+        self.embeddings = ProSSTXEmbeddings(config)
+        self.encoder = ProSSTXEncoder(config)
         self.post_init()
 
     def forward(
@@ -882,7 +882,7 @@ class ProSSTModel(ProSSTPreTrainedModel):
         )
 
 
-class ProSSTPredictionHeadTransform(nn.Module):
+class ProSSTXPredictionHeadTransform(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
@@ -900,11 +900,11 @@ class ProSSTPredictionHeadTransform(nn.Module):
         return hidden_states
 
 
-class ProSSTLMPredictionHead(nn.Module):
+class ProSSTXLMPredictionHead(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
-        self.transform = ProSSTPredictionHeadTransform(config)
+        self.transform = ProSSTXPredictionHeadTransform(config)
         self.embedding_size = config.hidden_size
         self.decoder = nn.Linear(self.embedding_size, config.vocab_size, bias=False)
 
@@ -914,24 +914,24 @@ class ProSSTLMPredictionHead(nn.Module):
         return hidden_states
 
 
-class ProSSTOnlyMLMHead(nn.Module):
+class ProSSTXOnlyMLMHead(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.predictions = ProSSTLMPredictionHead(config)
+        self.predictions = ProSSTXLMPredictionHead(config)
 
     def forward(self, sequence_output):
         prediction_scores = self.predictions(sequence_output)
         return prediction_scores
 
 
-class ProSSTPreTrainedModel(PreTrainedModel):
+class ProSSTXPreTrainedModel(PreTrainedModel):
     """
     An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
     models.
     """
 
-    config_class = ProSSTConfig
-    base_model_prefix = "ProSST"
+    config_class = ProSSTXConfig
+    base_model_prefix = "ProSSTX"
     _keys_to_ignore_on_load_unexpected = ["position_embeddings"]
     supports_gradient_checkpointing = True
 
@@ -949,11 +949,11 @@ class ProSSTPreTrainedModel(PreTrainedModel):
                 module.weight.data[module.padding_idx].zero_()
 
     def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, ProSSTEncoder):
+        if isinstance(module, ProSSTXEncoder):
             module.gradient_checkpointing = value
 
 
-class ProSSTForMaskedLM(ProSSTPreTrainedModel):
+class ProSSTXForMaskedLM(ProSSTXPreTrainedModel):
     _tied_weights_keys = [
         "cls.predictions.decoder.weight",
         "cls.predictions.decoder.bias",
@@ -961,8 +961,8 @@ class ProSSTForMaskedLM(ProSSTPreTrainedModel):
 
     def __init__(self, config):
         super().__init__(config)
-        self.prosst = ProSSTModel(config)
-        self.cls = ProSSTOnlyMLMHead(config)
+        self.prosst = ProSSTXModel(config)
+        self.cls = ProSSTXOnlyMLMHead(config)
         self.post_init()
 
     def forward(
@@ -1005,14 +1005,14 @@ class ProSSTForMaskedLM(ProSSTPreTrainedModel):
         )
 
 
-class ProSSTForSequenceClassification(ProSSTPreTrainedModel):
+class ProSSTXForSequenceClassification(ProSSTXPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
         num_labels = getattr(config, "num_labels", 2)
         self.num_labels = num_labels
         self.scale_hidden = getattr(config, "scale_hidden", 1)
-        self.prosst = ProSSTModel(config)
+        self.prosst = ProSSTXModel(config)
         self.pooler = ContextPooler(config)
         output_dim = self.pooler.output_dim * self.scale_hidden
 
@@ -1125,12 +1125,12 @@ class ProSSTForSequenceClassification(ProSSTPreTrainedModel):
         )
 
 
-class ProSSTForTokenClassification(ProSSTPreTrainedModel):
+class ProSSTXForTokenClassification(ProSSTXPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
 
-        self.prosst = ProSSTModel(config)
+        self.prosst = ProSSTXModel(config)
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
@@ -1190,9 +1190,9 @@ class ProSSTForTokenClassification(ProSSTPreTrainedModel):
         )
 
 
-ProSSTModel.register_for_auto_class("AutoModel")
-ProSSTForMaskedLM.register_for_auto_class("AutoModelForMaskedLM")
-ProSSTForSequenceClassification.register_for_auto_class(
+ProSSTXModel.register_for_auto_class("AutoModel")
+ProSSTXForMaskedLM.register_for_auto_class("AutoModelForMaskedLM")
+ProSSTXForSequenceClassification.register_for_auto_class(
     "AutoModelForSequenceClassification"
 )
-ProSSTForTokenClassification.register_for_auto_class("AutoModelForTokenClassification")
+ProSSTXForTokenClassification.register_for_auto_class("AutoModelForTokenClassification")