






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


__lowerCamelCase :      List[str]										=				logging.getLogger(__name__)



class 							SCREAMING_SNAKE_CASE__						(    _a						):





					"""simple docstring"""






					a_     =		"token-classification"



					def __init__( self					:  Dict							,		__A					:  Dict	):
						if type(_a	) == dict:
							snake_case__			:				Union[str, Any]       =  Namespace(**_a	)
						snake_case__			:				Union[str, Any]       =  import_module("tasks"	)
						try:
							snake_case__			:				List[Any]       =  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__			:				Union[str, Any]       =  self.token_classification_task.get_labels(hparams.labels	)
						snake_case__			:				Union[str, Any]       =  CrossEntropyLoss().ignore_index
						super().__init__(_a							,		len(self.labels	)							,		self.mode	)



					def 	_lowercase					( self					:  Union[str, Any]							,		**__A					:  List[Any]	):
						return self.model(**_a	)



					def 	_lowercase					( self					:  List[Any]							,		__A					:  Dict							,		__A					:  int	):
						snake_case__			:				Union[str, Any]       =  {"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__			:				Optional[int]       =  self(**_a	)
						snake_case__			:				str       =  outputs[0]
						# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
						return {"loss": loss}



					def 	_lowercase					( self					:  int	):
						snake_case__			:				Dict       =  self.hparams
						for mode in ["train", "dev", "test"]:
							snake_case__			:				str       =  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__			:				Tuple       =  torch.load(_a	)
							else:
								logger.info("Creating features from dataset file at %s"							,		args.data_dir	)
								snake_case__			:				List[Any]       =  self.token_classification_task.read_examples_from_file(args.data_dir							,		_a	)
								snake_case__			:				Tuple       =  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 	_lowercase					( self					:  Optional[Any]							,		__A					:  int							,		__A					:  int							,		__A					:  bool = False	):
						snake_case__			:				Optional[int]       =  self._feature_file(_a	)
						logger.info("Loading features from cached file %s"							,		_a	)
						snake_case__			:				int       =  torch.load(_a	)
						snake_case__			:				Union[str, Any]       =  torch.tensor([f.input_ids for f in features]							,		dtype=torch.long	)
						snake_case__			:				Union[str, Any]       =  torch.tensor([f.attention_mask for f in features]							,		dtype=torch.long	)
						if features[0].token_type_ids is not None:
							snake_case__			:				Dict       =  torch.tensor([f.token_type_ids for f in features]							,		dtype=torch.long	)
						else:
							snake_case__			:				str       =  torch.tensor([0 for f in features]							,		dtype=torch.long	)
							# HACK(we will not use this anymore soon)
						snake_case__			:				List[str]       =  torch.tensor([f.label_ids for f in features]							,		dtype=torch.long	)
						return DataLoader(
						    TensorDataset(_a							,		_a							,		_a							,		_a	)							,		batch_size=_a	)



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



					def 	_lowercase					( self					:  int							,		__A					:  Union[str, Any]	):
						snake_case__			:				Dict       =  torch.stack([x["val_loss"] for x in outputs]	).mean()
						snake_case__			:				Optional[Any]       =  np.concatenate([x["pred"] for x in outputs]							,		axis=0	)
						snake_case__			:				Any       =  np.argmax(_a							,		axis=2	)
						snake_case__			:				List[Any]       =  np.concatenate([x["target"] for x in outputs]							,		axis=0	)

						snake_case__			:				Tuple       =  dict(enumerate(self.labels	)	)
						snake_case__			:				List[Any]       =  [[] for _ in range(out_label_ids.shape[0]	)]
						snake_case__			:				str       =  [[] 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__			:				str       =  {
						    "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__			:				Union[str, Any]       =  results
						return ret, preds_list, out_label_list



					def 	_lowercase					( self					:  Optional[int]							,		__A					:  Tuple	):
						# when stable
						snake_case__			:				Tuple       =  self._eval_end(_a	)
						snake_case__			:				int       =  ret["log"]
						return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



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











from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch



@require_torch
@require_tf
class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def 	_lowercase					( self					:  Optional[Any]	):
						snake_case__			:				Optional[Any]       =  SMALL_MODEL_IDENTIFIER
						snake_case__			:				Any       =  "pt"
						snake_case__			:				Any       =  "tf"



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  List[Any]	):
						snake_case__			:				int       =  AutoModel.from_pretrained(self.test_model	)
						model_pt.save_pretrained(__A	)



					def 	_lowercase					( self					:  Optional[int]							,		__A					:  Tuple	):
						snake_case__			:				List[Any]       =  TFAutoModel.from_pretrained(self.test_model							,		from_pt=__A	)
						model_tf.save_pretrained(__A	)



					def 	_lowercase					( self					:  str	):
						snake_case__			:				Optional[Any]       =  "mock_framework"

						# Framework provided - return whatever the user provides
						snake_case__			:				Optional[Any]       =  FeaturesManager.determine_framework(self.test_model							,		__A	)
						self.assertEqual(__A							,		__A	)

						# Local checkpoint and framework provided - return provided framework
						# PyTorch checkpoint
						with TemporaryDirectory() as local_pt_ckpt:
							self._setup_pt_ckpt(__A	)
							snake_case__			:				Optional[int]       =  FeaturesManager.determine_framework(__A							,		__A	)
							self.assertEqual(__A							,		__A	)

						# TensorFlow checkpoint
						with TemporaryDirectory() as local_tf_ckpt:
							self._setup_tf_ckpt(__A	)
							snake_case__			:				int       =  FeaturesManager.determine_framework(__A							,		__A	)
							self.assertEqual(__A							,		__A	)



					def 	_lowercase					( self					:  Dict	):
						# PyTorch checkpoint
						with TemporaryDirectory() as local_pt_ckpt:
							self._setup_pt_ckpt(__A	)
							snake_case__			:				List[str]       =  FeaturesManager.determine_framework(__A	)
							self.assertEqual(__A							,		self.framework_pt	)

						# TensorFlow checkpoint
						with TemporaryDirectory() as local_tf_ckpt:
							self._setup_tf_ckpt(__A	)
							snake_case__			:				Tuple       =  FeaturesManager.determine_framework(__A	)
							self.assertEqual(__A							,		self.framework_tf	)

						# Invalid local checkpoint
						with TemporaryDirectory() as local_invalid_ckpt:
							with self.assertRaises(__A	):
								snake_case__			:				int       =  FeaturesManager.determine_framework(__A	)






					def 	_lowercase					( self					:  Dict	):
						snake_case__			:				Dict       =  MagicMock(return_value=__A	)
						with patch("transformers.onnx.features.is_tf_available"							,		__A	):
							snake_case__			:				List[Any]       =  FeaturesManager.determine_framework(self.test_model	)
							self.assertEqual(__A							,		self.framework_pt	)

						# PyTorch not in environment -> use TensorFlow
						snake_case__			:				Tuple       =  MagicMock(return_value=__A	)
						with patch("transformers.onnx.features.is_torch_available"							,		__A	):
							snake_case__			:				int       =  FeaturesManager.determine_framework(self.test_model	)
							self.assertEqual(__A							,		self.framework_tf	)

						# Both in environment -> use PyTorch
						snake_case__			:				Dict       =  MagicMock(return_value=__A	)
						snake_case__			:				Optional[int]       =  MagicMock(return_value=__A	)
						with patch("transformers.onnx.features.is_tf_available"							,		__A	), patch(
						    "transformers.onnx.features.is_torch_available"							,		__A	):
							snake_case__			:				Optional[Any]       =  FeaturesManager.determine_framework(self.test_model	)
							self.assertEqual(__A							,		self.framework_pt	)

						# Both not in environment -> raise error
						snake_case__			:				List[str]       =  MagicMock(return_value=__A	)
						snake_case__			:				Optional[Any]       =  MagicMock(return_value=__A	)
						with patch("transformers.onnx.features.is_tf_available"							,		__A	), patch(
						    "transformers.onnx.features.is_torch_available"							,		__A	):
							with self.assertRaises(__A	):
								snake_case__			:				Optional[Any]       =  FeaturesManager.determine_framework(self.test_model	)



