Making sure RoBERTa layers have all required arguments

qarac/corpora/CombinedCorpus.py

@@ -185,7 +185,11 @@ class CombinedCorpus(keras.utils.Sequence):
         maxlen = max((len(sample) for sample in batch))
         for sample in batch:
             sample.pad(maxlen,pad_id=self.pad_token)
-        return tensorflow.constant([sample.ids
-                                    for sample in batch])
+        input_ids = tensorflow.constant([sample.ids
+                                         for sample in batch])
+        attention_mask = tensorflow.constant(input_ids.numpy().apply(lambda x: 0.0 if x==self.pad_token
+                                                                     else 1.0))
+        return {'input_ids':input_ids,
+                'attention_mask':attention_mask}
     
     

qarac/models/QaracDecoderModel.py

@@ -49,7 +49,13 @@ class QaracDecoderHead(keras.layers.Layer):
         """
         self.built = True
         
-    def call(self,inputs):
+        
+        
+        
+    def call(self,
+             vector,
+             hidden_states,
+             attention_mask=None,training=False):
         """
         Predicts text fron vector and hidden states of base model
 
@@ -64,9 +70,20 @@ class QaracDecoderHead(keras.layers.Layer):
             Predicted text
 
         """
-        vectors = self.concat(inputs)
-        l0 = self.layer_0(vectors)
-        return self.head(self.layer1(l0.last_hidden_state[:,1:]))
+        vectors = self.concat(vector, hidden_states)
+        attentions = attention_mask if attention_mask is None else self.concat(tensorflow.ones((hidden_states.shape(0),
+                                                                                                1)),
+                                                                               attention_mask)
+        l0 = self.layer_0(vectors,
+                          attentions,
+                          None,
+                          False,
+                          training)
+        return self.head(self.layer1(l0.last_hidden_state[:,1:],
+                                     attention_mask,
+                                     None,
+                                     False,
+                                     training))
 
 class QaracDecoderModel(transformers.TFPreTrainedModel,transformers.generation_tf_utils.TFGenerationMixin):
     
@@ -114,7 +131,8 @@ class QaracDecoderModel(transformers.TFPreTrainedModel,transformers.generation_t
         (v,s) = (kwargs['vector'],inputs) if 'vector' in kwargs else inputs
         
         return self.decoder_head((tensorflow.expand_dims(v,1),
-                                  self.base_model(s).last_hidden_state))
+                                  self.base_model(s)),
+                                 training = kwargs.get('training',False))
     
     
         

qarac/models/QaracEncoderModel.py

@@ -45,7 +45,9 @@ class QaracEncoderModel(transformers.TFPreTrainedModel):
         """
         self.built=True
         
-    def call(self,inputs):
+    def call(self,input_ids,
+             attention_mask=None,
+             training=False):
         """
         Vectorizes a tokenised text
 
@@ -61,7 +63,11 @@ class QaracEncoderModel(transformers.TFPreTrainedModel):
 
         """
 
-        return self.head(self.base_model(inputs).last_hidden_state)
+        return self.head(self.base_model(input_ids,
+                                         attention_mask,
+                                         training=training).last_hidden_state,
+                         attention_mask,
+                         training)
   
     
     

qarac/models/layers/GlobalAttentionPoolingHead.py

@@ -57,7 +57,7 @@ class GlobalAttentionPoolingHead(keras.layers.Layer):
                                     self.local_projection,
                                     axes=1)
         
-    def call(self,X,training=None):
+    def call(self,X,attention_mask=None,training=None):
         """
         
 
@@ -65,6 +65,8 @@ class GlobalAttentionPoolingHead(keras.layers.Layer):
         ----------
         X : tensorflow.Tensor
             Base model vectors to apply pooling to.
+        attention_mask: tensorflow.Tensor, optional
+            mask for pad values
         training : bool, optional
             Not used. The default is None.
 
@@ -83,5 +85,7 @@ class GlobalAttentionPoolingHead(keras.layers.Layer):
                                                                       X),
                                             axis=2)
         attention = tensorflow.vectorized_map(dot_prod,(lp,gp))
+        if attention_mask is None:
+            attention_mask = tensorflow.ones_like(attention)
         return tensorflow.vectorized_map(dot_prod,
-                                         (attention,X))
+                                         (attention * attention_mask,X))