


import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
       import tensorflow as tf
       from tensorflow.python.framework.errors_impl import ResourceExhaustedError

       from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
       import pyanvml.pyanvml as nvml

__UpperCamelCase				:  Optional[int]      					=							logging.get_logger(__name__)
def      snake_case_						( __lowercase			,       __lowercase       ):
   def run_func(__lowercase       ):
      @wraps(__lowercase       )
      def run_in_eager_mode(*__lowercase			,       **__lowercase       ):
         return func(*__lowercase			,       **__lowercase       )

      @wraps(__lowercase       )
      @tf.function(experimental_compile=__lowercase       )
      def run_in_graph_mode(*__lowercase			,       **__lowercase       ):
         return func(*__lowercase			,       **__lowercase       )

      if do_eager_mode is True:
         if use_xla is not False:
            raise ValueError(
                '''Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'''       )
         return run_in_eager_mode
      else:
         return run_in_graph_mode

   return run_func
def      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      int						=   random.Random()
   UpperCAmelCase_							:      str						=   [rng.randint(0			,       vocab_size - 1       ) for i in range(batch_size * sequence_length       )]
   return tf.constant(__lowercase			,       shape=(batch_size, sequence_length)			,       dtype=tf.intaa       )





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     TensorFlowBenchmarkArguments
    A_  :     PretrainedConfig
    A_  :     str			= "TensorFlow"





    @property
    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       return tf.__version__





    def 					_lowerCamelCase (	self				:   Dict       ,   __snake_case				:   str       ,   __snake_case				:   int       ,   __snake_case				:   int	):





       '''simple docstring'''
       # initialize GPU on separate process
       UpperCAmelCase_							:      Optional[Any]						=   self.args.strategy
       if strategy is None:
          raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''	)
       UpperCAmelCase_							:      int						=   self._prepare_inference_func(__snake_case       ,   __snake_case       ,   __snake_case	)
       return self._measure_speed(_inference	)





    def 					_lowerCamelCase (	self				:   str       ,   __snake_case				:   str       ,   __snake_case				:   int       ,   __snake_case				:   int	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   self.args.strategy
       if strategy is None:
          raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''	)
       UpperCAmelCase_							:      int						=   self._prepare_train_func(__snake_case       ,   __snake_case       ,   __snake_case	)
       return self._measure_speed(_train	)





    def 					_lowerCamelCase (	self				:   Any       ,   __snake_case				:   str       ,   __snake_case				:   int       ,   __snake_case				:   int	):





       '''simple docstring'''
       # initialize GPU on separate process
       if self.args.is_gpu:
          tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]       ,   __snake_case	)
       UpperCAmelCase_							:      str						=   self.args.strategy
       if strategy is None:
          raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''	)
       UpperCAmelCase_							:      Tuple						=   self._prepare_inference_func(__snake_case       ,   __snake_case       ,   __snake_case	)
       return self._measure_memory(_inference	)





    def 					_lowerCamelCase (	self				:   List[Any]       ,   __snake_case				:   str       ,   __snake_case				:   int       ,   __snake_case				:   int	):





       '''simple docstring'''
       if self.args.is_gpu:
          tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]       ,   __snake_case	)
       UpperCAmelCase_							:      Union[str, Any]						=   self.args.strategy
       if strategy is None:
          raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''	)

       UpperCAmelCase_							:      Union[str, Any]						=   self._prepare_train_func(__snake_case       ,   __snake_case       ,   __snake_case	)
       return self._measure_memory(_train	)





    def 					_lowerCamelCase (	self				:   Optional[Any]       ,   __snake_case				:   str       ,   __snake_case				:   int       ,   __snake_case				:   int	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[int]						=   self.config_dict[model_name]

       if self.args.fpaa:
          raise NotImplementedError('''Mixed precision is currently not supported.'''	)

       UpperCAmelCase_							:      Optional[int]						=   (
           hasattr(__snake_case       ,   '''architectures'''	)
           and isinstance(config.architectures       ,   __snake_case	)
           and len(config.architectures	) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
          try:
             UpperCAmelCase_							:      Optional[Any]						=   '''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
             UpperCAmelCase_							:      Optional[Any]						=   __import__('''transformers'''       ,   fromlist=[model_class]	)
             UpperCAmelCase_							:      Optional[Any]						=   getattr(__snake_case       ,   __snake_case	)
             UpperCAmelCase_							:      Tuple						=   model_cls(__snake_case	)
          except ImportError:
             raise ImportError(
                 f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                 ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
       else:
          UpperCAmelCase_							:      List[Any]						=   TF_MODEL_MAPPING[config.__class__](__snake_case	)

       # encoder-decoder has vocab size saved differently
       UpperCAmelCase_							:      Union[str, Any]						=   config.vocab_size if hasattr(__snake_case       ,   '''vocab_size'''	) else config.encoder.vocab_size
       UpperCAmelCase_							:      List[str]						=   random_input_ids(__snake_case       ,   __snake_case       ,   __snake_case	)

       @run_with_tf_optimizations(self.args.eager_mode       ,   self.args.use_xla	)
       def encoder_decoder_forward():
          return model(__snake_case       ,   decoder_input_ids=__snake_case       ,   training=__snake_case	)

       @run_with_tf_optimizations(self.args.eager_mode       ,   self.args.use_xla	)
       def encoder_forward():
          return model(__snake_case       ,   training=__snake_case	)

       UpperCAmelCase_							:      Union[str, Any]						=   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

       return _inference





    def 					_lowerCamelCase (	self				:   List[Any]       ,   __snake_case				:   str       ,   __snake_case				:   int       ,   __snake_case				:   int	):





       '''simple docstring'''
       UpperCAmelCase_							:      str						=   self.config_dict[model_name]

       if self.args.eager_mode is not False:
          raise ValueError('''Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.'''	)

       if self.args.fpaa:
          raise NotImplementedError('''Mixed precision is currently not supported.'''	)

       UpperCAmelCase_							:      Any						=   (
           hasattr(__snake_case       ,   '''architectures'''	)
           and isinstance(config.architectures       ,   __snake_case	)
           and len(config.architectures	) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
          try:
             UpperCAmelCase_							:      Optional[Any]						=   '''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
             UpperCAmelCase_							:      int						=   __import__('''transformers'''       ,   fromlist=[model_class]	)
             UpperCAmelCase_							:      List[Any]						=   getattr(__snake_case       ,   __snake_case	)
             UpperCAmelCase_							:      Optional[int]						=   model_cls(__snake_case	)
          except ImportError:
             raise ImportError(
                 f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                 ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
       else:
          UpperCAmelCase_							:      Optional[Any]						=   TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__snake_case	)

       # encoder-decoder has vocab size saved differently
       UpperCAmelCase_							:      Union[str, Any]						=   config.vocab_size if hasattr(__snake_case       ,   '''vocab_size'''	) else config.encoder.vocab_size
       UpperCAmelCase_							:      Union[str, Any]						=   random_input_ids(__snake_case       ,   __snake_case       ,   __snake_case	)

       @run_with_tf_optimizations(self.args.eager_mode       ,   self.args.use_xla	)
       def encoder_decoder_train():
          UpperCAmelCase_							:      Optional[int]						=   model(__snake_case       ,   decoder_input_ids=__snake_case       ,   labels=__snake_case       ,   training=__snake_case	)[0]
          UpperCAmelCase_							:      List[Any]						=   tf.gradients(__snake_case       ,   model.trainable_variables	)
          return gradients

       @run_with_tf_optimizations(self.args.eager_mode       ,   self.args.use_xla	)
       def encoder_train():
          UpperCAmelCase_							:      Optional[Any]						=   model(__snake_case       ,   labels=__snake_case       ,   training=__snake_case	)[0]
          UpperCAmelCase_							:      Any						=   tf.gradients(__snake_case       ,   model.trainable_variables	)
          return gradients

       UpperCAmelCase_							:      Dict						=   encoder_decoder_train if config.is_encoder_decoder else encoder_train

       return _train





    def 					_lowerCamelCase (	self				:   Dict       ,   __snake_case				:   Union[str, Any]	):





       '''simple docstring'''
       with self.args.strategy.scope():
          try:
             if self.args.is_tpu or self.args.use_xla:
                # run additional 10 times to stabilize compilation for tpu
                logger.info('''Do inference on TPU. Running model 5 times to stabilize compilation'''	)
                timeit.repeat(__snake_case       ,   repeat=1       ,   number=5	)

             # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
             UpperCAmelCase_							:      Tuple						=   timeit.repeat(
                 __snake_case       ,   repeat=self.args.repeat       ,   number=10       ,   )

             return min(__snake_case	) / 10.0
          except ResourceExhaustedError as e:
             self.print_fn(f'''Doesn\'t fit on GPU. {e}'''	)





    def 					_lowerCamelCase (	self				:   Dict       ,   __snake_case				:   Callable[[], None]	):





       '''simple docstring'''
       logger.info(
           '''Note that TensorFlow allocates more memory than '''
           '''it might need to speed up computation. '''
           '''The memory reported here corresponds to the memory '''
           '''reported by `nvidia-smi`, which can vary depending '''
           '''on total available memory on the GPU that is used.'''	)
       with self.args.strategy.scope():
          try:
             if self.args.trace_memory_line_by_line:
                if not self.args.eager_mode:
                   raise ValueError(
                       '''`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'''
                       ''' consumption line by line.'''	)
                UpperCAmelCase_							:      Optional[int]						=   start_memory_tracing('''transformers'''	)

             if self.args.is_tpu:
                # tpu
                raise NotImplementedError(
                    '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking'''
                    ''' with `args.memory=False`'''	)
             elif self.args.is_gpu:
                # gpu
                if not is_pyanvml_available():
                   logger.warning(
                       '''py3nvml not installed, we won\'t log GPU memory usage. '''
                       '''Install py3nvml (pip install py3nvml) to log information about GPU.'''	)
                   UpperCAmelCase_							:      Any						=   '''N/A'''
                else:
                   logger.info(
                       '''Measuring total GPU usage on GPU device. Make sure to not have additional processes'''
                       ''' running on the same GPU.'''	)
                   # init nvml
                   nvml.nvmlInit()
                   func()
                   UpperCAmelCase_							:      Optional[int]						=   nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
                   UpperCAmelCase_							:      Union[str, Any]						=   nvml.nvmlDeviceGetMemoryInfo(__snake_case	)
                   UpperCAmelCase_							:      Optional[int]						=   meminfo.used
                   UpperCAmelCase_							:      Union[str, Any]						=   Memory(__snake_case	)
                   # shutdown nvml
                   nvml.nvmlShutdown()
             else:
                # cpu
                if self.args.trace_memory_line_by_line:
                   logger.info(
                       '''When enabling line by line tracing, the max peak memory for CPU is inaccurate in'''
                       ''' TensorFlow.'''	)
                   UpperCAmelCase_							:      str						=   None
                else:
                   UpperCAmelCase_							:      Optional[Any]						=   measure_peak_memory_cpu(__snake_case	)
                   UpperCAmelCase_							:      Any						=   Memory(__snake_case	) if isinstance(__snake_case       ,   __snake_case	) else memory_bytes
             if self.args.trace_memory_line_by_line:
                UpperCAmelCase_							:      Optional[int]						=   stop_memory_tracing(__snake_case	)
                if memory is None:
                   UpperCAmelCase_							:      Tuple						=   summary.total
             else:
                UpperCAmelCase_							:      Dict						=   None

             return memory, summary
          except ResourceExhaustedError as e:
             self.print_fn(f'''Doesn\'t fit on GPU. {e}'''	)
             return "N/A", None



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				:  Dict      					=							logging.getLogger(__name__)





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     Optional[int]			= 'token-classification'





    def __init__(	self				:   Any       ,   __snake_case				:   Optional[Any]	):





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





    def 					_lowerCamelCase (	self				:   Optional[int]       ,   **__snake_case				:   Optional[Any]	):





       '''simple docstring'''
       return self.model(**__snake_case	)





    def 					_lowerCamelCase (	self				:   str       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   Optional[Any]	):





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

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





    def 					_lowerCamelCase (	self				:   Dict	):





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





    def 					_lowerCamelCase (	self				:   Any       ,   __snake_case				:   int       ,   __snake_case				:   int       ,   __snake_case				:   bool = False	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[str]						=   self._feature_file(__snake_case	)
       logger.info('''Loading features from cached file %s'''       ,   __snake_case	)
       UpperCAmelCase_							:      Optional[int]						=   torch.load(__snake_case	)
       UpperCAmelCase_							:      Any						=   torch.tensor([f.input_ids for f in features]       ,   dtype=torch.long	)
       UpperCAmelCase_							:      int						=   torch.tensor([f.attention_mask for f in features]       ,   dtype=torch.long	)
       if features[0].token_type_ids is not None:
          UpperCAmelCase_							:      Any						=   torch.tensor([f.token_type_ids for f in features]       ,   dtype=torch.long	)
       else:
          UpperCAmelCase_							:      Optional[int]						=   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(__snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case	)       ,   batch_size=__snake_case	)





    def 					_lowerCamelCase (	self				:   List[Any]       ,   __snake_case				:   Dict       ,   __snake_case				:   Union[str, Any]	):





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





    def 					_lowerCamelCase (	self				:   List[str]       ,   __snake_case				:   Optional[int]	):





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

       UpperCAmelCase_							:      Tuple						=   dict(enumerate(self.labels	)	)
       UpperCAmelCase_							:      Union[str, Any]						=   [[] for _ in range(out_label_ids.shape[0]	)]
       UpperCAmelCase_							:      Optional[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_							:      Union[str, Any]						=   {
           '''val_loss''': val_loss_mean,
           '''accuracy_score''': accuracy_score(__snake_case       ,   __snake_case	),
           '''precision''': precision_score(__snake_case       ,   __snake_case	),
           '''recall''': recall_score(__snake_case       ,   __snake_case	),
           '''f1''': fa_score(__snake_case       ,   __snake_case	),
       }

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





    def 					_lowerCamelCase (	self				:   List[str]       ,   __snake_case				:   int	):





       '''simple docstring'''
       # when stable
       UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_							:      str						=   self._eval_end(__snake_case	)
       UpperCAmelCase_							:      int						=   ret['''log''']
       return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}





    def 					_lowerCamelCase (	self				:   List[Any]       ,   __snake_case				:   Tuple	):





       '''simple docstring'''
       # updating to test_epoch_end instead of deprecated test_end
       UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_							:      Tuple						=   self._eval_end(__snake_case	)

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





       '''simple docstring'''
       # Add NER specific options
       BaseTransformer.add_model_specific_args(__snake_case       ,   __snake_case	)
       parser.add_argument(
           '''--task_type'''       ,   default='''NER'''       ,   type=__snake_case       ,   help='''Task type to fine tune in training (e.g. NER, POS, etc)'''	)
       parser.add_argument(
           '''--max_seq_length'''       ,   default=128       ,   type=__snake_case       ,   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=__snake_case       ,   help='''Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'''       ,   )
       parser.add_argument(
           '''--gpus'''       ,   default=0       ,   type=__snake_case       ,   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				:  str      					=							argparse.ArgumentParser()
       add_generic_args(parser, os.getcwd())
       __UpperCamelCase				:  Dict      					=							NERTransformer.add_model_specific_args(parser, os.getcwd())
       __UpperCamelCase				:  Tuple      					=							parser.parse_args()
       __UpperCamelCase				:  Optional[Any]      					=							NERTransformer(args)
       __UpperCamelCase				:  int      					=							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				:  Optional[int]      					=							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)