Update modeling_gptbert.py

modeling_gptbert.py

@@ -138,6 +138,24 @@ class Embedding(nn.Module):
         return self.dropout(word_embedding)
 
 
+class LMClassifier(nn.Module):
+    def __init__(self, config: GptBertConfig, n_labels: int):
+        super().__init__()
+
+        self.pre_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.projection = CastedLinearIn(config.hidden_size, config.hidden_size, bias=False)
+        self.post_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.emb2vocab = CastedLinearIn(config.hidden_size, n_labels, bias=True)
+
+    def forward(self, x: torch.Tensor):
+        x = self.pre_norm(x.float()).type_as(x)
+        x = self.projection(x)
+        x = gelu_new(x)
+        x = self.post_norm(x.float()).type_as(x)
+        x = self.emb2vocab(x)
+        return x
+
+
 class Classifier(nn.Module):
     def __init__(self, config: GptBertConfig, n_labels: int):
         super().__init__()
@@ -146,7 +164,7 @@ class Classifier(nn.Module):
         self.projection = CastedLinearIn(config.hidden_size, config.hidden_size, bias=False)
         self.post_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
         self.dropout = nn.Dropout(config.classifier_dropout)
-        self.emb2vocab = CastedLinearIn(config.hidden_size, n_labels, bias=True)
+        self.output_projection = CastedLinearIn(config.hidden_size, n_labels, bias=True)
 
     def forward(self, x: torch.Tensor):
         x = self.pre_norm(x.float()).type_as(x)
@@ -154,7 +172,7 @@ class Classifier(nn.Module):
         x = gelu_new(x)
         x = self.post_norm(x.float()).type_as(x)
         x = self.dropout(x)
-        x = self.emb2vocab(x)
+        x = self.output_projection(x)
         return x
 
 
@@ -571,7 +589,7 @@ class GptBertModel(GptBertPreTrainedModel):
 
         self.embedding = Embedding(config)
         self.encoder = Encoder(config)
-        self.classifier = Classifier(config, config.vocab_size) if add_mlm_layer else None
+        self.classifier = LMClassifier(config, config.vocab_size) if add_mlm_layer else None
         self.set_window_length(config)
         self.gradient_checkpointing = False
         self.post_init()