models/prediction.py

class Prediction:
    def __init__(self, label: str, score: float):
        self.label = label
        self.score = score

    @staticmethod
    def from_output(dictionary: dict):
        return Prediction(Prediction.get_label(dictionary["label"]), dictionary["score"])

    @staticmethod
    def get_label(label: str):
        return int(label.split('_')[1])