






"""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


UpperCamelCase_	:				str										= logging.getLogger(__name__)


class      _lowercase (     lowerCAmelCase						):
      _a   :     Optional[int]								=   '''token-classification'''


      def __init__(    self				:	Union[str, Any]							,   a				:	Optional[int]						):





            """simple docstring"""







            if type(a						) == dict:
                  __snake_case      :    Optional[int]  			=Namespace(**a						)
            __snake_case      :    Optional[int]  			=import_module('''tasks'''						)
            try:
                  __snake_case      :    Optional[int]  			=getattr(a							,   hparams.task_type						)
                  __snake_case      :    TokenClassificationTask  			=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__()}'''						)
            __snake_case      :    str  			=self.token_classification_task.get_labels(hparams.labels						)
            __snake_case      :    Tuple  			=CrossEntropyLoss().ignore_index
            super().__init__(a							,   len(self.labels						)							,   self.mode						)


      def _UpperCamelCase						(    self				:	Optional[Any]							,   **a				:	Tuple						):





            """simple docstring"""







            return self.model(**a						)


      def _UpperCamelCase						(    self				:	Any							,   a				:	str							,   a				:	Tuple						):





            """simple docstring"""







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

            __snake_case      :    Union[str, Any]  			=self(**a						)
            __snake_case      :    int  			=outputs[0]
            # tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
            return {"loss": loss}


      def _UpperCamelCase						(    self				:	str						):





            """simple docstring"""







            __snake_case      :    Optional[int]  			=self.hparams
            for mode in ["train", "dev", "test"]:
                  __snake_case      :    List[Any]  			=self._feature_file(a						)
                  if os.path.exists(a						) and not args.overwrite_cache:
                        logger.info('''Loading features from cached file %s'''							,   a						)
                        __snake_case      :    Optional[Any]  			=torch.load(a						)
                  else:
                        logger.info('''Creating features from dataset file at %s'''							,   args.data_dir						)
                        __snake_case      :    List[str]  			=self.token_classification_task.read_examples_from_file(args.data_dir							,   a						)
                        __snake_case      :    List[str]  			=self.token_classification_task.convert_examples_to_features(
                            a							,   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=a							,   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'''							,   a						)
                        torch.save(a							,   a						)


      def _UpperCamelCase						(    self				:	Optional[int]							,   a				:	int							,   a				:	int							,   a				:	bool = False						):





            """simple docstring"""







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


      def _UpperCamelCase						(    self				:	Optional[int]							,   a				:	int							,   a				:	int						):





            """simple docstring"""







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


      def _UpperCamelCase						(    self				:	Union[str, Any]							,   a				:	Optional[int]						):





            """simple docstring"""







            __snake_case      :    List[Any]  			=torch.stack([x['''val_loss'''] for x in outputs]						).mean()
            __snake_case      :    Optional[int]  			=np.concatenate([x['''pred'''] for x in outputs]							,   axis=0						)
            __snake_case      :    Union[str, Any]  			=np.argmax(a							,   axis=2						)
            __snake_case      :    str  			=np.concatenate([x['''target'''] for x in outputs]							,   axis=0						)

            __snake_case      :    Union[str, Any]  			=dict(enumerate(self.labels						)						)
            __snake_case      :    List[Any]  			=[[] for _ in range(out_label_ids.shape[0]						)]
            __snake_case      :    Any  			=[[] 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]]						)

            __snake_case      :    int  			={
                '''val_loss''': val_loss_mean,
                '''accuracy_score''': accuracy_score(a							,   a						),
                '''precision''': precision_score(a							,   a						),
                '''recall''': recall_score(a							,   a						),
                '''f1''': fa_score(a							,   a						),
            }

            __snake_case      :    Any  			=dict(results.items()						)
            __snake_case      :    int  			=results
            return ret, preds_list, out_label_list


      def _UpperCamelCase						(    self				:	List[str]							,   a				:	List[Any]						):





            """simple docstring"""







            __snake_case			,				__snake_case			,				__snake_case      :    str  			=self._eval_end(a						)
            __snake_case      :    Optional[Any]  			=ret['''log''']
            return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}


      def _UpperCamelCase						(    self				:	List[Any]							,   a				:	Any						):





            """simple docstring"""







            __snake_case			,				__snake_case			,				__snake_case      :    Tuple  			=self._eval_end(a						)

            # Converting to the dict required by pl
            # https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
            # pytorch_lightning/trainer/logging.py#L139
            __snake_case      :    Tuple  			=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						(    a				:	int							,   a				:	Optional[int]						):





            """simple docstring"""







            BaseTransformer.add_model_specific_args(a							,   a						)
            parser.add_argument(
                '''--task_type'''							,   default='''NER'''							,   type=a							,   help='''Task type to fine tune in training (e.g. NER, POS, etc)'''						)
            parser.add_argument(
                '''--max_seq_length'''							,   default=1_2_8							,   type=a							,   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=a							,   help='''Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'''							,   )
            parser.add_argument(
                '''--gpus'''							,   default=0							,   type=a							,   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__":
    UpperCamelCase_	:				List[Any]										= argparse.ArgumentParser()
    add_generic_args(parser, os.getcwd())
    UpperCamelCase_	:				Optional[Any]										= NERTransformer.add_model_specific_args(parser, os.getcwd())
    UpperCamelCase_	:				Optional[Any]										= parser.parse_args()
    UpperCamelCase_	:				List[Any]										= NERTransformer(args)
    UpperCamelCase_	:				str										= 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
        UpperCamelCase_	:				str										= sorted(glob.glob(os.path.join(args.output_dir, """checkpoint-epoch=*.ckpt"""), recursive=True))
        UpperCamelCase_	:				List[Any]										= model.load_from_checkpoint(checkpoints[-1])
        trainer.test(model)










"""simple docstring"""





import json
import os
import unittest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin


@require_tokenizers
class      _lowercase (     lowerCAmelCase		,  unittest.TestCase						):
      _a   :     List[str]								=   MgpstrTokenizer
      _a   :     int								=   False
      _a   :     List[str]								=   {}
      _a   :     Optional[Any]								=   False


      def _UpperCamelCase						(    self				:	str						):





            """simple docstring"""







            super().setUp()

            # fmt: off
            __snake_case      :    Union[str, Any]  			=['''[GO]''', '''[s]''', '''0''', '''1''', '''2''', '''3''', '''4''', '''5''', '''6''', '''7''', '''8''', '''9''', '''a''', '''b''', '''c''', '''d''', '''e''', '''f''', '''g''', '''h''', '''i''', '''j''', '''k''', '''l''', '''m''', '''n''', '''o''', '''p''', '''q''', '''r''', '''s''', '''t''', '''u''', '''v''', '''w''', '''x''', '''y''', '''z''']
            # fmt: on
            __snake_case      :    Optional[Any]  			=dict(zip(a							,   range(len(a						)						)						)						)

            __snake_case      :    int  			=os.path.join(self.tmpdirname							,   VOCAB_FILES_NAMES['''vocab_file''']						)
            with open(self.vocab_file							,   '''w'''							,   encoding='''utf-8'''						) as fp:
                  fp.write(json.dumps(a						) + '''\n'''						)


      def _UpperCamelCase						(    self				:	Any							,   **a				:	int						):





            """simple docstring"""







            return MgpstrTokenizer.from_pretrained(self.tmpdirname							,   **a						)


      def _UpperCamelCase						(    self				:	int							,   a				:	Any						):





            """simple docstring"""







            __snake_case      :    Dict  			='''tester'''
            __snake_case      :    str  			='''tester'''
            return input_text, output_text


      @unittest.skip('''MGP-STR always lower cases letters.'''						)
      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            pass


      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            __snake_case      :    List[Any]  			=self.get_tokenizers(do_lower_case=a						)
            for tokenizer in tokenizers:
                  with self.subTest(f'''{tokenizer.__class__.__name__}'''						):
                        __snake_case      :    List[Any]  			='''[SPECIAL_TOKEN]'''

                        tokenizer.add_special_tokens({'''cls_token''': special_token}						)
                        __snake_case      :    Tuple  			=tokenizer.encode([special_token]							,   add_special_tokens=a						)
                        self.assertEqual(len(a						)							,   1						)

                        __snake_case      :    int  			=tokenizer.decode(a							,   skip_special_tokens=a						)
                        self.assertTrue(special_token not in decoded						)


      def _UpperCamelCase						(    self				:	Tuple						):





            """simple docstring"""







            __snake_case      :    Union[str, Any]  			=self.get_tokenizers()
            for tokenizer in tokenizers:
                  with self.subTest(f'''{tokenizer.__class__.__name__}'''						):
                        __snake_case			,				__snake_case      :    Tuple  			=self.get_input_output_texts(a						)

                        __snake_case      :    int  			=tokenizer.tokenize(a						)
                        __snake_case      :    str  			=tokenizer.convert_tokens_to_ids(a						)
                        __snake_case      :    Union[str, Any]  			=tokenizer.encode(a							,   add_special_tokens=a						)
                        self.assertListEqual(a							,   a						)

                        __snake_case      :    Tuple  			=tokenizer.convert_ids_to_tokens(a						)
                        self.assertNotEqual(len(a						)							,   0						)
                        __snake_case      :    Optional[Any]  			=tokenizer.decode(a						)
                        self.assertIsInstance(a							,   a						)

                        self.assertEqual(text_a.replace(''' '''							,   ''''''						)							,   a						)


      @unittest.skip('''MGP-STR tokenizer only handles one sequence.'''						)
      def _UpperCamelCase						(    self				:	List[Any]						):





            """simple docstring"""







            pass







      @unittest.skip('''inputs cannot be pretokenized in MgpstrTokenizer'''						)
      def _UpperCamelCase						(    self				:	List[str]						):





            """simple docstring"""







            pass


