Converted QaracTrainerModel to use PyTorch

qarac/models/QaracTrainerModel.py

@@ -6,11 +6,13 @@ Created on Tue Sep  5 15:30:06 2023
 @author: peter
 """
 
-import keras
+import torch
 import qarac.models.QaracEncoderModel
 import qarac.models.QaracDecoderModel
 
-class QaracTrainerModel(keras.Model):
+EPSILON=1.0e-12
+
+class QaracTrainerModel(torch.nn.Module()):
     
     def __init__(self,base_encoder_model,base_decoder_model,tokenizer):
         """
@@ -18,9 +20,9 @@ class QaracTrainerModel(keras.Model):
 
         Parameters
         ----------
-        base_encoder_model : transformers.TFRobertaModel
+        base_encoder_model : transformers.RobertaModel
             Base model for encoders.
-        base_decoder_model : transformers.TFRobertaModel
+        base_decoder_model : transformers.RobertaModel
             Base model for decoder
         tokenizer : transformers.RobertaTokenizer
             Tokeniaer for decoder
@@ -33,54 +35,64 @@ class QaracTrainerModel(keras.Model):
         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.consistency = keras.layers.Dot(axes=1,normalize=True)
         
-    def call(self,inputs,training=None):
+    def forward(self,
+                all_text,
+                offset_text,
+                question,
+                answer,
+                proposition0,
+                proposition1,
+                conclusion_offset,
+                statement0,
+                statement1):
         """
-        Generates training objective outputs from training data
+        Generates training objectives from data
 
         Parameters
         ----------
-        inputs : dict[str,tensoflow.tensor]
-            Fields are
-            'all_text': Tokenized text to train answer encoder to produce vectors 
-                        and decoder to convert them back to text
-            'offset_text': Same text as in 'all_text', but preceded by <s>
-            'question': Tokenized text of questions for question answering 
-                        objective
-            'answer': Tokenized text of answers for question answering objective
-            'proposition0': tokenized proposition for reasoning objective
-            'proposition1': tokenized proposition for reasoning objective
-            'conclusion_offset': tokenized text of conclusions for reasoning 
-                                 objective, prefixed by '<s>'
-            'statement0': tokenized statement for consistency objective
-            'statement1: tokenized statement for consistency objective'
-        training : Bool, optional
-            Not used. The default is None.
+        all_text : torch.tensor
+            Tokenized text for encode-decode objective
+        offset_text : torch.tensor
+            As above, prefixed with <s>
+        question : torch.tensor
+            tokenized question for question ansering objective
+        answer : torch.tensor
+            tokenized answer for question answering objective
+        proposition0 : torch.tensor
+            tokenized proposition for reasoning objective.
+        proposition1 : otrch.tensor
+            tokenized proposition for reasoning objective
+        conclusion_offset : torch.tensor
+            tokeniaed conclusion for reasoning objective, prefixed with <s>
+        statement0 : torch.tensor
+            tokenized statement for consistency objective
+        statement1 : torch.tensor
+            tokenized.statement for consistency ogjective
 
         Returns
         -------
-        results : dict[str,tensorflow.tensor]
-            Fields are
-            'encode_decode': tokeniaed text from decoding of vectors produced by
-                             answer encoder from 'all_text'
-            'question_answering': difference between vector produced by question
-                                  encoder for 'question' and answer encoder for 
-                                  'answer'
-            'reasoning': tokenised text produced by decoder from sum of vectors 
-                         produced by answwr endocer for 'proposition0' and 
-                         'proposition1'
-            'consistency': cosine similarity of vectors produced by answer encoder 
-                           from 'statement0' and 'statement1'
+        encode_decode : transformers.modeling_outputs.CausalLMOutputWithCrossAttentions
+            Predicted text for encode-decode task
+        question_answering : torch.tensor
+            Difference between encoded question and encoded answeer
+        reasoning : transformers.modeling_outputs.CausalLMOutputWithCrossAttentions
+            Predicted text for reasoning objective
+        consistency : torch.tensor
+            Cosine similarity of vectorized statements
 
         """
-        results = {}
-        results['encode_decode'] = self.decoder((self.answer_encoder(inputs['all_text']),
-                                                inputs['offset_text']))
-        results['question_answering'] = self.question_encoder(inputs['question']) - self.answer_encoder(inputs['answer'])
-        results['reasoning'] = self.decoder((self.answer_encoder(inputs['proposition0'])
-                                             +self.answer_encoder(inputs['proposition1']),
-                                             inputs['conclusion_offset']))
-        results['consistency'] = self.consistency((self.answer_encoder(inputs['statement0']),
-                                                   self.answer_encoder(inputs['statement1'])))
-        return results
+        encode_decode = self.decoder((self.answer_encoder(all_text),
+                                      offset_text))
+        question_answering = self.question_encoder(question) - self.answer_encoder(answer)
+        reasoning = self.decoder((self.answer_encoder(proposition0)
+                                             +self.answer_encoder(proposition1),
+                                             conclusion_offset))
+        s0vec = self.answer_encoder(statement0)
+        s0norm = torch.max(torch.linalg.vector_norm(s0vec,dim=1),EPSILON)
+        s0 = s0vec/s0norm
+        s1vec = self.answer_encoder(statement1)
+        s1norm = torch.max(torch.linalg.vector_norm(s1vec,dim=1),EPSILON)
+        s1 = s1vec/s1norm
+        consistency = torch.einsum('ij,ij->i',s0,s1)
+        return (encode_decode,question_answering,reasoning,consistency)