ebd_reg

# -*- coding: utf-8 -*-
import tensorflow as tf
from tensorflow.keras import layers, Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import SparseCategoricalCrossentropy
from transformers import (BertTokenizer, TFBertModel,
                          RobertaTokenizer, TFRobertaModel,
                          AlbertTokenizer, TFAlbertModel,
                          DebertaTokenizer, TFDebertaModel,
                          FunnelTokenizer, TFFunnelModel)


class Transformer_EBD_Reg:
    def __init__(self):
        self.num_classes = 3
        self.shared_fc1 = layers.Dense(768, activation='tanh')  # Assuming 768 as the dimension of output embeddings
        self.shared_fc2 = layers.Dense(1, activation='sigmoid')
        self.shared_pooling = layers.GlobalMaxPool1D()
        self.shared_output_layer = layers.Dense(self.num_classes, activation="softmax")
        self.model_function = {'bert':self.load_bert,'albert':self.load_albert,'roberta':self.load_roberta,'deberta':self.load_deberta,'funnel_tf':self.load_funnel_tf}

    def _build_model(self, model, input_shapes, weight_path):
        inputs = [layers.Input(shape=shape, dtype=tf.int32, name=name)
                  for shape, name in zip(input_shapes.values(), input_shapes.keys())]
        bert_output = model(*inputs).last_hidden_state

        bert_output_transformed = self.shared_fc2(self.shared_fc1(bert_output))
        bert_output_multiplied = bert_output_transformed * bert_output

        norm = tf.norm(bert_output_multiplied, axis=1, keepdims=True)
        bert_output_multiplied_normalized = bert_output_multiplied / norm

        bert_output_pooled = self.shared_pooling(bert_output_multiplied_normalized)
        output = self.shared_output_layer(bert_output_pooled)

        new_model = Model(inputs=inputs, outputs=[output])
        new_model.load_weights(weight_path)

        loss = SparseCategoricalCrossentropy(from_logits=False)
        optimizer = Adam(learning_rate=1e-5)
        new_model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])

        return new_model

    def load_bert(self,model):
        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
        bert_model = TFBertModel.from_pretrained('bert-base-uncased')
        input_shapes = {"input_ids": (None,), "attention_mask": (None,), "token_type_ids": (None,)}
        weight_path = '{}_weights.h5'.format(model)
        return self._build_model(bert_model, input_shapes, weight_path), tokenizer

    def load_roberta(self,model):
        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
        roberta_model = TFRobertaModel.from_pretrained('roberta-base')
        input_shapes = {"input_ids": (None,), "attention_mask": (None,)}
        weight_path = '{}_weights.h5'.format(model)
        return self._build_model(roberta_model, input_shapes, weight_path), tokenizer

    def load_albert(self,model):
        tokenizer = AlbertTokenizer.from_pretrained("albert-base-v2")
        albert_model = TFAlbertModel.from_pretrained('albert-base-v2')
        input_shapes = {"input_ids": (None,), "attention_mask": (None,), "token_type_ids": (None,)}
        weight_path = '{}_weights.h5'.format(model)
        return self._build_model(albert_model, input_shapes, weight_path), tokenizer

    def load_deberta(self,model):
        tokenizer = DebertaTokenizer.from_pretrained("microsoft/deberta-base")
        deberta_model = TFDebertaModel.from_pretrained('microsoft/deberta-base')
        input_shapes = {"input_ids": (None,), "attention_mask": (None,), "token_type_ids": (None,)}
        weight_path = '{}_weights.h5'.format(model)
        return self._build_model(deberta_model, input_shapes, weight_path), tokenizer

    def load_funnel_tf(self,model):
        tokenizer = FunnelTokenizer.from_pretrained('funnel-transformer/small')
        funnel_model = TFFunnelModel.from_pretrained('funnel-transformer/small')
        input_shapes = {"input_ids": (None,), "attention_mask": (None,), "token_type_ids": (None,)}
        weight_path = '{}_weights.h5'.format(model)
        return self._build_model(funnel_model, input_shapes, weight_path), tokenizer

    def load_weights(self, model):
        return self.model_function[model](model)



if __name__ == '__main__':
    transformer_model = Transformer_EBD_Reg()
    roberta_model, roberta_tokenizer = transformer_model.load_weights('roberta')
    roberta_model.summary()