from anonymous_demo import TADCheckpointManager

from textattack.model_args import DEMO_MODELS
from textattack.reactive_defense.reactive_defender import ReactiveDefender


class TADReactiveDefender(ReactiveDefender):
    """ Transformers sentiment analysis pipeline returns a list of responses
        like

            [{'label': 'POSITIVE', 'score': 0.7817379832267761}]

        We need to convert that to a format TextAttack understands, like

            [[0.218262017, 0.7817379832267761]
    """

    def __init__(self, ckpt='tad-sst2', **kwargs):
        super().__init__(**kwargs)
        self.tad_classifier = TADCheckpointManager.get_tad_text_classifier(checkpoint=DEMO_MODELS[ckpt],
                                                                           auto_device=True)

    def reactive_defense(self, text, **kwargs):
        res = self.tad_classifier.infer(text, defense='pwws', print_result=False, **kwargs)
        return res