






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__	:			List[str]			   = logging.getLogger(__name__)
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				lowercase__											=       """token-classification"""




				def __init__(     self						:    List[Any],  lowerCamelCase						:    str  ):

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




				def 				lowercase__   (     self						:    Optional[Any],  **lowerCamelCase						:    Union[str, Any]  ):

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




				def 				lowercase__   (     self						:    Any,  lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    List[str]  ):

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

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




				def 				lowercase__   (     self						:    List[str]  ):

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

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




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    List[Any],  lowerCamelCase						:    Optional[int]  ):

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




				def 				lowercase__   (     self						:    Dict,  lowerCamelCase						:    Optional[Any]  ):

								'''simple docstring'''
								lowercase__    =			torch.stack([x['''val_loss'''] for x in outputs]  ).mean()
								lowercase__    =			np.concatenate([x['''pred'''] for x in outputs],  axis=0  )
								lowercase__    =			np.argmax(lowerCamelCase,  axis=2  )
								lowercase__    =			np.concatenate([x['''target'''] for x in outputs],  axis=0  )

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

								lowercase__    =			{
								    '''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  ),
								}

								lowercase__    =			dict(results.items()  )
								lowercase__    =			results
								return ret, preds_list, out_label_list




				def 				lowercase__   (     self						:    Dict,  lowerCamelCase						:    List[str]  ):

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




				def 				lowercase__   (     self						:    Any,  lowerCamelCase						:    int  ):

								'''simple docstring'''
								# updating to test_epoch_end instead of deprecated test_end
								lowercase__ ,						lowercase__ ,						lowercase__    =			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
								lowercase__    =			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 				lowercase__   (     lowerCamelCase						:    int,  lowerCamelCase						:    Any  ):

								'''simple docstring'''
								# Add NER specific options
								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=128,  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__":
							A__	:			Dict			   = argparse.ArgumentParser()
							add_generic_args(parser, os.getcwd())
							A__	:			Tuple			   = NERTransformer.add_model_specific_args(parser, os.getcwd())
							A__	:			List[Any]			   = parser.parse_args()
							A__	:			Union[str, Any]			   = NERTransformer(args)
							A__	:			Dict			   = 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__	:			Optional[Any]			   = model.load_from_checkpoint(checkpoints[-1])
														trainer.test(model)












import math
import time

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm
							import torch_xla.debug.metrics as met
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				def __init__(     self						:    Dict,  *lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    Any=None,  **lowerCamelCase						:    str  ):

								'''simple docstring'''
								super().__init__(*lowerCamelCase,  **lowerCamelCase  )
								lowercase__    =			eval_examples
								lowercase__    =			post_process_function




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    str=None,  lowerCamelCase						:    Optional[Any]=None,  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    str = "eval"  ):

								'''simple docstring'''
								lowercase__    =			self.eval_dataset if eval_dataset is None else eval_dataset
								lowercase__    =			self.get_eval_dataloader(lowerCamelCase  )
								lowercase__    =			self.eval_examples if eval_examples is None else eval_examples

								# Temporarily disable metric computation, we will do it in the loop here.
								lowercase__    =			self.compute_metrics
								lowercase__    =			None
								lowercase__    =			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								lowercase__    =			time.time()
								try:
												lowercase__    =			eval_loop(
												    lowerCamelCase,  description='''Evaluation''',  prediction_loss_only=True if compute_metrics is None else None,  ignore_keys=lowerCamelCase,  metric_key_prefix=lowerCamelCase,  )
								finally:
												lowercase__    =			compute_metrics
								lowercase__    =			self.args.eval_batch_size * self.args.world_size
								if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
												start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
								output.metrics.update(
								    speed_metrics(
								        lowerCamelCase,  lowerCamelCase,  num_samples=output.num_samples,  num_steps=math.ceil(output.num_samples / total_batch_size  ),  )  )
								if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
												# Only the main node write the results by default
												lowercase__    =			self.post_process_function(lowerCamelCase,  lowerCamelCase,  output.predictions  )
												lowercase__    =			self.compute_metrics(lowerCamelCase  )

												# Prefix all keys with metric_key_prefix + '_'
												for key in list(metrics.keys()  ):
																if not key.startswith(F"""{metric_key_prefix}_"""  ):
																				lowercase__    =			metrics.pop(lowerCamelCase  )
												metrics.update(output.metrics  )
								else:
												lowercase__    =			output.metrics

								if self.args.should_log:
												# Only the main node log the results by default
												self.log(lowerCamelCase  )

								if self.args.tpu_metrics_debug or self.args.debug:
												# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
												xm.master_print(met.metrics_report()  )

								lowercase__    =			self.callback_handler.on_evaluate(self.args,  self.state,  self.control,  lowerCamelCase  )
								return metrics




				def 				lowercase__   (     self						:    List[Any],  lowerCamelCase						:    Any,  lowerCamelCase						:    Dict,  lowerCamelCase						:    int=None,  lowerCamelCase						:    str = "test"  ):

								'''simple docstring'''
								lowercase__    =			self.get_test_dataloader(lowerCamelCase  )

								# Temporarily disable metric computation, we will do it in the loop here.
								lowercase__    =			self.compute_metrics
								lowercase__    =			None
								lowercase__    =			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								lowercase__    =			time.time()
								try:
												lowercase__    =			eval_loop(
												    lowerCamelCase,  description='''Prediction''',  prediction_loss_only=True if compute_metrics is None else None,  ignore_keys=lowerCamelCase,  metric_key_prefix=lowerCamelCase,  )
								finally:
												lowercase__    =			compute_metrics
								lowercase__    =			self.args.eval_batch_size * self.args.world_size
								if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
												start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
								output.metrics.update(
								    speed_metrics(
								        lowerCamelCase,  lowerCamelCase,  num_samples=output.num_samples,  num_steps=math.ceil(output.num_samples / total_batch_size  ),  )  )

								if self.post_process_function is None or self.compute_metrics is None:
												return output

								lowercase__    =			self.post_process_function(lowerCamelCase,  lowerCamelCase,  output.predictions,  '''predict'''  )
								lowercase__    =			self.compute_metrics(lowerCamelCase  )

								# Prefix all keys with metric_key_prefix + '_'
								for key in list(metrics.keys()  ):
												if not key.startswith(F"""{metric_key_prefix}_"""  ):
																lowercase__    =			metrics.pop(lowerCamelCase  )
								metrics.update(output.metrics  )
								return PredictionOutput(predictions=predictions.predictions,  label_ids=predictions.label_ids,  metrics=lowerCamelCase  )




