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






			def __init__(self  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	Any      ,SCREAMING_SNAKE_CASE_  :	Any=13      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=7      ,SCREAMING_SNAKE_CASE_  :	str=True      ,SCREAMING_SNAKE_CASE_  :	Any=True      ,SCREAMING_SNAKE_CASE_  :	List[Any]=True      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=True      ,SCREAMING_SNAKE_CASE_  :	Any=99      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]=24      ,SCREAMING_SNAKE_CASE_  :	int=2      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=6      ,SCREAMING_SNAKE_CASE_  :	List[Any]=37      ,SCREAMING_SNAKE_CASE_  :	str="gelu"      ,SCREAMING_SNAKE_CASE_  :	List[str]=0.1      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=0.1      ,SCREAMING_SNAKE_CASE_  :	List[str]=512      ,SCREAMING_SNAKE_CASE_  :	Any=16      ,SCREAMING_SNAKE_CASE_  :	List[str]=2      ,SCREAMING_SNAKE_CASE_  :	int=0.02      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]=3      ,SCREAMING_SNAKE_CASE_  :	Dict=None      ,SCREAMING_SNAKE_CASE_  :	Dict=1_000      ,)       ->	List[Any]:





				"""simple docstring"""


				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 					UpperCAmelCase							(self  :	Optional[int]			)       ->	Dict:





				"""simple docstring"""


				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 					UpperCAmelCase							(self  :	List[str]			)       ->	Optional[Any]:





				"""simple docstring"""


				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 					UpperCAmelCase							(self  :	Dict      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	int      ,)       ->	str:





				"""simple docstring"""


				lowerCAmelCase												=  LiltModel(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,bbox=SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,token_type_ids=SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,bbox=SCREAMING_SNAKE_CASE_      ,token_type_ids=SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,bbox=SCREAMING_SNAKE_CASE_			)

				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 					UpperCAmelCase							(self  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	int      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,)       ->	str:





				"""simple docstring"""


				lowerCAmelCase												=  self.num_labels
				lowerCAmelCase												=  LiltForTokenClassification(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(
				    SCREAMING_SNAKE_CASE_      ,bbox=SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,token_type_ids=SCREAMING_SNAKE_CASE_      ,labels=SCREAMING_SNAKE_CASE_			)
				self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.num_labels)			)






			def 					UpperCAmelCase							(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,)       ->	Tuple:





				"""simple docstring"""


				lowerCAmelCase												=  LiltForQuestionAnswering(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(
				    SCREAMING_SNAKE_CASE_      ,bbox=SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,token_type_ids=SCREAMING_SNAKE_CASE_      ,start_positions=SCREAMING_SNAKE_CASE_      ,end_positions=SCREAMING_SNAKE_CASE_      ,)
				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 					UpperCAmelCase							(self  :	int			)       ->	Optional[Any]:





				"""simple docstring"""


				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 			lowercase					( lowercase__   ,lowercase__   ,lowercase__   ,unittest.TestCase		):
			lowercase        =     (
			    (
			        LiltModel,
			        LiltForSequenceClassification,
			        LiltForTokenClassification,
			        LiltForQuestionAnswering,
			    )
			    if is_torch_available()
			    else ()
			)
			lowercase        =     (
			    {
			        '''feature-extraction''': LiltModel,
			        '''question-answering''': LiltForQuestionAnswering,
			        '''text-classification''': LiltForSequenceClassification,
			        '''token-classification''': LiltForTokenClassification,
			        '''zero-shot''': LiltForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowercase        =     False
			lowercase        =     False






			def 					UpperCAmelCase							(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	Optional[int]			)       ->	List[str]:





				"""simple docstring"""


				return True






			def 					UpperCAmelCase							(self  :	Optional[Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				lowerCAmelCase												=  LiltModelTester(self			)
				lowerCAmelCase												=  ConfigTester(self      ,config_class=SCREAMING_SNAKE_CASE_      ,hidden_size=37			)






			def 					UpperCAmelCase							(self  :	int			)       ->	Dict:





				"""simple docstring"""


				self.config_tester.run_common_tests()






			def 					UpperCAmelCase							(self  :	Dict			)       ->	Optional[int]:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	List[str]			)       ->	Dict:





				"""simple docstring"""


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






			def 					UpperCAmelCase							(self  :	List[str]			)       ->	int:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	Tuple			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_			)






			@slow
			def 					UpperCAmelCase							(self  :	List[Any]			)       ->	Dict:





				"""simple docstring"""


				for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCAmelCase												=  LiltModel.from_pretrained(SCREAMING_SNAKE_CASE_			)
					self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)





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






			def 					UpperCAmelCase							(self  :	Optional[Any]			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  LiltModel.from_pretrained('''SCUT-DLVCLab/lilt-roberta-en-base'''			).to(SCREAMING_SNAKE_CASE_			)

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

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

				lowerCAmelCase												=  torch.Size([1, 2, 768]			)
				lowerCAmelCase												=  torch.tensor(
				    [[-0.06_53, 0.09_50, -0.00_61], [-0.05_45, 0.09_26, -0.03_24]]      ,device=SCREAMING_SNAKE_CASE_      ,)

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




"""simple docstring"""






from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


UpperCAmelCase        =				logging.get_logger(__name__)


if is_torch_available():
					import torch
					from torch import nn
					from torch.utils.data import Dataset

					from transformers import Trainer



					class 			lowercase					( lowercase__		):






								def __init__(self  :	int      ,SCREAMING_SNAKE_CASE_  :	int = 101			)       ->	List[str]:





									"""simple docstring"""


									lowerCAmelCase												=  length






								def __len__(self  :	Any			)       ->	List[Any]:





									"""simple docstring"""


									return self.length






								def __getitem__(self  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	int:





									"""simple docstring"""


									return i





					class 			lowercase					:






								def __call__(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[Any]			)       ->	Optional[Any]:





									"""simple docstring"""


									return {"input_ids": torch.tensor(SCREAMING_SNAKE_CASE_			), "labels": torch.tensor(SCREAMING_SNAKE_CASE_			)}





					class 			lowercase					( nn.Module		):






								def __init__(self  :	List[str]			)       ->	List[Any]:





									"""simple docstring"""


									super().__init__()
									# Add some (unused) params otherwise DDP will complain.
									lowerCAmelCase												=  nn.Linear(120      ,80			)






								def 					UpperCAmelCase							(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=None			)       ->	Union[str, Any]:





									"""simple docstring"""


									if labels is not None:
										return torch.tensor(0.0      ,device=input_ids.device			), input_ids
									else:
										return input_ids



class 			lowercase					( lowercase__		):






			@require_torch_neuroncore
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call





class 			lowercase					( lowercase__		):






			@require_torch_multi_gpu
			def 					UpperCAmelCase							(self  :	List[str]			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
					# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
					#
					# PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

					UpperCAmelCase        =				HfArgumentParser((TrainingArguments,))
					UpperCAmelCase        =				parser.parse_args_into_dataclasses()[0]

					logger.warning(
					    f"""Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, """
					    f"""distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}"""
					)

					# Essentially, what we want to verify in the distributed case is that we get all samples back,
					# in the right order. (this is crucial for prediction for instance)
					for dataset_length in [1_0_1, 4_0, 7]:
										UpperCAmelCase        =				DummyDataset(dataset_length)




										def __magic_name__	(				_lowerCamelCase:		EvalPrediction			)					->						Dict:

											'''simple docstring'''




											lowerCAmelCase												=  list(range(len(_lowerCamelCase			)			)			)
											lowerCAmelCase												=  p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
											if not success and training_args.local_rank == 0:
												logger.warning(
												    '''Predictions and/or labels do not match expected results:\n  - predictions: '''
												    F"""{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"""			)
											return {"success": success}

										UpperCAmelCase        =				Trainer(
										    model=DummyModel(),
										    args=training_args,
										    data_collator=DummyDataCollator(),
										    eval_dataset=dataset,
										    compute_metrics=compute_metrics,
										)
										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)

										UpperCAmelCase        =				2

										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)



										UpperCAmelCase        =				None


