'''simple docstring'''


import argparse
import glob
import logging
import os
from argparse import Namespace
from importlib import import_module

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from utils_ner import TokenClassificationTask


snake_case__					:      Union[str, Any]				      =  logging.getLogger(__name__)






class 					__SCREAMING_SNAKE_CASE						(	_a	):







  '''simple docstring'''






  lowerCamelCase_		:str						   =		"""token-classification"""


  def __init__(			self				,		snake_case_						):




       '''simple docstring'''

       if type(__lowerCAmelCase						) == dict:
            UpperCAmelCase_							:       Union[str, Any]       =				Namespace(**__lowerCAmelCase						)
       UpperCAmelCase_							:       List[Any]       =				import_module('tasks'						)
       try:
            UpperCAmelCase_							:       List[str]       =				getattr(__lowerCAmelCase				,		hparams.task_type						)
            UpperCAmelCase_							:       Optional[Any]       =				token_classification_task_clazz()
       except AttributeError:
            raise ValueError(
                F'''Task {hparams.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. '''
                F'''Available tasks classes are: {TokenClassificationTask.__subclasses__()}'''						)
       UpperCAmelCase_							:       Union[str, Any]       =				self.token_classification_task.get_labels(hparams.labels						)
       UpperCAmelCase_							:       Optional[int]       =				CrossEntropyLoss().ignore_index
       super().__init__(__lowerCAmelCase				,		len(self.labels						)				,		self.mode						)


  def _UpperCamelCase					(			self				,		**snake_case_						):




       '''simple docstring'''

       return self.model(**__lowerCAmelCase						)


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:       Tuple       =				{'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}
       if self.config.model_type != "distilbert":
            UpperCAmelCase_							:       int       =				(
                batch[2] if self.config.model_type in ['bert', 'xlnet'] else None
            )  # XLM and RoBERTa don"t use token_type_ids

       UpperCAmelCase_							:       Dict       =				self(**__lowerCAmelCase						)
       UpperCAmelCase_							:       Dict       =				outputs[0]
       # tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
       return {"loss": loss}


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       List[Any]       =				self.hparams
       for mode in ["train", "dev", "test"]:
            UpperCAmelCase_							:       List[str]       =				self._feature_file(__lowerCAmelCase						)
            if os.path.exists(__lowerCAmelCase						) and not args.overwrite_cache:
                 logger.info('Loading features from cached file %s'				,		__lowerCAmelCase						)
                 UpperCAmelCase_							:       Optional[int]       =				torch.load(__lowerCAmelCase						)
            else:
                 logger.info('Creating features from dataset file at %s'				,		args.data_dir						)
                 UpperCAmelCase_							:       Union[str, Any]       =				self.token_classification_task.read_examples_from_file(args.data_dir				,		__lowerCAmelCase						)
                 UpperCAmelCase_							:       Optional[Any]       =				self.token_classification_task.convert_examples_to_features(
                     __lowerCAmelCase				,		self.labels				,		args.max_seq_length				,		self.tokenizer				,		cls_token_at_end=bool(self.config.model_type in ['xlnet']						)				,		cls_token=self.tokenizer.cls_token				,		cls_token_segment_id=2 if self.config.model_type in ['xlnet'] else 0				,		sep_token=self.tokenizer.sep_token				,		sep_token_extra=__lowerCAmelCase				,		pad_on_left=bool(self.config.model_type in ['xlnet']						)				,		pad_token=self.tokenizer.pad_token_id				,		pad_token_segment_id=self.tokenizer.pad_token_type_id				,		pad_token_label_id=self.pad_token_label_id				,		)
                 logger.info('Saving features into cached file %s'				,		__lowerCAmelCase						)
                 torch.save(__lowerCAmelCase				,		__lowerCAmelCase						)


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_				,		snake_case_ = False						):




       '''simple docstring'''

       UpperCAmelCase_							:       Tuple       =				self._feature_file(__lowerCAmelCase						)
       logger.info('Loading features from cached file %s'				,		__lowerCAmelCase						)
       UpperCAmelCase_							:       Dict       =				torch.load(__lowerCAmelCase						)
       UpperCAmelCase_							:       Dict       =				torch.tensor([f.input_ids for f in features]				,		dtype=torch.long						)
       UpperCAmelCase_							:       Any       =				torch.tensor([f.attention_mask for f in features]				,		dtype=torch.long						)
       if features[0].token_type_ids is not None:
            UpperCAmelCase_							:       Union[str, Any]       =				torch.tensor([f.token_type_ids for f in features]				,		dtype=torch.long						)
       else:
            UpperCAmelCase_							:       Optional[Any]       =				torch.tensor([0 for f in features]				,		dtype=torch.long						)
            # HACK(we will not use this anymore soon)
       UpperCAmelCase_							:       int       =				torch.tensor([f.label_ids for f in features]				,		dtype=torch.long						)
       return DataLoader(
           TensorDataset(__lowerCAmelCase				,		__lowerCAmelCase				,		__lowerCAmelCase				,		__lowerCAmelCase						)				,		batch_size=__lowerCAmelCase						)


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_						):




       '''simple docstring'''

       """Compute validation""" ""
       UpperCAmelCase_							:       Dict       =				{'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}
       if self.config.model_type != "distilbert":
            UpperCAmelCase_							:       List[Any]       =				(
                batch[2] if self.config.model_type in ['bert', 'xlnet'] else None
            )  # XLM and RoBERTa don"t use token_type_ids
       UpperCAmelCase_							:       Union[str, Any]       =				self(**__lowerCAmelCase						)
       UpperCAmelCase_       ,      UpperCAmelCase_							:       Optional[Any]       =				outputs[:2]
       UpperCAmelCase_							:       int       =				logits.detach().cpu().numpy()
       UpperCAmelCase_							:       Any       =				inputs['labels'].detach().cpu().numpy()
       return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}


  def _UpperCamelCase					(			self				,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:       List[str]       =				torch.stack([x['val_loss'] for x in outputs]						).mean()
       UpperCAmelCase_							:       List[str]       =				np.concatenate([x['pred'] for x in outputs]				,		axis=0						)
       UpperCAmelCase_							:       List[Any]       =				np.argmax(__lowerCAmelCase				,		axis=2						)
       UpperCAmelCase_							:       Dict       =				np.concatenate([x['target'] for x in outputs]				,		axis=0						)

       UpperCAmelCase_							:       Dict       =				dict(enumerate(self.labels						)						)
       UpperCAmelCase_							:       Optional[int]       =				[[] for _ in range(out_label_ids.shape[0]						)]
       UpperCAmelCase_							:       int       =				[[] for _ in range(out_label_ids.shape[0]						)]

       for i in range(out_label_ids.shape[0]						):
            for j in range(out_label_ids.shape[1]						):
                 if out_label_ids[i, j] != self.pad_token_label_id:
                      out_label_list[i].append(label_map[out_label_ids[i][j]]						)
                      preds_list[i].append(label_map[preds[i][j]]						)

       UpperCAmelCase_							:       Any       =				{
           'val_loss': val_loss_mean,
           'accuracy_score': accuracy_score(__lowerCAmelCase				,		__lowerCAmelCase						),
           'precision': precision_score(__lowerCAmelCase				,		__lowerCAmelCase						),
           'recall': recall_score(__lowerCAmelCase				,		__lowerCAmelCase						),
           'f1': fa_score(__lowerCAmelCase				,		__lowerCAmelCase						),
       }

       UpperCAmelCase_							:       List[str]       =				dict(results.items()						)
       UpperCAmelCase_							:       Dict       =				results
       return ret, preds_list, out_label_list


  def _UpperCamelCase					(			self				,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_							:       List[Any]       =				self._eval_end(__lowerCAmelCase						)
       UpperCAmelCase_							:       Optional[int]       =				ret['log']
       return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}


  def _UpperCamelCase					(			self				,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_							:       Optional[int]       =				self._eval_end(__lowerCAmelCase						)

       # Converting to the dict required by pl
       # https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
       # pytorch_lightning/trainer/logging.py#L139
       UpperCAmelCase_							:       Union[str, Any]       =				ret['log']
       # `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
       return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}






  @staticmethod
  def _UpperCamelCase					(			snake_case_				,		snake_case_						):




       '''simple docstring'''

       BaseTransformer.add_model_specific_args(__lowerCAmelCase				,		__lowerCAmelCase						)
       parser.add_argument(
           '--task_type'				,		default='NER'				,		type=__lowerCAmelCase				,		help='Task type to fine tune in training (e.g. NER, POS, etc)'						)
       parser.add_argument(
           '--max_seq_length'				,		default=1_2_8				,		type=__lowerCAmelCase				,		help=(
               'The maximum total input sequence length after tokenization. Sequences longer '
               'than this will be truncated, sequences shorter will be padded.'
           )				,		)

       parser.add_argument(
           '--labels'				,		default=''				,		type=__lowerCAmelCase				,		help='Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'				,		)
       parser.add_argument(
           '--gpus'				,		default=0				,		type=__lowerCAmelCase				,		help='The number of GPUs allocated for this, it is by default 0 meaning none'				,		)

       parser.add_argument(
           '--overwrite_cache'				,		action='store_true'				,		help='Overwrite the cached training and evaluation sets'						)

       return parser


if __name__ == "__main__":
 snake_case__					:      Union[str, Any]				      =  argparse.ArgumentParser()
 add_generic_args(parser, os.getcwd())
 snake_case__					:      Union[str, Any]				      =  NERTransformer.add_model_specific_args(parser, os.getcwd())
 snake_case__					:      Dict				      =  parser.parse_args()
 snake_case__					:      Tuple				      =  NERTransformer(args)
 snake_case__					:      Tuple				      =  generic_train(model, args)

 if args.do_predict:
  # See https://github.com/huggingface/transformers/issues/3159
  # pl use this default format to create a checkpoint:
  # https://github.com/PyTorchLightning/pytorch-lightning/blob/master\
  # /pytorch_lightning/callbacks/model_checkpoint.py#L322
  snake_case__					:      Any				      =  sorted(glob.glob(os.path.join(args.output_dir, '''checkpoint-epoch=*.ckpt'''), recursive=True))
  snake_case__					:      Any				      =  model.load_from_checkpoint(checkpoints[-1])
  trainer.test(model)







'''simple docstring'''


from sklearn.metrics import fa_score

import datasets


snake_case__					:      str				      =  '''
The F1 score is the harmonic mean of the precision and recall. It can be computed with the equation:
F1 = 2 * (precision * recall) / (precision + recall)
'''


snake_case__					:      int				      =  '''
Args:
    predictions (`list` of `int`): Predicted labels.
    references (`list` of `int`): Ground truth labels.
    labels (`list` of `int`): The set of labels to include when `average` is not set to `\'binary\'`, and the order of the labels if `average` is `None`. Labels present in the data can be excluded, for example to calculate a multiclass average ignoring a majority negative class. Labels not present in the data will result in 0 components in a macro average. For multilabel targets, labels are column indices. By default, all labels in `predictions` and `references` are used in sorted order. Defaults to None.
    pos_label (`int`): The class to be considered the positive class, in the case where `average` is set to `binary`. Defaults to 1.
    average (`string`): This parameter is required for multiclass/multilabel targets. If set to `None`, the scores for each class are returned. Otherwise, this determines the type of averaging performed on the data. Defaults to `\'binary\'`.

        - \'binary\': Only report results for the class specified by `pos_label`. This is applicable only if the classes found in `predictions` and `references` are binary.
        - \'micro\': Calculate metrics globally by counting the total true positives, false negatives and false positives.
        - \'macro\': Calculate metrics for each label, and find their unweighted mean. This does not take label imbalance into account.
        - \'weighted\': Calculate metrics for each label, and find their average weighted by support (the number of true instances for each label). This alters `\'macro\'` to account for label imbalance. This option can result in an F-score that is not between precision and recall.
        - \'samples\': Calculate metrics for each instance, and find their average (only meaningful for multilabel classification).
    sample_weight (`list` of `float`): Sample weights Defaults to None.

Returns:
    f1 (`float` or `array` of `float`): F1 score or list of f1 scores, depending on the value passed to `average`. Minimum possible value is 0. Maximum possible value is 1. Higher f1 scores are better.

Examples:

    Example 1-A simple binary example
        >>> f1_metric = datasets.load_metric("f1")
        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0])
        >>> print(results)
        {\'f1\': 0.5}

    Example 2-The same simple binary example as in Example 1, but with `pos_label` set to `0`.
        >>> f1_metric = datasets.load_metric("f1")
        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0], pos_label=0)
        >>> print(round(results[\'f1\'], 2))
        0.67

    Example 3-The same simple binary example as in Example 1, but with `sample_weight` included.
        >>> f1_metric = datasets.load_metric("f1")
        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0], sample_weight=[0.9, 0.5, 3.9, 1.2, 0.3])
        >>> print(round(results[\'f1\'], 2))
        0.35

    Example 4-A multiclass example, with different values for the `average` input.
        >>> predictions = [0, 2, 1, 0, 0, 1]
        >>> references = [0, 1, 2, 0, 1, 2]
        >>> results = f1_metric.compute(predictions=predictions, references=references, average="macro")
        >>> print(round(results[\'f1\'], 2))
        0.27
        >>> results = f1_metric.compute(predictions=predictions, references=references, average="micro")
        >>> print(round(results[\'f1\'], 2))
        0.33
        >>> results = f1_metric.compute(predictions=predictions, references=references, average="weighted")
        >>> print(round(results[\'f1\'], 2))
        0.27
        >>> results = f1_metric.compute(predictions=predictions, references=references, average=None)
        >>> print(results)
        {\'f1\': array([0.8, 0. , 0. ])}
'''


snake_case__					:      Tuple				      =  '''
@article{scikit-learn,
    title={Scikit-learn: Machine Learning in {P}ython},
    author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
           and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
           and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
           Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
    journal={Journal of Machine Learning Research},
    volume={12},
    pages={2825--2830},
    year={2011}
}
'''






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,     _KWARGS_DESCRIPTION	)
class 					__SCREAMING_SNAKE_CASE						(	datasets.Metric	):







  '''simple docstring'''






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return datasets.MetricInfo(
           description=_DESCRIPTION				,		citation=_CITATION				,		inputs_description=_KWARGS_DESCRIPTION				,		features=datasets.Features(
               {
                   'predictions': datasets.Sequence(datasets.Value('int32'						)						),
                   'references': datasets.Sequence(datasets.Value('int32'						)						),
               }
               if self.config_name == 'multilabel'
               else {
                   'predictions': datasets.Value('int32'						),
                   'references': datasets.Value('int32'						),
               }						)				,		reference_urls=['https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html']				,		)






  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_				,		snake_case_=None				,		snake_case_=1				,		snake_case_="binary"				,		snake_case_=None						):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[int]       =				fa_score(
           snake_case_				,		snake_case_				,		labels=snake_case_				,		pos_label=snake_case_				,		average=snake_case_				,		sample_weight=snake_case_						)
       return {"f1": float(snake_case_						) if score.size == 1 else score}






