Encoder, Decoder and Trainer models (assuming RoBERTa base models)

qarac/models/QaracDecoderModel.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Sep  5 10:29:03 2023
+
+@author: peter
+"""
+
+import keras
+import transformers
+
+class QaracDecoderHead(keras.layers.Layer):
+    
+    def __init__(self,config):
+        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 call(self,inputs):
+        vectors = self.concat(inputs)
+        l0 = self.layer_0(vectors)
+        return self.head(self.layer1(l0.last_hidden_state[:,1:]))
+
+class QaracDecoderModel(transformers.TFPretrainedModel):
+    
+    def __init__(self,base_model):
+        super(QaracDecoderModel,self).__init__()
+        self.base_model = base_model
+        self.decoder_head = QaracDecoderHead(self.base_model.config)
+        
+    def call(self,inputs):
+        (v,s) = inputs
+        return self.decoder_head((v,self.base_model(s)))
+        
+    

qarac/models/QaracEncoderModel.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Sep  5 10:01:39 2023
+
+@author: peter
+"""
+
+import transformers
+import qarac.layers.GlobalAttentionPoolingHead
+
+class QaracEncoderModel(transformers.TFPretrainedModel):
+    
+    def __init__(self,base_model):
+        super(QaracEncoderModel,self).__init__()
+        self.base_model = base_model
+        self.head = qarac.layers.GlobalAttentionPoolingHead.GlobalAttentionPoolingHead()
+        
+    def call(self,inputs):
+        return self.head(self.base_model(inputs).last_hidden_state)
+  
+    
+    

qarac/models/QaracTrainerModel.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Sep  5 15:30:06 2023
+
+@author: peter
+"""
+
+import keras
+import QaracEncoderModel
+import QaracDecoderModel
+
+class QuaracTrainerModel(keras.Model):
+    
+    def __init__(self,base_encoder_model,base_decoder_model):
+        
+        self.question_encoder = QaracEncoderModel.QaracEncoderModel(base_encoder_model)
+        self.answer_encoder = QaracEncoderModel.QaracEncoderModel(base_encoder_model)
+        self.decoder = QaracDecoderModel.QaracDecoderModel(base_decoder_model)
+        self.consistency = keras.layers.Dot(axes=1,normalize=True)
+        
+    def call(self,inputs,training=None):
+        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']),
+                                             self.inputs['conclusion_offset']))
+        results['consistency'] = self.consistency((self.answer_encoder(inputs['statement0']),
+                                                   self.answer_encoder(inputs['statement1'])))
+        return results