More work on models

qarac/models/QaracDecoderModel.py

@@ -7,31 +7,111 @@ Created on Tue Sep  5 10:29:03 2023
 """
 
 import keras
+import tensorflow
 import transformers
 
 class QaracDecoderHead(keras.layers.Layer):
     
     def __init__(self,config):
+        """
+        Creates the Decoder head
+
+        Parameters
+        ----------
+        config : transformers.RobertaConfig
+            Config for the RobertaModel that this head will be attached to.
+
+        Returns
+        -------
+        None.
+
+        """
         super(QaracDecoderHead,self).__init__()
         self.concat = keras.layers.Concatenate(axis=1)
         self.layer_0 = transformers.TFRobertaLayer(config)
         self.layer_1 = transformers.TFRobertalayer(config)
         self.head = transformers.TFRobertaLMHead(config)
         
+    def build(self,input_shape):
+        """
+        
+
+        Parameters
+        ----------
+        input_shape : tuple
+            Input shape.
+
+        Returns
+        -------
+        None.
+
+        """
+        self.built = True
+        
     def call(self,inputs):
+        """
+        Predicts text fron vector and hidden states of base model
+
+        Parameters
+        ----------
+        inputs : tuple of tensorflow.Tensors
+            Vector to be decoded and last hidden states of base model
+
+        Returns
+        -------
+        transformers.modeling_tf_outputs.TFCausalLMOutputWithCrossAttentions
+            Predicted text
+
+        """
         vectors = self.concat(inputs)
         l0 = self.layer_0(vectors)
         return self.head(self.layer1(l0.last_hidden_state[:,1:]))
 
-class QaracDecoderModel(transformers.TFPretrainedModel):
+class QaracDecoderModel(transformers.TFPretrainedModel,transformers.TFGenerationMixin):
     
-    def __init__(self,base_model):
+    def __init__(self,base_model,tokenizer):
+        """
+        Creates decoder model from base model
+
+        Parameters
+        ----------
+        base_model : transformers.TFRobertaModel
+            The base model
+
+        Returns
+        -------
+        None.
+
+        """
         super(QaracDecoderModel,self).__init__()
         self.base_model = base_model
         self.decoder_head = QaracDecoderHead(self.base_model.config)
+        self.tokenizer = tokenizer
+        self.start=None
+        self.end=None
+        self.pad=None
         
-    def call(self,inputs):
-        (v,s) = inputs
-        return self.decoder_head((v,self.base_model(s)))
+    def call(self,inputs,**kwargs):
+        """
+        Predicts text from inputs
+
+        Parameters
+        ----------
+        inputs : tuple of Tensorflow.Tensors OR tensorflow.Tensor
+            Vector to be converted to text and seed text ORtokenized seed text
+        kwargs : optional keyword arguments
+            vector : tensorflow.Tensor vector to be decoded. May be supplied 
+                     via a keyword argument when this is invoked by .generate
+
+        Returns
+        -------
+        transformers.modeling_tf_outputs.TFCausalLMOutputWithCrossAttentions
+            Predicted text
+
+        """
+        (v,s) = (kwargs['vector'],inputs) if 'vector' in kwargs else inputs
+        return self.decoder_head((v,self.base_model(s).last_hidden_state))
+    
+    
         
     

qarac/models/QaracEncoderModel.py

@@ -7,16 +7,59 @@ Created on Tue Sep  5 10:01:39 2023
 """
 
 import transformers
-import qarac.layers.GlobalAttentionPoolingHead
+import qarac.models.layers.GlobalAttentionPoolingHead
 
 class QaracEncoderModel(transformers.TFPretrainedModel):
     
     def __init__(self,base_model):
+        """
+        Creates the endocer model
+
+        Parameters
+        ----------
+        base_model : transformers.TFRobertaModel
+            The base model
+
+        Returns
+        -------
+        None.
+
+        """
         super(QaracEncoderModel,self).__init__()
         self.base_model = base_model
-        self.head = qarac.layers.GlobalAttentionPoolingHead.GlobalAttentionPoolingHead()
+        self.head = qarac.models.layers.GlobalAttentionPoolingHead.GlobalAttentionPoolingHead()
+        
+    def build(self,input_shape):
+        """
+        
+
+        Parameters
+        ----------
+        input_shape : tuple
+            shape of input data.
+
+        Returns
+        -------
+        None.
+
+        """
+        self.built=True
         
     def call(self,inputs):
+        """
+        Vectorizes a tokenised text
+
+        Parameters
+        ----------
+        inputs : tensorflow.Tensor
+            tokenized text to endode
+
+        Returns
+        -------
+        tensorflow.Tensor
+            Vector representing the document
+
+        """
         return self.head(self.base_model(inputs).last_hidden_state)
   
     

qarac/models/QaracTrainerModel.py

@@ -12,14 +12,56 @@ import QaracDecoderModel
 
 class QuaracTrainerModel(keras.Model):
     
-    def __init__(self,base_encoder_model,base_decoder_model):
-        
+    def __init__(self,base_encoder_model,base_decoder_model,tokenizer):
+        """
+        Sets up the Trainer model
+
+        Parameters
+        ----------
+        base_encoder_model : transformers.TFRobertaModel
+            Base model for encoders.
+        base_decoder_model : transformers.TFRobertaModel
+            Base model for decoder
+        tokenizer : transformers.RobertaTokenizer
+            Tokeniaer for decoder
+        Returns
+        -------
+        None.
+
+        """
         self.question_encoder = QaracEncoderModel.QaracEncoderModel(base_encoder_model)
         self.answer_encoder = QaracEncoderModel.QaracEncoderModel(base_encoder_model)
-        self.decoder = QaracDecoderModel.QaracDecoderModel(base_decoder_model)
+        self.decoder = QaracDecoderModel.QaracDecoderModel(base_decoder_model,tokenizer)
         self.consistency = keras.layers.Dot(axes=1,normalize=True)
         
     def call(self,inputs,training=None):
+        """
+        Generates training objective outputs from training 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
+        training : Bool, optional
+            Not used. The default is None.
+
+        Returns
+        -------
+        results : TYPE
+            DESCRIPTION.
+
+        """
         results = {}
         results['encode_decode'] = self.decoder((self.answer_encoder(inputs['all_text']),
                                                 inputs['offset_text']))

qarac/models/layers/GlobalAttentionPoolingHead.py

@@ -12,37 +12,66 @@ import tensorflow
 class GlobalAttentionPoolingHead(keras.layers.Layer):
     
     def __init__(self):
+        """
+        Creates the layer
+
+        Returns
+        -------
+        None.
+
+        """
         super(GlobalAttentionPoolingHead,self).__init__()
         self.global_projection = None
         self.local_projection = None
         
         
     def build(self,input_shape):
+        """
+        Initialises layer weights
+
+        Parameters
+        ----------
+        input_shape : tuple
+            Shape of the input layer
+
+        Returns
+        -------
+        None.
+
+        """
         width = input_shape[-1]
         self.global_projection = self.add_weight('global projection',shape=(width,width))
         self.local_projection = self.add_weight('local projection',shape=(width,width))
-        self.build=True
+        self.built=True
     
-    @tensorflow.function
-    def project(self,X):
-        return tensorflow.tensordot(X,self.local_projection,axes=1)
     
-    def attention_function(self,gp):
-        @tensorflow.function
-        def inner(lp):
-            return tensorflow.tensordot(lp,gp,axes=1)
-        return inner
         
     def call(self,X,training=None):
+        """
+        
+
+        Parameters
+        ----------
+        X : tensorflow.Tensor
+            Base model vectors to apply pooling to.
+        training : bool, optional
+            Not used. The default is None.
+
+        Returns
+        -------
+        tensorflow.Tensor
+            The pooled value.
+
+        """
         gp = tensorflow.linalg.l2_normalize(tensorflow.tensordot([tensorflow.reduce_sum(X,
                                                                                        axis=1),
                                                                   self.global_projection],
                                                                  axes=1),
                                             axis=1)
-        lp = tensorflow.linalg.l2_normalize(tensorflow.ragged.map_flat_values(self.project,
-                                                                              X),
+        lp = tensorflow.linalg.l2_normalize(tensorflow.tensordot(X,
+                                                                 self.local_projection,
+                                                                 axes=1),
                                             axis=2)
-        attention = tensorflow.ragged.map_flat_values(self.attention_function(gp), 
-                                                      lp)
+        attention = tensorflow.tensordot(lp,gp,axes=1)
         return tensorflow.reduce_sum(attention *X,
                                      axis=1)

scripts.py

@@ -9,7 +9,10 @@ import qarac.models.qarac_base_model
 import keras
 import tensorflow
 
-#tensorflow.debugging.disable_traceback_filtering()
+def decoder_loss(y_true,y_pred):
+    return keras.losses.sparse_categorical_crossentropy(y_true,
+                                                        y_pred.logits,
+                                                        logits=True)