There's a simpler way of doing this, I hope

qarac/models/QaracDecoderModel.py

@@ -92,9 +92,11 @@ class QaracDecoderModel(transformers.RobertaModel,
         None.
 
         """
-        super(QaracDecoderModel,self).from_pretrained(model_path,config=config)
-        self.decoder_head = QaracDecoderHead(self.base_model.config,
-                                             self.base_model.roberta.get_input_embeddings())
+        super(QaracDecoderModel,self).__init__(config)
+        self.decoder_base = transformers.RobertaModel.from_pretrained(model_path,
+                                                                      config=config)
+        self.decoder_head = QaracDecoderHead(self.config,
+                                             self.decoder_base.roberta.get_input_embeddings())
         self.tokenizer = tokenizer
 
         
@@ -119,7 +121,7 @@ class QaracDecoderModel(transformers.RobertaModel,
         (v,s) = (kwargs['vector'],inputs) if 'vector' in kwargs else inputs
         
         return self.decoder_head(torch.unsqueeze(v,1),
-                                  self.base_model(s).last_hidden_state,
+                                  self.decoder_base(s).last_hidden_state,
                                  training = kwargs.get('training',False))
     
     def prepare_inputs_for_generation(self, 

qarac/models/QaracEncoderModel.py

@@ -9,9 +9,9 @@ Created on Tue Sep  5 10:01:39 2023
 import transformers
 import qarac.models.layers.GlobalAttentionPoolingHead
 
-class QaracEncoderModel(transformers.RobertaModel):
+class QaracEncoderModel(transformers.PreTrainedModel):
     
-    def __init__(self,base_model):
+    def __init__(self,path):
         """
         Creates the endocer model
 
@@ -25,8 +25,9 @@ class QaracEncoderModel(transformers.RobertaModel):
         None.
 
         """
-        config = transformers.PretrainedConfig.from_pretrained(base_model)
-        super(QaracEncoderModel,self).from_pretrained(base_model,config=config)
+        config = transformers.PretrainedConfig.from_pretrained(path)
+        super(QaracEncoderModel,self).__init__(config)
+        self.encoder = transformers.RobertaModel(path)
         self.head = qarac.models.layers.GlobalAttentionPoolingHead.GlobalAttentionPoolingHead(config)
         
         
@@ -47,8 +48,8 @@ class QaracEncoderModel(transformers.RobertaModel):
 
         """
 
-        return self.head(self.base_model(input_ids,
-                                         attention_mask).last_hidden_state,
+        return self.head(self.encoder(input_ids,
+                                      attention_mask).last_hidden_state,
                          attention_mask)
     
     @property