



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


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



class 				__snake_case						(     __lowerCamelCase						):
							_a			      = '''token-classification'''




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




							def    UpperCAmelCase__			(       self     :       Union[str, Any]				,				**A_     :       Union[str, Any]):
									return self.model(**A_)




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

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




							def    UpperCAmelCase__			(       self     :       List[str]):
									lowerCAmelCase_			:     Tuple									= self.hparams
									for mode in ["train", "dev", "test"]:
											lowerCAmelCase_			:     Union[str, Any]									= self._feature_file(A_)
											if os.path.exists(A_) and not args.overwrite_cache:
													logger.info('''Loading features from cached file %s'''				,				A_)
													lowerCAmelCase_			:     Tuple									= torch.load(A_)
											else:
													logger.info('''Creating features from dataset file at %s'''				,				args.data_dir)
													lowerCAmelCase_			:     List[Any]									= self.token_classification_task.read_examples_from_file(args.data_dir				,				A_)
													lowerCAmelCase_			:     Optional[int]									= 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     :       List[str]				,				A_     :       int				,				A_     :       int				,				A_     :       bool = False):
									lowerCAmelCase_			:     List[str]									= self._feature_file(A_)
									logger.info('''Loading features from cached file %s'''				,				A_)
									lowerCAmelCase_			:     int									= torch.load(A_)
									lowerCAmelCase_			:     Optional[int]									= torch.tensor([f.input_ids for f in features]				,				dtype=torch.long)
									lowerCAmelCase_			:     Optional[Any]									= torch.tensor([f.attention_mask for f in features]				,				dtype=torch.long)
									if features[0].token_type_ids is not None:
											lowerCAmelCase_			:     Dict									= torch.tensor([f.token_type_ids for f in features]				,				dtype=torch.long)
									else:
											lowerCAmelCase_			:     Tuple									= torch.tensor([0 for f in features]				,				dtype=torch.long)
											# HACK(we will not use this anymore soon)
									lowerCAmelCase_			:     Optional[int]									= 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     :       Tuple				,				A_     :       Optional[Any]				,				A_     :       int):
									"""Compute validation""" ""
									lowerCAmelCase_			:     List[str]									= {"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
									if self.config.model_type != "distilbert":
											lowerCAmelCase_			:     Union[str, Any]									= (
											    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
											)  # XLM and RoBERTa don"t use token_type_ids
									lowerCAmelCase_			:     List[str]									= self(**A_)
									lowerCAmelCase_			:     int									= outputs[:2]
									lowerCAmelCase_			:     List[str]									= logits.detach().cpu().numpy()
									lowerCAmelCase_			:     Any									= inputs["""labels"""].detach().cpu().numpy()
									return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}




							def    UpperCAmelCase__			(       self     :       Dict				,				A_     :       Dict):
									lowerCAmelCase_			:     Union[str, Any]									= torch.stack([x['''val_loss'''] for x in outputs]).mean()
									lowerCAmelCase_			:     Tuple									= np.concatenate([x['''pred'''] for x in outputs]				,				axis=0)
									lowerCAmelCase_			:     List[str]									= np.argmax(A_				,				axis=2)
									lowerCAmelCase_			:     Optional[Any]									= np.concatenate([x['''target'''] for x in outputs]				,				axis=0)

									lowerCAmelCase_			:     Tuple									= dict(enumerate(self.labels))
									lowerCAmelCase_			:     Dict									= [[] for _ in range(out_label_ids.shape[0])]
									lowerCAmelCase_			:     Tuple									= [[] 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]])

									lowerCAmelCase_			:     Dict									= {
									    """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_),
									}

									lowerCAmelCase_			:     Dict									= dict(results.items())
									lowerCAmelCase_			:     Dict									= results
									return ret, preds_list, out_label_list




							def    UpperCAmelCase__			(       self     :       Tuple				,				A_     :       int):
									# when stable
									lowerCAmelCase_			:     List[str]									= self._eval_end(A_)
									lowerCAmelCase_			:     int									= ret["""log"""]
									return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}




							def    UpperCAmelCase__			(       self     :       Dict				,				A_     :       Optional[int]):
									# updating to test_epoch_end instead of deprecated test_end
									lowerCAmelCase_			:     int									= 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
									lowerCAmelCase_			:     List[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__			(       A_     :       Any				,				A_     :       int):
									# Add NER specific options
									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__":
						A__  :  Tuple								=					argparse.ArgumentParser()
						add_generic_args(parser, os.getcwd())
						A__  :  Tuple								=					NERTransformer.add_model_specific_args(parser, os.getcwd())
						A__  :  List[str]								=					parser.parse_args()
						A__  :  Optional[int]								=					NERTransformer(args)
						A__  :  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
												A__  :  int								=					sorted(glob.glob(os.path.join(args.output_dir, '''checkpoint-epoch=*.ckpt'''), recursive=True))
												A__  :  Any								=					model.load_from_checkpoint(checkpoints[-1])
												trainer.test(model)












import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
      import tensorflow as tf

      from transformers import TensorFlowBenchmark, TensorFlowBenchmarkArguments







@require_tf
class     __snake_case       (							unittest.TestCase			):





       '''simple docstring'''







       def 							UpperCAmelCase__ (  self :    Any							,				A :    Union[str, Any]		):
         for model_result in results.values():
           for batch_size, sequence_length in zip(model_result["""bs"""]							,				model_result["""ss"""]		):
             __snake_case:						Dict   =						model_result["""result"""][batch_size][sequence_length]
             self.assertIsNotNone(A		)




       def 							UpperCAmelCase__ (  self :    List[Any]		):
         __snake_case:						List[Any]   =						"""sshleifer/tiny-gpt2"""
         __snake_case:						Union[str, Any]   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				eager_mode=A							,				multi_process=A							,				)
         __snake_case:						Dict   =						TensorFlowBenchmark(A		)
         __snake_case:						List[str]   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result		)
         self.check_results_dict_not_empty(results.memory_inference_result		)




       def 							UpperCAmelCase__ (  self :    Optional[int]		):
         __snake_case:						Optional[int]   =						"""sgugger/tiny-distilbert-classification"""
         __snake_case:						str   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				multi_process=A							,				only_pretrain_model=A							,				)
         __snake_case:						Any   =						TensorFlowBenchmark(A		)
         __snake_case:						List[Any]   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result		)
         self.check_results_dict_not_empty(results.memory_inference_result		)




       def 							UpperCAmelCase__ (  self :    int		):
         __snake_case:						Dict   =						"""sshleifer/tiny-gpt2"""
         __snake_case:						Union[str, Any]   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				multi_process=A							,				)
         __snake_case:						int   =						TensorFlowBenchmark(A		)
         __snake_case:						str   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result		)
         self.check_results_dict_not_empty(results.memory_inference_result		)




       def 							UpperCAmelCase__ (  self :    List[Any]		):
         __snake_case:						List[Any]   =						"""sshleifer/tiny-gpt2"""
         __snake_case:						int   =						AutoConfig.from_pretrained(A		)
         __snake_case:						Optional[Any]   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				eager_mode=A							,				multi_process=A							,				)
         __snake_case:						Union[str, Any]   =						TensorFlowBenchmark(A							,				[config]		)
         __snake_case:						Tuple   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result		)
         self.check_results_dict_not_empty(results.memory_inference_result		)




       def 							UpperCAmelCase__ (  self :    Any		):
         __snake_case:						Optional[int]   =						"""sshleifer/tiny-gpt2"""
         __snake_case:						Tuple   =						AutoConfig.from_pretrained(A		)
         __snake_case:						Union[str, Any]   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				multi_process=A							,				)
         __snake_case:						str   =						TensorFlowBenchmark(A							,				[config]		)
         __snake_case:						str   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result		)
         self.check_results_dict_not_empty(results.memory_inference_result		)




       def 							UpperCAmelCase__ (  self :    Any		):
         __snake_case:						Tuple   =						"""sshleifer/tiny-gpt2"""
         __snake_case:						str   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				multi_process=A							,				)
         __snake_case:						Tuple   =						TensorFlowBenchmark(A		)
         __snake_case:						Tuple   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_train_result		)
         self.check_results_dict_not_empty(results.memory_train_result		)




       def 							UpperCAmelCase__ (  self :    int		):
         __snake_case:						Dict   =						"""sshleifer/tiny-gpt2"""
         __snake_case:						Union[str, Any]   =						AutoConfig.from_pretrained(A		)
         __snake_case:						Optional[int]   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				multi_process=A							,				)
         __snake_case:						List[str]   =						TensorFlowBenchmark(A							,				[config]		)
         __snake_case:						Dict   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_train_result		)
         self.check_results_dict_not_empty(results.memory_train_result		)




       def 							UpperCAmelCase__ (  self :    List[str]		):
         __snake_case:						List[Any]   =						"""patrickvonplaten/t5-tiny-random"""
         __snake_case:						List[str]   =						AutoConfig.from_pretrained(A		)
         __snake_case:						List[str]   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				multi_process=A							,				)
         __snake_case:						Optional[int]   =						TensorFlowBenchmark(A							,				configs=[config]		)
         __snake_case:						Tuple   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result		)
         self.check_results_dict_not_empty(results.memory_inference_result		)




       @unittest.skipIf(is_tf_available() and len(tf.config.list_physical_devices("""GPU"""		)		) == 0							,				"""Cannot do xla on CPU."""		)
       def 							UpperCAmelCase__ (  self :    Optional[Any]		):
         __snake_case:						Optional[Any]   =						"""sshleifer/tiny-gpt2"""
         __snake_case:						List[Any]   =						TensorFlowBenchmarkArguments(
             models=[MODEL_ID]							,				training=A							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				use_xla=A							,				multi_process=A							,				)
         __snake_case:						Union[str, Any]   =						TensorFlowBenchmark(A		)
         __snake_case:						List[Any]   =						benchmark.run()
         self.check_results_dict_not_empty(results.time_inference_result		)
         self.check_results_dict_not_empty(results.memory_inference_result		)




       def 							UpperCAmelCase__ (  self :    Union[str, Any]		):
         __snake_case:						List[str]   =						"""sshleifer/tiny-gpt2"""
         with tempfile.TemporaryDirectory() as tmp_dir:
           __snake_case:						int   =						TensorFlowBenchmarkArguments(
               models=[MODEL_ID]							,				inference=A							,				save_to_csv=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				inference_time_csv_file=os.path.join(A							,				"""inf_time.csv"""		)							,				inference_memory_csv_file=os.path.join(A							,				"""inf_mem.csv"""		)							,				env_info_csv_file=os.path.join(A							,				"""env.csv"""		)							,				multi_process=A							,				)
           __snake_case:						Tuple   =						TensorFlowBenchmark(A		)
           benchmark.run()
           self.assertTrue(Path(os.path.join(A							,				"""inf_time.csv"""		)		).exists()		)
           self.assertTrue(Path(os.path.join(A							,				"""inf_mem.csv"""		)		).exists()		)
           self.assertTrue(Path(os.path.join(A							,				"""env.csv"""		)		).exists()		)






       def 							UpperCAmelCase__ (  self :    Optional[int]		):
         __snake_case:						Any   =						"""sshleifer/tiny-gpt2"""

         def _check_summary_is_not_empty(A :    Union[str, Any]		):
           self.assertTrue(hasattr(A							,				"""sequential"""		)		)
           self.assertTrue(hasattr(A							,				"""cumulative"""		)		)
           self.assertTrue(hasattr(A							,				"""current"""		)		)
           self.assertTrue(hasattr(A							,				"""total"""		)		)

         with tempfile.TemporaryDirectory() as tmp_dir:
           __snake_case:						Union[str, Any]   =						TensorFlowBenchmarkArguments(
               models=[MODEL_ID]							,				inference=A							,				sequence_lengths=[8]							,				batch_sizes=[1]							,				log_filename=os.path.join(A							,				"""log.txt"""		)							,				log_print=A							,				trace_memory_line_by_line=A							,				eager_mode=A							,				multi_process=A							,				)
           __snake_case:						Dict   =						TensorFlowBenchmark(A		)
           __snake_case:						List[str]   =						benchmark.run()
           _check_summary_is_not_empty(result.inference_summary		)
           self.assertTrue(Path(os.path.join(A							,				"""log.txt"""		)		).exists()		)






