import tensorflow as tf

from tensorflow.keras.models import load_model
from  transformers import TFBertModel

class MultiTaskModel:
    """
    A class used to represent a MultiTaskModel,  which classifies the emotion and toxicity of Valorant chat messages
    """
    def __init__(self, is_multilingual = False, preprocessor = None):
        if is_multilingual:
            pre_model = TFBertModel.from_pretrained('bert-base-multilingual-cased')
            self.model = load_model('model_with_bert_multilingual.h5', custom_objects={'TFBertModel': TFBertModel})
        else:
            pre_model = TFBertModel.from_pretrained('bert-base-uncased')
            self.model = load_model('model_with_bert_base.h5', custom_objects={'TFBertModel': pre_model})
        self.load_preprocess(preprocessor)

    def load_preprocess(self, prep):
        self.preprocessor = prep

    def predict(self, text):
        preptext= self.preprocessor.preprocess_text(text)
        return self.model.predict(preptext)
    
    def decode(self, pred):
        return self.preprocessor.decoder.toxicity(pred[1]), self.preprocessor.decoder.emotion(pred[0])