from model.model import Model


class EmotionAnalysis(Model):
    def __init__(self) -> None:
        self.model_name = "thoriqfy/indobert-emotion-classification"
        self.tasks = "emotion"
        self.load_model(model_name=self.model_name, tasks=self.tasks)
        
    def predict(self, sentences):
        outputs = super().predict(sentences, self.tasks)
        
        return {
            "result": outputs["label"], 
            "score": outputs["score"]
        }
    
emotion = EmotionAnalysis()