Rename flaubert2_model.py to jargon_model.py

flaubert2_model.py → jargon_model.py

@@ -1,42 +1,44 @@
-#from transformers import RobertaModel, RobertaConfig, RobertaForMaskedLM, RobertaLMHead
-#from linformer import LinformerTransformerEncoder, LinformerTransformerEncoderLayer, LinformerTransformerEncoderFS, LinformerTransformerEncoderLayerFS
-#import linformer
-from .linformer import LinformerTransformerEncoderLayer
-from .flaubert2_configuration import Flaubert2Config
-from transformers.models.roberta.modeling_roberta import RobertaEncoder, RobertaConfig, RobertaModel, RobertaLMHead, RobertaForMaskedLM, RobertaEmbeddings, RobertaForTokenClassification, RobertaForSequenceClassification
-import torch.nn as nn
 import math
-import torch.nn.functional as F
-from torch.nn import LayerNorm
-import torch
 from typing import List, Optional, Tuple, Union
 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import LayerNorm
 from fairseq.models.roberta import (
     RobertaModel as RobertModel, 
     RobertaEncoder as RobertaEncoderFS
 )
-
+from transformers.models.roberta.modeling_roberta import (
+    RobertaEncoder,
+    RobertaConfig,
+    RobertaModel,
+    RobertaLMHead,
+    RobertaForMaskedLM,
+    RobertaEmbeddings,
+    RobertaForTokenClassification,
+    RobertaForSequenceClassification
+)
 from transformers.modeling_outputs import (
     MaskedLMOutput,   
     BaseModelOutputWithPastAndCrossAttentions,
     BaseModelOutputWithPoolingAndCrossAttentions,
 )
 
+from .linformer import LinformerTransformerEncoderLayer
+from .jargon_configuration import JargonConfig
+
 
-class  Flaubert2ModelForSequenceClassification(RobertaForSequenceClassification):
+class JargonForSequenceClassification(RobertaForSequenceClassification):
 
-    config_class = Flaubert2Config
-    auto_map = {"test": "test3"}
+    config_class = JargonConfig
     
     def __init__(self, config,  **kwargs):
-        base_model_prefix = "flaubert2"
+        base_model_prefix = "jargon"
         
         super().__init__(config, **kwargs)
 
-        #self.encoder = Flaubert2Model(config, add_pooling_layer=False)
-        self.roberta = Flaubert2Model(config, add_pooling_layer=False)
-        #self.encoder = LinformerTransformerEncoder(config)
-        #self.encoder = LinformerTransformerEncoder(config)
+        self.roberta = JargonModel(config, add_pooling_layer=False)
         self.sbo_head = self.build_sbo_head(config)
 
     def build_sbo_head(self, config):
@@ -50,19 +52,16 @@ class  Flaubert2ModelForSequenceClassification(RobertaForSequenceClassification)
         )
 
 
-class Flaubert2ModelForTokenClassification(RobertaForTokenClassification):
+class JargonForTokenClassification(RobertaForTokenClassification):
 
-    config_class = Flaubert2Config
+    config_class = JargonConfig
     
     def __init__(self, config,  **kwargs):
-        base_model_prefix = "flaubert2"
+        base_model_prefix = "jargon"
         
         super().__init__(config, **kwargs)
 
-        #self.encoder = Flaubert2Model(config, add_pooling_layer=False)
-        self.roberta = Flaubert2Model(config, add_pooling_layer=False)
-        #self.encoder = LinformerTransformerEncoder(config)
-        #self.encoder = LinformerTransformerEncoder(config)
+        self.roberta = JargonModel(config, add_pooling_layer=False)
         self.sbo_head = self.build_sbo_head(config)
 
     def build_sbo_head(self, config):
@@ -76,19 +75,16 @@ class Flaubert2ModelForTokenClassification(RobertaForTokenClassification):
         )
         
 
-class Flaubert2ModelForMaskedLM(RobertaForMaskedLM):
+class JargonForMaskedLM(RobertaForMaskedLM):
 
-    config_class = Flaubert2Config
+    config_class = JargonConfig
     
     def __init__(self, config,  **kwargs):
-        base_model_prefix = "flaubert2"
+        base_model_prefix = "jargon"
         
         super().__init__(config, **kwargs)
 
-        #self.encoder = Flaubert2Model(config, add_pooling_layer=False)
-        self.roberta = Flaubert2Model(config, add_pooling_layer=False)
-        #self.encoder = LinformerTransformerEncoder(config)
-        #self.encoder = LinformerTransformerEncoder(config)
+        self.roberta = JargonModel(config, add_pooling_layer=False)
         self.sbo_head = self.build_sbo_head(config)
 
     def build_sbo_head(self, config):
@@ -101,20 +97,16 @@ class Flaubert2ModelForMaskedLM(RobertaForMaskedLM):
             )
         )
 
-class Flaubert2ModelForMaskedLMFS(RobertaForMaskedLM):
+
+class JargonForMaskedLMFS(RobertaForMaskedLM):
 
     def __init__(self, config, dictionary, **kwargs):
-        config_class = Flaubert2Config
-        base_model_prefix = "flaubert2"
+        config_class = JargonConfig
+        base_model_prefix = "jargon"
         
         super().__init__(config, **kwargs)
 
-        #self.encoder = Flaubert2Model(config, add_pooling_layer=False)
-        #self.roberta = Flaubert2ModelFS(config, dictionary, add_pooling_layer=False)
-        self.roberta =FlaubertEncoder(config, dictionary)
-        #self.encoder = 
-        #self.encoder = LinformerTransformerEncoder(config)
-        #self.sbo_head = self.build_sbo_head(config)
+        self.roberta = FlaubertEncoder(config, dictionary)
 
     def build_sbo_head(self, config):
         return SBOHead(
@@ -127,12 +119,11 @@ class Flaubert2ModelForMaskedLMFS(RobertaForMaskedLM):
         )
 
 
-
-class Flaubert2Embeddings(RobertaEmbeddings):
+class JargonEmbeddings(RobertaEmbeddings):
 
     def __init__(self, config, **kwargs):
-        config_class = Flaubert2Config
-        base_model_prefix = "flaubert2"
+        config_class = JargonConfig
+        base_model_prefix = "jargon"
         super().__init__(config, **kwargs)
     
     def forward(
@@ -168,16 +159,14 @@ class Flaubert2Embeddings(RobertaEmbeddings):
         token_type_embeddings = self.token_type_embeddings(token_type_ids)
         
         embeddings = inputs_embeds + token_type_embeddings
-        #if self.position_embedding_type == "absolute":
         position_embeddings = self.position_embeddings(position_ids)
-        #else:
 
         embeddings += position_embeddings
-        #embeddings = self.LayerNorm(embeddings)
         embeddings = self.dropout(embeddings)
         return embeddings
 
-class Flaubert2Encoder(RobertaEncoder):
+
+class JargonEncoder(RobertaEncoder):
 
     def __init__(self, args):
         compress_layer = None
@@ -258,15 +247,15 @@ class Flaubert2Encoder(RobertaEncoder):
             return lm_out
 
     
-class Flaubert2Model(RobertaModel):
-    config_class = Flaubert2Config
+class JargonModel(RobertaModel):
+    config_class = JargonConfig
     def __init__(self, config, **kwargs):
-        
-        base_model_prefix = "flaubert2"
+        config_class = JargonConfig
+        base_model_prefix = "jargon"
         
         super().__init__(config, **kwargs)
-        self.embeddings = Flaubert2Embeddings(config)
-        self.encoder = Flaubert2Encoder(config)
+        self.embeddings = JargonEmbeddings(config)
+        self.encoder = JargonEncoder(config)
     # Copied from modeling_roberta.py
     # Add transpose of embeddings as implemented in fairseq
     def forward(
@@ -406,6 +395,7 @@ class Flaubert2Model(RobertaModel):
             cross_attentions=encoder_outputs.cross_attentions,
         )
 
+
 class SBOLayer(nn.Module):
 
     def __init__(self, input_size, hidden_size, activation, export):
@@ -417,6 +407,7 @@ class SBOLayer(nn.Module):
     def forward(self, x):
         return self.norm(self.activ(self.layer(x)))
 
+
 class SBONetwork(nn.Module):
 
     def __init__(self, input_size, hidden_size, activation, export):
@@ -514,6 +505,7 @@ def get_activation_fn(activation):
     else:
         raise RuntimeError("--activation-fn {} not supported".format(activation))
 
+
 def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
     """
     Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols