






import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

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


__UpperCamelCase    :						Optional[Any]		     =  logging.getLogger(__name__)

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(		UpperCamelCase__		):
							def __init__(   self:					str				,		*UpperCamelCase:					Optional[int]				,		UpperCamelCase:					Optional[Any]=None				,		UpperCamelCase:					Optional[int]=None				,		UpperCamelCase:					Tuple=None				,		**UpperCamelCase:					Any  )						->						Union[str, Any]:
											super().__init__(*UpperCamelCase_				,		**UpperCamelCase_  )
											snake_case__													=   eval_examples
											snake_case__													=   post_process_function
											snake_case__													=   quant_trainer_args
											snake_case__													=   1_28  # default number of calibration samples
							def lowerCAmelCase_       (   self:					Union[str, Any]				,		UpperCamelCase:					Optional[Any]=None  )						->						Any:







											if calib_dataset is None and self.calib_dataset is None:
															raise ValueError('Trainer: calibration requires an calib_dataset.'  )
											snake_case__													=   calib_dataset if calib_dataset is not None else self.calib_dataset

											snake_case__													=   self._remove_unused_columns(UpperCamelCase_				,		description='Calibration'  )

											return DataLoader(
											    UpperCamelCase_				,		batch_size=self.args.eval_batch_size				,		collate_fn=self.data_collator				,		drop_last=self.args.dataloader_drop_last				,		num_workers=self.args.dataloader_num_workers				,		pin_memory=self.args.dataloader_pin_memory				,		shuffle=UpperCamelCase_				,		)
							def lowerCAmelCase_       (   self:					Union[str, Any]				,		UpperCamelCase:					str=None  )						->						Dict:
											snake_case__													=   self.train_dataset if calib_dataset is None else calib_dataset
											snake_case__													=   self.get_calib_dataloader(UpperCamelCase_  )

											snake_case__													=   self.model
											quant_trainer.configure_model(UpperCamelCase_				,		self.quant_trainer_args				,		calib=UpperCamelCase_  )
											model.eval()
											quant_trainer.enable_calibration(UpperCamelCase_  )

											logger.info('***** Running calibration *****'  )
											logger.info(F'''  Num examples = {self.calib_num}'''  )
											logger.info(F'''  Batch size = {calib_dataloader.batch_size}'''  )

											for step, inputs in enumerate(UpperCamelCase_  ):
															# Prediction step
															snake_case__							,  snake_case__							,  snake_case__													=   self.prediction_step(UpperCamelCase_				,		UpperCamelCase_				,		prediction_loss_only=UpperCamelCase_  )
															if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
																			break

											quant_trainer.finish_calibration(UpperCamelCase_				,		self.quant_trainer_args  )
											snake_case__													=   model
							def lowerCAmelCase_       (   self:					Union[str, Any]				,		UpperCamelCase:					Union[str, Any]=None				,		UpperCamelCase:					List[Any]=None				,		UpperCamelCase:					List[str]=None				,		UpperCamelCase:					str = "eval"  )						->						Tuple:
											snake_case__													=   self.eval_dataset if eval_dataset is None else eval_dataset
											snake_case__													=   self.get_eval_dataloader(UpperCamelCase_  )
											snake_case__													=   self.eval_examples if eval_examples is None else eval_examples

											# Temporarily disable metric computation, we will do it in the loop here.
											snake_case__													=   self.compute_metrics
											snake_case__													=   None
											snake_case__													=   self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
															snake_case__													=   eval_loop(
															    UpperCamelCase_				,		description='Evaluation'				,		prediction_loss_only=True if compute_metrics is None else None				,		ignore_keys=UpperCamelCase_				,		)
											finally:
															snake_case__													=   compute_metrics

											if self.post_process_function is not None and self.compute_metrics is not None:
															snake_case__													=   self.post_process_function(UpperCamelCase_				,		UpperCamelCase_				,		output.predictions  )
															snake_case__													=   self.compute_metrics(UpperCamelCase_  )

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

															self.log(UpperCamelCase_  )
											else:
															snake_case__													=   {}

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

											snake_case__													=   self.callback_handler.on_evaluate(self.args				,		self.state				,		self.control				,		UpperCamelCase_  )
											return metrics
							def lowerCAmelCase_       (   self:					int				,		UpperCamelCase:					Optional[Any]				,		UpperCamelCase:					Tuple				,		UpperCamelCase:					List[Any]=None				,		UpperCamelCase:					str = "test"  )						->						Optional[int]:
											snake_case__													=   self.get_test_dataloader(UpperCamelCase_  )

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

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

											snake_case__													=   self.post_process_function(UpperCamelCase_				,		UpperCamelCase_				,		output.predictions				,		'predict'  )
											snake_case__													=   self.compute_metrics(UpperCamelCase_  )

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

											return PredictionOutput(predictions=predictions.predictions				,		label_ids=predictions.label_ids				,		metrics=UpperCamelCase_  )


							def lowerCAmelCase_       (   self:					int				,		UpperCamelCase:					Tuple="./"  )						->						str:
											snake_case__													=   self.eval_dataset
											snake_case__													=   self.get_eval_dataloader(UpperCamelCase_  )

											snake_case__													=   next(iter(UpperCamelCase_  )  )

											# saving device - to make it consistent
											snake_case__													=   torch.device('cuda' if torch.cuda.is_available() else 'cpu'  )

											# convert to tuple
											snake_case__													=   tuple(v.to(UpperCamelCase_  ) for k, v in batch.items()  )

											logger.info('Converting model to be onnx compatible'  )
											from pytorch_quantization.nn import TensorQuantizer

											snake_case__													=   True

											snake_case__													=   self.model.to(UpperCamelCase_  )

											model.eval()
											model.float()

											snake_case__													=   model.module if hasattr(UpperCamelCase_				,		'module'  ) else model
											quant_trainer.configure_model(UpperCamelCase_				,		self.quant_trainer_args  )

											snake_case__													=   os.path.join(UpperCamelCase_				,		'model.onnx'  )
											logger.info(F'''exporting model to {output_model_file}'''  )

											snake_case__													=   {0: 'batch_size', 1: 'seq_len'}

											torch.onnx.export(
											    UpperCamelCase_				,		UpperCamelCase_				,		UpperCamelCase_				,		export_params=UpperCamelCase_				,		opset_version=13				,		do_constant_folding=UpperCamelCase_				,		input_names=['input_ids', 'attention_mask', 'token_type_ids']				,		output_names=['output_start_logits', 'output_end_logits']				,		dynamic_axes={
											        'input_ids': axes,
											        'attention_mask': axes,
											        'token_type_ids': axes,
											        'output_start_logits': axes,
											        'output_end_logits': axes,
											    }				,		verbose=UpperCamelCase_				,		)
											logger.info('onnx export finished'  )








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


class        _a							(			UpperCamelCase__			):
							_lowercase	:			Union[str, Any]        = '''token-classification'''

							def __init__(							self:	int    ,						UpperCamelCase_:	Optional[Any]							)			->   Dict:





									"""simple docstring"""







									if type(UpperCamelCase_							) == dict:
											lowercase__            =      Namespace(**UpperCamelCase_							)
									lowercase__            =      import_module('''tasks'''							)
									try:
											lowercase__            =      getattr(UpperCamelCase_    ,						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__(UpperCamelCase_    ,						len(self.labels							)    ,						self.mode							)

							def       lowerCamelCase_ (							self:	Tuple    ,						**UpperCamelCase_:	Optional[int]							)			->   str:





									"""simple docstring"""







									return self.model(**UpperCamelCase_							)

							def       lowerCamelCase_ (							self:	List[str]    ,						UpperCamelCase_:	List[Any]    ,						UpperCamelCase_:	int							)			->   int:





									"""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(**UpperCamelCase_							)
									lowercase__            =      outputs[0]
									# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
									return {"loss": loss}

							def       lowerCamelCase_ (							self:	Optional[Any]							)			->   Optional[int]:





									"""simple docstring"""







									lowercase__            =      self.hparams
									for mode in ["train", "dev", "test"]:
											lowercase__            =      self._feature_file(UpperCamelCase_							)
											if os.path.exists(UpperCamelCase_							) and not args.overwrite_cache:
													logger.info('''Loading features from cached file %s'''    ,						UpperCamelCase_							)
													lowercase__            =      torch.load(UpperCamelCase_							)
											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    ,						UpperCamelCase_							)
													lowercase__            =      self.token_classification_task.convert_examples_to_features(
													    UpperCamelCase_    ,						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=UpperCamelCase_    ,						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'''    ,						UpperCamelCase_							)
													torch.save(UpperCamelCase_    ,						UpperCamelCase_							)

							def       lowerCamelCase_ (							self:	Union[str, Any]    ,						UpperCamelCase_:	int    ,						UpperCamelCase_:	int    ,						UpperCamelCase_:	bool = False							)			->   DataLoader:





									"""simple docstring"""







									lowercase__            =      self._feature_file(UpperCamelCase_							)
									logger.info('''Loading features from cached file %s'''    ,						UpperCamelCase_							)
									lowercase__            =      torch.load(UpperCamelCase_							)
									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(UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_							)    ,						batch_size=UpperCamelCase_							)

							def       lowerCamelCase_ (							self:	Any    ,						UpperCamelCase_:	int    ,						UpperCamelCase_:	List[Any]							)			->   Union[str, Any]:





									"""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(**UpperCamelCase_							)
									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       lowerCamelCase_ (							self:	Dict    ,						UpperCamelCase_:	List[str]							)			->   int:





									"""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(UpperCamelCase_    ,						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(UpperCamelCase_    ,						UpperCamelCase_							),
									    '''precision''': precision_score(UpperCamelCase_    ,						UpperCamelCase_							),
									    '''recall''': recall_score(UpperCamelCase_    ,						UpperCamelCase_							),
									    '''f1''': fa_score(UpperCamelCase_    ,						UpperCamelCase_							),
									}

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

							def       lowerCamelCase_ (							self:	List[Any]    ,						UpperCamelCase_:	List[Any]							)			->   Dict:





									"""simple docstring"""







									lowercase__     ,		lowercase__     ,		lowercase__            =      self._eval_end(UpperCamelCase_							)
									lowercase__            =      ret['''log''']
									return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}

							def       lowerCamelCase_ (							self:	List[str]    ,						UpperCamelCase_:	Dict							)			->   Dict:





									"""simple docstring"""







									lowercase__     ,		lowercase__     ,		lowercase__            =      self._eval_end(UpperCamelCase_							)

									# 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       lowerCamelCase_ (							UpperCamelCase_:	Optional[Any]    ,						UpperCamelCase_:	str							)			->   Optional[Any]:





									"""simple docstring"""







									BaseTransformer.add_model_specific_args(UpperCamelCase_    ,						UpperCamelCase_							)
									parser.add_argument(
									    '''--task_type'''    ,						default='''NER'''    ,						type=UpperCamelCase_    ,						help='''Task type to fine tune in training (e.g. NER, POS, etc)'''							)
									parser.add_argument(
									    '''--max_seq_length'''    ,						default=128    ,						type=UpperCamelCase_    ,						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=UpperCamelCase_    ,						help='''Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'''    ,						)
									parser.add_argument(
									    '''--gpus'''    ,						default=0    ,						type=UpperCamelCase_    ,						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       =							argparse.ArgumentParser()
			add_generic_args(parser, os.getcwd())
			lowerCAmelCase       =							NERTransformer.add_model_specific_args(parser, os.getcwd())
			lowerCAmelCase       =							parser.parse_args()
			lowerCAmelCase       =							NERTransformer(args)
			lowerCAmelCase       =							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       =							sorted(glob.glob(os.path.join(args.output_dir, 'checkpoint-epoch=*.ckpt'), recursive=True))
						lowerCAmelCase       =							model.load_from_checkpoint(checkpoints[-1])
						trainer.test(model)


