from typing import Dict

from service.AbstractTextClassificationService import AbstractTextClassificationService
from transformers import pipeline


class ZeroShotTextClassificationService(AbstractTextClassificationService):

    def __init__(self, modelid):
        self.classifier = pipeline("zero-shot-classification", model=modelid)
        self.candidate_labels = ["Information Technology", "Economics", "Healthcare", "History", "Education"]

    def classify(self, text: str) -> Dict[str, float]:
        result = self.classifier(text, self.candidate_labels)
        classification_results: Dict[str, float] = dict(zip(result["labels"], result["scores"]))
        return classification_results


# if __name__ == '__main__':
#     # let's test it
#     text = "Minimax (sometimes Minmax, MM[1] or saddle point[2]) is a decision rule used in artificial intelligence, decision theory, game theory, statistics, and philosophy for minimizing the possible loss for a worst case (maximum loss) scenario. When dealing with gains, it is referred to as maximin – to maximize the minimum gain. Originally formulated for several-player zero-sum game theory, covering both the cases where players take alternate moves and those where they make simultaneous moves, it has also been extended to more complex games and to general decision-making in the presence of uncertainty."
#     classifier = ZeroShotTextClassificationService("roberta-large-mnli")
#     result = classifier.classify(text)
#     print(result)







