



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


UpperCAmelCase     		=   logging.getLogger(__name__)
class __snake_case(				_lowerCAmelCase ):


						'''simple docstring'''



						UpperCAmelCase							:   Union[str, Any]									=		"token-classification"


						def __init__( self     ,	A_				)     ->     List[Any]:
							if type(A_				) == dict:
								lowerCAmelCase   =     Namespace(**A_				)
							lowerCAmelCase   =     import_module("""tasks"""				)
							try:
								lowerCAmelCase   =     getattr(A_     ,	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__(A_     ,	len(self.labels				)     ,	self.mode				)


						def     __snake_case   ( self     ,	**A_				)     ->     Tuple:
							return self.model(**A_				)


						def     __snake_case   ( self     ,	A_     ,	A_				)     ->     Optional[int]:
							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(**A_				)
							lowerCAmelCase   =     outputs[0]
							# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
							return {"loss": loss}


						def     __snake_case   ( self				)     ->     Tuple:
							lowerCAmelCase   =     self.hparams
							for mode in ["train", "dev", "test"]:
								lowerCAmelCase   =     self._feature_file(A_				)
								if os.path.exists(A_				) and not args.overwrite_cache:
									logger.info("""Loading features from cached file %s"""     ,	A_				)
									lowerCAmelCase   =     torch.load(A_				)
								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     ,	A_				)
									lowerCAmelCase   =     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     __snake_case   ( self     ,	A_     ,	A_     ,	A_ = False				)     ->     DataLoader:
							lowerCAmelCase   =     self._feature_file(A_				)
							logger.info("""Loading features from cached file %s"""     ,	A_				)
							lowerCAmelCase   =     torch.load(A_				)
							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(A_     ,	A_     ,	A_     ,	A_				)     ,	batch_size=A_				)


						def     __snake_case   ( self     ,	A_     ,	A_				)     ->     Tuple:
							"""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(**A_				)
							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     __snake_case   ( self     ,	A_				)     ->     Dict:
							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(A_     ,	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(A_     ,	A_				),
							    """precision""": precision_score(A_     ,	A_				),
							    """recall""": recall_score(A_     ,	A_				),
							    """f1""": fa_score(A_     ,	A_				),
							}

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


						def     __snake_case   ( self     ,	A_				)     ->     Tuple:
							# when stable
							lowerCAmelCase,       lowerCAmelCase,       lowerCAmelCase   =     self._eval_end(A_				)
							lowerCAmelCase   =     ret["""log"""]
							return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}


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




'''simple docstring'''

import unittest

from transformers import LiltConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import (
				    LiltForQuestionAnswering,
				    LiltForSequenceClassification,
				    LiltForTokenClassification,
				    LiltModel,
				)
				from transformers.models.lilt.modeling_lilt import LILT_PRETRAINED_MODEL_ARCHIVE_LIST

class __snake_case:


						'''simple docstring'''



						def __init__( self     ,	A_     ,	A_=13     ,	A_=7     ,	A_=True     ,	A_=True     ,	A_=True     ,	A_=True     ,	A_=99     ,	A_=24     ,	A_=2     ,	A_=6     ,	A_=37     ,	A_="gelu"     ,	A_=0.1     ,	A_=0.1     ,	A_=512     ,	A_=16     ,	A_=2     ,	A_=0.0_2     ,	A_=3     ,	A_=None     ,	A_=1000     ,	)     ->     Union[str, Any]:
							lowerCAmelCase   =     parent
							lowerCAmelCase   =     batch_size
							lowerCAmelCase   =     seq_length
							lowerCAmelCase   =     is_training
							lowerCAmelCase   =     use_input_mask
							lowerCAmelCase   =     use_token_type_ids
							lowerCAmelCase   =     use_labels
							lowerCAmelCase   =     vocab_size
							lowerCAmelCase   =     hidden_size
							lowerCAmelCase   =     num_hidden_layers
							lowerCAmelCase   =     num_attention_heads
							lowerCAmelCase   =     intermediate_size
							lowerCAmelCase   =     hidden_act
							lowerCAmelCase   =     hidden_dropout_prob
							lowerCAmelCase   =     attention_probs_dropout_prob
							lowerCAmelCase   =     max_position_embeddings
							lowerCAmelCase   =     type_vocab_size
							lowerCAmelCase   =     type_sequence_label_size
							lowerCAmelCase   =     initializer_range
							lowerCAmelCase   =     num_labels
							lowerCAmelCase   =     scope
							lowerCAmelCase   =     range_bbox


						def     __snake_case   ( self				)     ->     List[Any]:
							lowerCAmelCase   =     ids_tensor([self.batch_size, self.seq_length]     ,	self.vocab_size				)

							lowerCAmelCase   =     ids_tensor([self.batch_size, self.seq_length, 4]     ,	self.range_bbox				)
							# Ensure that bbox is legal
							for i in range(bbox.shape[0]				):
								for j in range(bbox.shape[1]				):
									if bbox[i, j, 3] < bbox[i, j, 1]:
										lowerCAmelCase   =     bbox[i, j, 3]
										lowerCAmelCase   =     bbox[i, j, 1]
										lowerCAmelCase   =     t
									if bbox[i, j, 2] < bbox[i, j, 0]:
										lowerCAmelCase   =     bbox[i, j, 2]
										lowerCAmelCase   =     bbox[i, j, 0]
										lowerCAmelCase   =     t

							lowerCAmelCase   =     None
							if self.use_input_mask:
								lowerCAmelCase   =     ids_tensor([self.batch_size, self.seq_length]     ,	vocab_size=2				)

							lowerCAmelCase   =     None
							if self.use_token_type_ids:
								lowerCAmelCase   =     ids_tensor([self.batch_size, self.seq_length]     ,	self.type_vocab_size				)

							lowerCAmelCase   =     None
							lowerCAmelCase   =     None
							if self.use_labels:
								lowerCAmelCase   =     ids_tensor([self.batch_size]     ,	self.type_sequence_label_size				)
								lowerCAmelCase   =     ids_tensor([self.batch_size, self.seq_length]     ,	self.num_labels				)

							lowerCAmelCase   =     self.get_config()

							return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels


						def     __snake_case   ( self				)     ->     int:
							return LiltConfig(
							    vocab_size=self.vocab_size     ,	hidden_size=self.hidden_size     ,	num_hidden_layers=self.num_hidden_layers     ,	num_attention_heads=self.num_attention_heads     ,	intermediate_size=self.intermediate_size     ,	hidden_act=self.hidden_act     ,	hidden_dropout_prob=self.hidden_dropout_prob     ,	attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,	max_position_embeddings=self.max_position_embeddings     ,	type_vocab_size=self.type_vocab_size     ,	initializer_range=self.initializer_range     ,	)


						def     __snake_case   ( self     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	)     ->     int:
							lowerCAmelCase   =     LiltModel(config=A_				)
							model.to(A_				)
							model.eval()
							lowerCAmelCase   =     model(A_     ,	bbox=A_     ,	attention_mask=A_     ,	token_type_ids=A_				)
							lowerCAmelCase   =     model(A_     ,	bbox=A_     ,	token_type_ids=A_				)
							lowerCAmelCase   =     model(A_     ,	bbox=A_				)

							self.parent.assertEqual(result.last_hidden_state.shape     ,	(self.batch_size, self.seq_length, self.hidden_size)				)
							self.parent.assertEqual(result.pooler_output.shape     ,	(self.batch_size, self.hidden_size)				)


						def     __snake_case   ( self     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	)     ->     Union[str, Any]:
							lowerCAmelCase   =     self.num_labels
							lowerCAmelCase   =     LiltForTokenClassification(config=A_				)
							model.to(A_				)
							model.eval()
							lowerCAmelCase   =     model(
							    A_     ,	bbox=A_     ,	attention_mask=A_     ,	token_type_ids=A_     ,	labels=A_				)
							self.parent.assertEqual(result.logits.shape     ,	(self.batch_size, self.seq_length, self.num_labels)				)


						def     __snake_case   ( self     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_     ,	)     ->     Optional[Any]:
							lowerCAmelCase   =     LiltForQuestionAnswering(config=A_				)
							model.to(A_				)
							model.eval()
							lowerCAmelCase   =     model(
							    A_     ,	bbox=A_     ,	attention_mask=A_     ,	token_type_ids=A_     ,	start_positions=A_     ,	end_positions=A_     ,	)
							self.parent.assertEqual(result.start_logits.shape     ,	(self.batch_size, self.seq_length)				)
							self.parent.assertEqual(result.end_logits.shape     ,	(self.batch_size, self.seq_length)				)




						def     __snake_case   ( self				)     ->     Optional[int]:
							lowerCAmelCase   =     self.prepare_config_and_inputs()
							(
							    (
							    lowerCAmelCase
							),       (
							    lowerCAmelCase
							),       (
							    lowerCAmelCase
							),       (
							    lowerCAmelCase
							),       (
							    lowerCAmelCase
							),       (
							    lowerCAmelCase
							),       (
							    lowerCAmelCase
							),       
							)   =     config_and_inputs
							lowerCAmelCase   =     {
							    """input_ids""": input_ids,
							    """bbox""": bbox,
							    """token_type_ids""": token_type_ids,
							    """attention_mask""": input_mask,
							}
							return config, inputs_dict







@require_torch
class __snake_case(				_lowerCAmelCase		,							_lowerCAmelCase		,							_lowerCAmelCase		,							unittest.TestCase ):


						'''simple docstring'''



						UpperCAmelCase							:   int									=		(
						    (
						        LiltModel,
						        LiltForSequenceClassification,
						        LiltForTokenClassification,
						        LiltForQuestionAnswering,
						    )
						    if is_torch_available()
						    else ()
						)
						UpperCAmelCase							:   Union[str, Any]									=		(
						    {
						        "feature-extraction": LiltModel,
						        "question-answering": LiltForQuestionAnswering,
						        "text-classification": LiltForSequenceClassification,
						        "token-classification": LiltForTokenClassification,
						        "zero-shot": LiltForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						UpperCAmelCase							:   Optional[int]									=		False
						UpperCAmelCase							:   Dict									=		False


						def     __snake_case   ( self     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_				)     ->     Union[str, Any]:
							return True


						def     __snake_case   ( self				)     ->     int:
							lowerCAmelCase   =     LiltModelTester(self				)
							lowerCAmelCase   =     ConfigTester(self     ,	config_class=A_     ,	hidden_size=37				)


						def     __snake_case   ( self				)     ->     Any:
							self.config_tester.run_common_tests()


						def     __snake_case   ( self				)     ->     str:
							lowerCAmelCase   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*A_				)


						def     __snake_case   ( self				)     ->     int:
							lowerCAmelCase   =     self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
								lowerCAmelCase   =     type
								self.model_tester.create_and_check_model(*A_				)


						def     __snake_case   ( self				)     ->     List[Any]:
							lowerCAmelCase   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*A_				)


						def     __snake_case   ( self				)     ->     str:
							lowerCAmelCase   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*A_				)




						@slow
						def     __snake_case   ( self				)     ->     Union[str, Any]:
							for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								lowerCAmelCase   =     LiltModel.from_pretrained(A_				)
								self.assertIsNotNone(A_				)







@require_torch
@slow
class __snake_case(				unittest.TestCase ):


						'''simple docstring'''



						def     __snake_case   ( self				)     ->     Tuple:
							lowerCAmelCase   =     LiltModel.from_pretrained("""SCUT-DLVCLab/lilt-roberta-en-base"""				).to(A_				)

							lowerCAmelCase   =     torch.tensor([[1, 2]]     ,	device=A_				)
							lowerCAmelCase   =     torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]     ,	device=A_				)

							# forward pass
							with torch.no_grad():
								lowerCAmelCase   =     model(input_ids=A_     ,	bbox=A_				)

							lowerCAmelCase   =     torch.Size([1, 2, 768]				)
							lowerCAmelCase   =     torch.tensor(
							    [[-0.0_6_5_3, 0.0_9_5_0, -0.0_0_6_1], [-0.0_5_4_5, 0.0_9_2_6, -0.0_3_2_4]]     ,	device=A_     ,	)

							self.assertTrue(outputs.last_hidden_state.shape     ,	A_				)
							self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :, :3]     ,	A_     ,	atol=1e-3				)				)