# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_frame_labels.ipynb.

# %% auto 0
__all__ = ['FramingLabels']

# %% ../nbs/01_frame_labels.ipynb 0
import torch
from transformers import pipeline
from functools import partial
import numpy as np
import pandas as pd
import tqdm
import matplotlib.pyplot as plt
import seaborn as sns

class FramingLabels:
    def __init__(self, base_model, candidate_labels, batch_size=16):
        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        self.base_pipeline = pipeline("zero-shot-classification", model=base_model, device=device)
        self.candidate_labels = candidate_labels
        self.classifier = partial(self.base_pipeline, candidate_labels=candidate_labels, multi_label=True, batch_size=batch_size)

    def order_scores(self, dic):
        indices_order = [dic["labels"].index(l) for l in self.candidate_labels]
        scores_ordered = np.array(dic["scores"])[indices_order].tolist()
        return scores_ordered

    def get_ordered_scores(self, sequence_to_classify):
        if type(sequence_to_classify) == list:
            res = []
            for out in tqdm.tqdm(self.classifier(sequence_to_classify)):
                res.append(out)
        else:
            res = self.classifier(sequence_to_classify)
        if type(res) == list:
            scores_ordered = list(map(self.order_scores, res))
            scores_ordered = list(map(list, zip(*scores_ordered)))  # reorder
        else:
            scores_ordered = self.order_scores(res)
        return scores_ordered

    def get_label_names(self):
        label_names = [l.split(":")[0].split(" ")[0] for l in self.candidate_labels]
        return label_names

    def __call__(self, sequence_to_classify):
        scores = self.get_ordered_scores(sequence_to_classify)
        label_names = self.get_label_names()
        return dict(zip(label_names, scores))

    def visualize(self, name_to_score_dict, threshold=0.5, **kwargs):
        fig, ax = plt.subplots()

        cp = sns.color_palette()

        scores_ordered = list(name_to_score_dict.values())
        label_names = list(name_to_score_dict.keys())

        colors = [cp[0] if s > 0.5 else cp[1] for s in scores_ordered]
        ax.barh(label_names[::-1], scores_ordered[::-1], color=colors[::-1], **kwargs)
        plt.xlim(left=0)
        plt.tight_layout()
        return fig, ax
