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




class     _SCREAMING_SNAKE_CASE     (      snake_case_						):
	lowerCAmelCase__         = 'token-classification'



	def __init__(    self						,  lowercase )				->       Optional[Any]:
					if type(lowercase ) == dict:
									lowerCamelCase_						=			Namespace(**lowercase )
					lowerCamelCase_						=			import_module("tasks" )
					try:
									lowerCamelCase_						=			getattr(lowercase						,  hparams.task_type )
									lowerCamelCase_						=			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_						=			self.token_classification_task.get_labels(hparams.labels )
					lowerCamelCase_						=			CrossEntropyLoss().ignore_index
					super().__init__(lowercase						,  len(self.labels )						,  self.mode )



	def     SCREAMING_SNAKE_CASE_(    self						,  **lowercase )				->       List[str]:
					return self.model(**lowercase )



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase )				->       str:
					lowerCamelCase_						=			{"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
					if self.config.model_type != "distilbert":
									lowerCamelCase_						=			(
									    batch[2] if self.config.model_type in ["bert", "xlnet"] else None
									)  # XLM and RoBERTa don"t use token_type_ids

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



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



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase						,  lowercase = False )				->       DataLoader:
					lowerCamelCase_						=			self._feature_file(lowercase )
					logger.info("Loading features from cached file %s"						,  lowercase )
					lowerCamelCase_						=			torch.load(lowercase )
					lowerCamelCase_						=			torch.tensor([f.input_ids for f in features]						,  dtype=torch.long )
					lowerCamelCase_						=			torch.tensor([f.attention_mask for f in features]						,  dtype=torch.long )
					if features[0].token_type_ids is not None:
									lowerCamelCase_						=			torch.tensor([f.token_type_ids for f in features]						,  dtype=torch.long )
					else:
									lowerCamelCase_						=			torch.tensor([0 for f in features]						,  dtype=torch.long )
									# HACK(we will not use this anymore soon)
					lowerCamelCase_						=			torch.tensor([f.label_ids for f in features]						,  dtype=torch.long )
					return DataLoader(
					    TensorDataset(lowercase						,  lowercase						,  lowercase						,  lowercase )						,  batch_size=lowercase )



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



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase )				->       List[str]:
					lowerCamelCase_						=			torch.stack([x["val_loss"] for x in outputs] ).mean()
					lowerCamelCase_						=			np.concatenate([x["pred"] for x in outputs]						,  axis=0 )
					lowerCamelCase_						=			np.argmax(lowercase						,  axis=2 )
					lowerCamelCase_						=			np.concatenate([x["target"] for x in outputs]						,  axis=0 )

					lowerCamelCase_						=			dict(enumerate(self.labels ) )
					lowerCamelCase_						=			[[] for _ in range(out_label_ids.shape[0] )]
					lowerCamelCase_						=			[[] 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_						=			{
					    "val_loss": val_loss_mean,
					    "accuracy_score": accuracy_score(lowercase						,  lowercase ),
					    "precision": precision_score(lowercase						,  lowercase ),
					    "recall": recall_score(lowercase						,  lowercase ),
					    "f1": fa_score(lowercase						,  lowercase ),
					}

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



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase )				->       Dict:
					# when stable
					lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_						=			self._eval_end(lowercase )
					lowerCamelCase_						=			ret["log"]
					return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase )				->       List[Any]:
					# updating to test_epoch_end instead of deprecated test_end
					lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_						=			self._eval_end(lowercase )

					# Converting to the dict required by pl
					# https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
					# pytorch_lightning/trainer/logging.py#L139
					lowerCamelCase_						=			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     SCREAMING_SNAKE_CASE_(    lowercase						,  lowercase )				->       Any:
					# Add NER specific options
					BaseTransformer.add_model_specific_args(lowercase						,  lowercase )
					parser.add_argument(
					    "--task_type"						,  default="NER"						,  type=lowercase						,  help="Task type to fine tune in training (e.g. NER, POS, etc)" )
					parser.add_argument(
					    "--max_seq_length"						,  default=128						,  type=lowercase						,  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=lowercase						,  help="Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."						,  )
					parser.add_argument(
					    "--gpus"						,  default=0						,  type=lowercase						,  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				=argparse.ArgumentParser()
			add_generic_args(parser, os.getcwd())
			__A				=NERTransformer.add_model_specific_args(parser, os.getcwd())
			__A				=parser.parse_args()
			__A				=NERTransformer(args)
			__A				=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				=sorted(glob.glob(os.path.join(args.output_dir, '''checkpoint-epoch=*.ckpt'''), recursive=True))
						__A				=model.load_from_checkpoint(checkpoints[-1])
						trainer.test(model)




import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, 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     _SCREAMING_SNAKE_CASE     (      snake_case_						):



	def __init__(    self						,  *lowercase						,  lowercase=None						,  lowercase=None						,  **lowercase )				->       List[str]:
					super().__init__(*lowercase						,  **lowercase )
					lowerCamelCase_						=			eval_examples
					lowerCamelCase_						=			post_process_function



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase = None						,  lowercase=None						,  lowercase = None						,  lowercase = "eval"						,  **lowercase						,  )				->       Dict[str, float]:
					lowerCamelCase_						=			gen_kwargs.copy()
					lowerCamelCase_						=			(
					    gen_kwargs["max_length"] if gen_kwargs.get("max_length" ) is not None else self.args.generation_max_length
					)
					lowerCamelCase_						=			(
					    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams" ) is not None else self.args.generation_num_beams
					)
					lowerCamelCase_						=			gen_kwargs

					lowerCamelCase_						=			self.eval_dataset if eval_dataset is None else eval_dataset
					lowerCamelCase_						=			self.get_eval_dataloader(lowercase )
					lowerCamelCase_						=			self.eval_examples if eval_examples is None else eval_examples

					# Temporarily disable metric computation, we will do it in the loop here.
					lowerCamelCase_						=			self.compute_metrics
					lowerCamelCase_						=			None
					lowerCamelCase_						=			time.time()
					lowerCamelCase_						=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
									lowerCamelCase_						=			eval_loop(
									    lowercase						,  description="Evaluation"						,  prediction_loss_only=True if compute_metrics is None else None						,  ignore_keys=lowercase						,  metric_key_prefix=lowercase						,  )
					finally:
									lowerCamelCase_						=			compute_metrics
					lowerCamelCase_						=			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(
					        lowercase						,  lowercase						,  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
									lowerCamelCase_						=			self.post_process_function(lowercase						,  lowercase						,  lowercase )
									lowerCamelCase_						=			self.compute_metrics(lowercase )

									# Prefix all keys with metric_key_prefix + '_'
									for key in list(metrics.keys() ):
													if not key.startswith(f'{metric_key_prefix}_' ):
																	lowerCamelCase_						=			metrics.pop(lowercase )

									metrics.update(output.metrics )
					else:
									lowerCamelCase_						=			output.metrics

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

					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() )

					lowerCamelCase_						=			self.callback_handler.on_evaluate(self.args						,  self.state						,  self.control						,  lowercase )
					return metrics



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase						,  lowercase=None						,  lowercase = "test"						,  **lowercase )				->       Union[str, Any]:
					lowerCamelCase_						=			gen_kwargs.copy()

					lowerCamelCase_						=			self.get_test_dataloader(lowercase )

					# Temporarily disable metric computation, we will do it in the loop here.
					lowerCamelCase_						=			self.compute_metrics
					lowerCamelCase_						=			None
					lowerCamelCase_						=			time.time()
					lowerCamelCase_						=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
									lowerCamelCase_						=			eval_loop(
									    lowercase						,  description="Prediction"						,  prediction_loss_only=True if compute_metrics is None else None						,  ignore_keys=lowercase						,  metric_key_prefix=lowercase						,  )
					finally:
									lowerCamelCase_						=			compute_metrics

					lowerCamelCase_						=			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(
					        lowercase						,  lowercase						,  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

					lowerCamelCase_						=			self.post_process_function(lowercase						,  lowercase						,  lowercase						,  "predict" )
					lowerCamelCase_						=			self.compute_metrics(lowercase )

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


