# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/17_training.token_classification.ipynb (unless otherwise specified).

__all__ = ['NERMetric', 'encode_tags', 'TokenClassificationDatasets', 'SeqEvalMetrics', 'TokenClassificationTuner']

# Cell
import pandas as pd
import numpy as np
from fastcore.meta import delegates
from fastcore.xtras import Path, range_of

from fastai.basics import *
from fastai.data.core import DataLoaders
from fastai.data.transforms import get_files

from .core import *
from .arrow_utils import TextNoNewLineDatasetReader
from ..inference.token_classification import TransformersTokenTagger, TokenClassificationResult, DetailLevel

from transformers import default_data_collator, AutoModelForTokenClassification, AutoTokenizer
from datasets import Dataset, load_metric

from typing import List
from ast import literal_eval

# Cell
#nbdev_comment _all_ = ['NERMetric']

# Cell
def encode_tags(tags, encodings):
    # create an empty array of -100
    encoded_labels = np.ones(len(encodings['offset_mapping']),dtype=int) * -100
    offset_array = np.array(encodings['offset_mapping'])

    # set labels whose first offset position is 0 and the second is not 0
    encoded_labels[(offset_array[:,0] == 0) & (offset_array[:,1] != 0)] = tags

    return list(encoded_labels)

# Cell
def _batch_tokenize(batch, tokenizer, tokenize_kwargs):
    updated_tags = [] # list of lists
    updated_batch = {} # dict of lists

    for tokens, tags in zip(batch['tokens'], batch['ner_tags']):
        encodings = tokenizer(tokens, **tokenize_kwargs)
        updated_tags.append(encode_tags(tags, encodings))

        for feat in encodings.keys():
            if (feat != 'offset_mapping'):
                if feat in updated_batch:
                    updated_batch[feat].append(encodings[feat])
                else:
                    updated_batch[feat] = [encodings[feat]]

    updated_batch['labels'] = updated_tags
    return updated_batch


# Cell
class TokenClassificationDatasets(TaskDatasets):
    """
    A set of datasets designed for token classification
    """
    def __init__(
        self,
        train_dset:Dataset, # A training dataset
        valid_dset:Dataset, # A validation dataset
        tokenizer_name:str, # The name of a tokenizer
        tokenize:bool, # Whether to tokenize immediately
        tokenize_kwargs:dict, # kwargs for the tokenize function
        auto_kwargs:dict, # AutoTokenizer.from_pretrained kwargs
        remove_columns:list, # The columns to remove when tokenizing
        entity_mapping:dict, # A mapping of entity names to encoded labels
    ):

        "Constructs TaskDatasets, should not be called explicitly"
        super().__init__(
            train_dset,
            valid_dset,
            tokenizer_name,
            tokenize,
            _batch_tokenize,
            tokenize_kwargs,
            auto_kwargs,
            remove_columns,
        )
        self.entity_mapping = entity_mapping

    @classmethod
    def from_dfs(
        cls,
        train_df:pd.DataFrame, # A training dataframe
        token_col:str, # The name of the token column
        tag_col:str, # The name of the tag column
        entity_mapping:dict, # A mapping of entity names to encoded labels
        tokenizer_name:str, # The name of the tokenizer
        tokenize:bool=True, # Whether to tokenize immediately
        valid_df=None, # An optional validation dataframe
        split_func=None, # Optionally a splitting function similar to RandomSplitter
        split_pct=.2, # What % to split the train_df
        tokenize_kwargs:dict={}, # kwargs for the tokenize function
        auto_kwargs:dict={} # kwargs for the AutoTokenizer.from_pretrained constructor
    ):
        "Builds `TokenClassificationDatasets` from a `DataFrame` or set of `DataFrames`"
        if split_func is None: split_func = RandomSplitter(split_pct)
        if valid_df is None:
            train_idxs, valid_idxs = split_func(train_df)
            valid_df = train_df.iloc[valid_idxs]
            train_df = train_df.iloc[train_idxs]

        train_df = train_df[[token_col,tag_col]]
        valid_df = valid_df[[token_col,tag_col]]
        train_df = train_df.rename(columns={token_col:'tokens', tag_col: 'ner_tags'})
        valid_df = valid_df.rename(columns={token_col:'tokens', tag_col: 'ner_tags'})

        train_dset = Dataset.from_dict(train_df.to_dict('list'))
        valid_dset = Dataset.from_dict(valid_df.to_dict('list'))

        return cls(
            train_dset,
            valid_dset,
            tokenizer_name,
            tokenize,
            tokenize_kwargs,
            auto_kwargs,
            remove_columns=['tokens', 'ner_tags'],
            entity_mapping = entity_mapping
        )

    @classmethod
    def from_csvs(
        cls,
        train_csv:Path, # A training csv file
        token_col:str, # The name of the token column
        tag_col:str, # The name of the tag column
        entity_mapping:dict, # A mapping of entity names to encoded labels
        tokenizer_name:str, # The name of the tokenizer
        tokenize:bool=True, # Whether to tokenize immediately
        valid_csv:Path=None, # An optional validation csv
        split_func=None, # Optionally a splitting function similar to RandomSplitter
        split_pct=.2, # What % to split the train df
        tokenize_kwargs:dict={}, # kwargs for the tokenize function
        auto_kwargs:dict={}, # kwargs for the AutoTokenizer.from_pretrained constructor
        **kwargs, # kwargs for `pd.read_csv`
    ):
        "Builds `SequenceClassificationDatasets` from a single csv or set of csvs. A convience constructor for `from_dfs`"
        train_df = pd.read_csv(train_csv, **kwargs)
        if valid_csv is not None: valid_df = pd.read_csv(valid_csv, **kwargs)
        else: valid_df = None
        return cls.from_dfs(
            train_df,
            token_col,
            tag_col,
            entity_mapping,
            tokenizer_name,
            tokenize,
            valid_df,
            split_func,
            split_pct,
            tokenize_kwargs,
            auto_kwargs
        )


    @delegates(DataLoaders)
    def dataloaders(
        self,
        batch_size:int=16, # A batch size
        shuffle_train:bool=True, # Whether to shuffle the training dataset
        collate_fn:callable=None, # A custom collation function
        **kwargs # Torch DataLoader kwargs
    ):
        "Build DataLoaders from `self`"
        dls = super().dataloaders(batch_size, shuffle_train, collate_fn, **kwargs)
        dls.entity_mapping = self.entity_mapping
        return dls

# Cell
class SeqEvalMetrics():
    """
    Multi-label classification metrics for NER, using seqeval metric from HuggingFace
    """
    def __init__(
        self,
        entity_mapping:dict, # A mapping of entity names to encoded labels
    ):
        self.metric = load_metric("seqeval")
        self.label_list = list(entity_mapping.values())

    # source: https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/token_classification.ipynb#scrollTo=-1bS-ATSqPq2
    def compute_metrics(
        self,
        preds, # Predictions output by the model
        labels # Ground truth target values
    ):
        "Computes multi-label classification metrics for NER"
        predictions = preds.argmax(2)

        # 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)
        ]

        results = self.metric.compute(predictions=true_predictions, references=true_labels)
        return results

    def accuracy(
        self,
        preds, # Predictions output by the model
        labels # Ground truth target values
    ):
        "Computes multi-label accuracy for NER"
        results = self.compute_metrics(preds, labels)
        return results['overall_accuracy']

    def precision(
        self,
        preds, # Predictions output by the model
        labels # Ground truth target values
    ):
        "Computes multi-label precision for NER"
        results = self.compute_metrics(preds, labels)
        return results['overall_precision']

    def recall(
        self,
        preds, # Predictions output by the model
        labels # Ground truth target values
    ):
        "Computes multi-label recall for NER"
        results = self.compute_metrics(preds, labels)
        return results['overall_recall']

    def f1(
        self,
        preds, # Predictions output by the model
        labels # Ground truth target values
    ):
        "Computes multi-label F1 for NER"
        results = self.compute_metrics(preds, labels)
        return results['overall_f1']

# Cell
mk_class('NERMetric', **{prop.title():prop for prop in dir(SeqEvalMetrics)
                         if not prop.startswith('_') and prop != 'compute_metrics'},
        doc="Class for all valid NER metrics usable during fine-tuning with typo-proofing")

# Cell
class TokenClassificationTuner(AdaptiveTuner):
    """
    An `AdaptiveTuner` with good defaults for Token Classification tasks

    **Valid kwargs and defaults:**
      - `lr`:float = 0.001
      - `splitter`:function = `trainable_params`
      - `cbs`:list = None
      - `path`:Path = None
      - `model_dir`:Path = 'models'
      - `wd`:float = None
      - `wd_bn_bias`:bool = False
      - `train_bn`:bool = True
      - `moms`: tuple(float) = (0.95, 0.85, 0.95)

    """
    def __init__(
        self,
        dls:DataLoaders, # A set of DataLoaders
        model_name:str, # A HuggingFace model
        tokenizer = None, # A HuggingFace tokenizer
        loss_func = CrossEntropyLossFlat(), # A loss function
        metrics:List[NERMetric] = [NERMetric.Accuracy, NERMetric.F1], # Metrics to monitor the training with
        opt_func = Adam, # A fastai or torch Optimizer
        additional_cbs = None, # Additional Callbacks to have always tied to the Tuner
        expose_fastai_api = False, # Whether to expose the fastai API
        num_classes:int=None, # The number of classes
        entity_mapping:dict=None, # A mapping of entity names to encoded labels
        **kwargs, # kwargs for `Learner.__init__`
    ):
        additional_cbs = listify(additional_cbs)
        for arg in 'dls,model,loss_func,metrics,opt_func,cbs,expose_fastai'.split(','):
            if arg in kwargs.keys(): kwargs.pop(arg) # Pop all existing kwargs
        if hasattr(dls, 'entity_mapping'): num_classes = len(dls.entity_mapping)
        else: raise ValueError('Could not extract entity mapping from DataLoaders, please pass it in as a param')
        if num_classes is None: raise ValueError('Could not extrapolate number of classes, please pass it in as a param')
        model = AutoModelForTokenClassification.from_pretrained(
            model_name,
            num_labels=num_classes,
            id2label=dls.entity_mapping
        )
        if tokenizer is None: tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.num_classes = num_classes

        self._base_metrics = SeqEvalMetrics(dls.entity_mapping)
        new_metrics = []
        for met in metrics:
            ner_met = getattr(NERMetric, met.title(), None)
            if not ner_met: raise ValueError('Metric not supported')
            else: new_metrics.append(getattr(self._base_metrics, ner_met))

        super().__init__(
            expose_fastai_api,
            dls = dls,
            model = model,
            tokenizer = tokenizer,
            loss_func = loss_func,
            metrics = new_metrics,
            opt_func = opt_func,
            cbs=additional_cbs,
            **kwargs
        )

    def predict(
        self,
        text:Union[List[str], str], # Some text or list of texts to do inference with
        bs:int=64, # A batch size to use for multiple texts
        grouped_entities: bool = True, # Return whole entity span strings
        detail_level:DetailLevel = DetailLevel.Low, # A detail level to return on the predictions
    ) -> dict: # A dictionary of filtered predictions
        "Predict some `text` for token classification with the currently loaded model"
        if getattr(self, '_inferencer', None) is None: self._inferencer = TransformersTokenTagger(self.tokenizer, self.model)

        preds = self._inferencer.predict(text, bs, grouped_entities, detail_level)

        return preds