PyTorch implementation of HugggingFace PreTrainedModel class does not allow direct setting of base_model. Rejig constructors accordingly

qarac/models/QaracDecoderModel.py

@@ -11,7 +11,7 @@ import transformers
 
 class QaracDecoderHead(torch.nn.Module):
     
-    def __init__(self,config,input_embeddings):
+    def __init__(self,base_model,config,input_embeddings):
         """
         Creates the Decoder head
 
@@ -25,7 +25,7 @@ class QaracDecoderHead(torch.nn.Module):
         None.
 
         """
-        super(QaracDecoderHead,self).__init__()
+        super(QaracDecoderHead,self).from_pretrained(base_model,config)
         self.layer_0 = transformers.models.roberta.modeling_roberta.RobertaLayer(config)
         self.layer_1 = transformers.models.roberta.modeling_roberta.RobertaLayer(config)
         self.head = transformers.models.roberta.modeling_roberta.RobertaLMHead(config,
@@ -77,7 +77,7 @@ class QaracDecoderHead(torch.nn.Module):
 
 class QaracDecoderModel(transformers.PreTrainedModel,transformers.generation_utils.GenerationMixin):
     
-    def __init__(self,base_model,tokenizer):
+    def __init__(self,config,tokenizer):
         """
         Creates decoder model from base model
 
@@ -91,7 +91,7 @@ class QaracDecoderModel(transformers.PreTrainedModel,transformers.generation_uti
         None.
 
         """
-        super(QaracDecoderModel,self).__init__(base_model.config)
+        super(QaracDecoderModel,self).__init__(config)
         self.base_model = base_model
         self.decoder_head = QaracDecoderHead(self.base_model.config,
                                              self.base_model.roberta.get_input_embeddings())

qarac/models/QaracEncoderModel.py

@@ -25,9 +25,8 @@ class QaracEncoderModel(transformers.PreTrainedModel):
         None.
 
         """
-        super(QaracEncoderModel,self).__init__(base_model.config)
-        self.base_model = base_model
-        self.head = qarac.models.layers.GlobalAttentionPoolingHead.GlobalAttentionPoolingHead(base_model.config)
+        super(QaracEncoderModel,self).from_pretrained(base_model)
+        self.head = qarac.models.layers.GlobalAttentionPoolingHead.GlobalAttentionPoolingHead(self.base_model.config)
         
         
     def forward(self,input_ids,
@@ -50,6 +49,10 @@ class QaracEncoderModel(transformers.PreTrainedModel):
         return self.head(self.base_model(input_ids,
                                          attention_mask).last_hidden_state,
                          attention_mask)
+    
+    @property
+    def config(self):
+        return self.base_model.config
   
     
     

qarac/models/QaracTrainerModel.py

@@ -14,7 +14,7 @@ EPSILON=1.0e-12
 
 class QaracTrainerModel(torch.nn.Module):
     
-    def __init__(self,base_encoder_model,base_decoder_model,tokenizer):
+    def __init__(self,base_model_path,tokenizer):
         """
         Sets up the Trainer model
 
@@ -32,9 +32,13 @@ class QaracTrainerModel(torch.nn.Module):
 
         """
         super(QaracTrainerModel,self).__init__()
-        self.question_encoder = qarac.models.QaracEncoderModel.QaracEncoderModel(base_encoder_model)
-        self.answer_encoder = qarac.models.QaracEncoderModel.QaracEncoderModel(base_encoder_model)
-        self.decoder = qarac.models.QaracDecoderModel.QaracDecoderModel(base_decoder_model,tokenizer)
+        self.question_encoder = qarac.models.QaracEncoderModel.QaracEncoderModel(base_model_path)
+        self.answer_encoder = qarac.models.QaracEncoderModel.QaracEncoderModel(base_model_path)
+        config = self.answer_encoder.config
+        config.is_decoder = True
+        self.decoder = qarac.models.QaracDecoderModel.QaracDecoderModel(base_model_path,
+                                                                        config,
+                                                                        tokenizer)
         
     def forward(self,
                 all_text,

scripts.py

@@ -123,14 +123,8 @@ def prepare_training_datasets():
     consistency.to_csv('corpora/consistency.csv')
     
 def train_models(path):
-    encoder_base = transformers.TFRobertaModel.from_pretrained('roberta-base')
-    config = encoder_base.config
-    config.is_decoder = True
-    decoder_base = transformers.TFRobertaModel.from_pretrained('roberta-base',
-                                                               config=config)
-    tokenizer = tokenizers.Tokenizer.from_pretrained('roberta-base')
-    trainer = qarac.models.QaracTrainerModel.QaracTrainerModel(encoder_base, 
-                                                               decoder_base, 
+    tokenizer = tokenizers.from_pretrained('roberta-base')
+    trainer = qarac.models.QaracTrainerModel.QaracTrainerModel('roberta_base', 
                                                                tokenizer)
     loss_fn = CombinedLoss()
     optimizer = torch.optim.NAdam(trainer.parameters(),lr=5.0e-5)