# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_metrics.ipynb (unless otherwise specified).

__all__ = ['MetricCallback', 'RougeScore', 'Seqeval']

# Cell
import torch
from fastcore.all import *
from fastai.basics import Transform, ItemTransform, Callback, ValueMetric
from functools import partial
from typing import List, Tuple
from datasets import load_metric
try:
    import nltk
except: pass
from .learner import TransCallback

# Cell
class MetricCallback(Callback):
    """
    Computes metric scores using HuggingFace datasets metric
    """
    order = TransCallback.order + 2
    run_train, run_valid = False, True
    score_xtra_keys = {}

    def before_fit(self): self.add_metrics()

    def after_loss(self):
        if self.predict_with_generate:
            preds = self.generated_tokens
        else:
            preds = self.pred.argmax(-1)
        labels = self.yb[0] if len(self.yb) else self.trans.labels[0]

        preproc_preds, preproc_labels = self.preprocess(preds, labels)
        self.metric.add_batch(predictions=preproc_preds, references=preproc_labels)

    def after_validate(self):
        self.res = self.metric.compute()

    def preprocess(self, predictions, labels):
        return predictions, labels

    @staticmethod
    def _get_score(obj, score, **kwargs):
        return obj.res[score]

    def _register_value_funcs(self):
        for score in self.scores:
            func = partial(self._get_score, self, score, **self.score_xtra_keys)
            setattr(self, score, func)

    def add_metrics(self):
        learn_metrics = set([m.name for m in self.metrics])
        for score_name in self.scores:
            if not (score_name in learn_metrics):
                self.learn.metrics += [ValueMetric(getattr(self, score_name), score_name)]

# Cell
def _get_score(obj, score_name, agg, measure):
    return getattr(getattr(obj.res[score_name], agg), measure)

# Cell
class RougeScore(Callback):
    """
    Computes ROUGE score using HF datasets metric. Adds `ValueMetric`
    for each score in `scores`
    """
    order = TransCallback.order + 2
    run_train, run_valid = False, True
    def __init__(self, tokenizer, argmax:bool=True, agg:str='mid', measure:str='fmeasure',
                 scores:List[str]=['rouge1', 'rouge2', 'rougeL', 'rougeLsum']):
        nltk.download('punkt', quiet=True)
        self.metric = load_metric("rouge")
        store_attr()
        self._register_value_funcs()

    def before_fit(self): self.add_metrics()

    def after_loss(self):
        if self.predict_with_generate:
            preds = self.generated_tokens
        else:
            preds = self.pred.argmax(-1)
        labels = self.yb[0] if len(self.yb) else self.trans.labels[0]
        # decode preds and labels
        decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True)
        # Replace -100 in the labels as we can't decode them.
        labels = torch.where(labels != -100, labels, self.tokenizer.pad_token_id)
        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True)
        # Rouge expects a newline after each sentence
        decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]
        decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]
        self.metric.add_batch(predictions=decoded_preds, references=decoded_labels)

    def after_validate(self):
        self.res = self.metric.compute()

    def _register_value_funcs(self):
        for score in self.scores:
            func = partial(_get_score, self, score, self.agg, self.measure)
            setattr(self, score, func)

    def add_metrics(self):
        learn_metrics = set([m.name for m in self.metrics])
        for score_name in self.scores:
            if not (score_name in learn_metrics):
                self.learn.metrics += [ValueMetric(getattr(self, score_name), score_name)]


# Cell
class Seqeval(MetricCallback):
    "Computes scores using `seqeval` https://github.com/chakki-works/seqeval"
    def __init__(self, label_list, scores:Tuple[str]=('accuracy', 'f1', 'precision', 'recall')):
        self.metric = load_metric('seqeval')
        store_attr()
        self._register_value_funcs()

    @staticmethod
    def _get_score(obj, score, **kwargs):
        return obj.res[f"overall_{score}"]

    def preprocess(self, predictions, labels):
        "Preprocess predictions and labels for seqeval, ref: https://github.com/huggingface/notebooks/blob/master/examples/token_classification.ipynb"
        # Remove ignored index (special tokens)
        true_predictions = [
            [self.label_list[p] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]
        true_labels = [
            [self.label_list[l] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]
        return true_predictions, true_labels