


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






	"""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.02						,    __a=3						,    __a=None						,    __a=1000						,    ):

						'''simple docstring'''





						__a      :	str    =  parent
						__a      :	Optional[int]    =  batch_size
						__a      :	List[Any]    =  seq_length
						__a      :	Dict    =  is_training
						__a      :	Any    =  use_input_mask
						__a      :	Optional[int]    =  use_token_type_ids
						__a      :	Optional[int]    =  use_labels
						__a      :	List[str]    =  vocab_size
						__a      :	Dict    =  hidden_size
						__a      :	Union[str, Any]    =  num_hidden_layers
						__a      :	List[str]    =  num_attention_heads
						__a      :	Any    =  intermediate_size
						__a      :	Any    =  hidden_act
						__a      :	Optional[Any]    =  hidden_dropout_prob
						__a      :	int    =  attention_probs_dropout_prob
						__a      :	Dict    =  max_position_embeddings
						__a      :	str    =  type_vocab_size
						__a      :	Dict    =  type_sequence_label_size
						__a      :	str    =  initializer_range
						__a      :	List[Any]    =  num_labels
						__a      :	str    =  scope
						__a      :	Tuple    =  range_bbox




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Any    =  ids_tensor([self.batch_size, self.seq_length]						,    self.vocab_size							)

						__a      :	int    =  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]:
																					__a      :	Optional[int]    =  bbox[i, j, 3]
																					__a      :	Optional[int]    =  bbox[i, j, 1]
																					__a      :	List[str]    =  t
																if bbox[i, j, 2] < bbox[i, j, 0]:
																					__a      :	Optional[int]    =  bbox[i, j, 2]
																					__a      :	Optional[Any]    =  bbox[i, j, 0]
																					__a      :	int    =  t

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

						__a      :	List[Any]    =  None
						if self.use_token_type_ids:
											__a      :	Union[str, Any]    =  ids_tensor([self.batch_size, self.seq_length]						,    self.type_vocab_size							)

						__a      :	Tuple    =  None
						__a      :	int    =  None
						if self.use_labels:
											__a      :	Any    =  ids_tensor([self.batch_size]						,    self.type_sequence_label_size							)
											__a      :	Any    =  ids_tensor([self.batch_size, self.seq_length]						,    self.num_labels							)

						__a      :	str    =  self.get_config()

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




	def 				__UpperCAmelCase							(       self							):

						'''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						,    __a						,    __a						,    __a						,    __a						,    __a						,    __a						,    __a						,    ):

						'''simple docstring'''





						__a      :	Optional[int]    =  LiltModel(config=__a							)
						model.to(__a							)
						model.eval()
						__a      :	Any    =  model(__a						,    bbox=__a						,    attention_mask=__a						,    token_type_ids=__a							)
						__a      :	List[str]    =  model(__a						,    bbox=__a						,    token_type_ids=__a							)
						__a      :	Any    =  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 				__UpperCAmelCase							(       self						,    __a						,    __a						,    __a						,    __a						,    __a						,    __a						,    __a						,    ):

						'''simple docstring'''





						__a      :	Optional[int]    =  self.num_labels
						__a      :	str    =  LiltForTokenClassification(config=__a							)
						model.to(__a							)
						model.eval()
						__a      :	List[Any]    =  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 				__UpperCAmelCase							(       self						,    __a						,    __a						,    __a						,    __a						,    __a						,    __a						,    __a						,    ):

						'''simple docstring'''





						__a      :	str    =  LiltForQuestionAnswering(config=__a							)
						model.to(__a							)
						model.eval()
						__a      :	int    =  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 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Optional[int]    =  self.prepare_config_and_inputs()
						(
						    __a    
						)      :	Tuple    =  config_and_inputs
						__a      :	int    =  {
						    'input_ids': input_ids,
						    'bbox': bbox,
						    'token_type_ids': token_type_ids,
						    'attention_mask': input_mask,
						}
						return config, inputs_dict





@require_torch
class 							__UpperCamelCase    (		lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				unittest.TestCase ):






	"""simple docstring"""




	A_           =   (
	    (
	        LiltModel,
	        LiltForSequenceClassification,
	        LiltForTokenClassification,
	        LiltForQuestionAnswering,
	    )
	    if is_torch_available()
	    else ()
	)
	A_           =   (
	    {
	        "feature-extraction": LiltModel,
	        "question-answering": LiltForQuestionAnswering,
	        "text-classification": LiltForSequenceClassification,
	        "token-classification": LiltForTokenClassification,
	        "zero-shot": LiltForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	A_           =   False
	A_           =   False




	def 				__UpperCAmelCase							(       self						,    __a						,    __a						,    __a						,    __a						,    __a							):

						'''simple docstring'''





						return True




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Dict    =  LiltModelTester(self							)
						__a      :	Union[str, Any]    =  ConfigTester(self						,    config_class=__a						,    hidden_size=37							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						self.config_tester.run_common_tests()




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Optional[int]    =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__a							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Optional[int]    =  self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
											__a      :	Dict    =  type
											self.model_tester.create_and_check_model(*__a							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Optional[Any]    =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*__a							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Dict    =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*__a							)





	@slow
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__a      :	Optional[int]    =  LiltModel.from_pretrained(__a							)
											self.assertIsNotNone(__a							)





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






	"""simple docstring"""




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Union[str, Any]    =  LiltModel.from_pretrained('SCUT-DLVCLab/lilt-roberta-en-base'							).to(__a							)

						__a      :	List[Any]    =  torch.tensor([[1, 2]]						,    device=__a							)
						__a      :	List[Any]    =  torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]						,    device=__a							)

						# forward pass
						with torch.no_grad():
											__a      :	Union[str, Any]    =  model(input_ids=__a						,    bbox=__a							)

						__a      :	Union[str, Any]    =  torch.Size([1, 2, 768]							)
						__a      :	Dict    =  torch.tensor(
						    [[-0.0653, 0.0950, -0.0061], [-0.0545, 0.0926, -0.0324]]						,    device=__a						,    )

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







'''simple docstring'''



from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput





@dataclass
class 							__UpperCamelCase    (		lowerCAmelCase_ ):
	A_           =   None
	A_           =   None
	A_           =   None
	A_           =   None





class 							__UpperCamelCase    (		lowerCAmelCase_ ):




	def __init__(       self						,    __a=1						,    __a=0						,    __a=2						,    __a=512						,    __a="cls"						,    __a=False						,    __a=True						,    **__a						,    ):

						'''simple docstring'''





						super().__init__(pad_token_id=__a						,    bos_token_id=__a						,    eos_token_id=__a						,    **__a							)
						__a      :	Any    =  project_dim
						__a      :	Optional[Any]    =  pooler_fn
						__a      :	int    =  learn_encoder
						__a      :	str    =  use_attention_mask





class 							__UpperCamelCase    (		lowerCAmelCase_ ):
	A_           =   [r"pooler", r"logit_scale"]
	A_           =   [r"position_ids", r"predictions.decoder.bias"]
	A_           =   "roberta"
	A_           =   RobertaSeriesConfig




	def __init__(       self						,    __a							):

						'''simple docstring'''





						super().__init__(__a							)
						__a      :	Optional[Any]    =  XLMRobertaModel(__a							)
						__a      :	str    =  nn.Linear(config.hidden_size						,    config.project_dim							)
						__a      :	Optional[int]    =  getattr(__a						,    'has_pre_transformation'						,    __a							)
						if self.has_pre_transformation:
											__a      :	int    =  nn.Linear(config.hidden_size						,    config.project_dim							)
											__a      :	List[str]    =  nn.LayerNorm(config.hidden_size						,    eps=config.layer_norm_eps							)
						self.post_init()





	def 				__UpperCAmelCase							(       self						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    ):

						'''simple docstring'''





						__a      :	Optional[Any]    =  return_dict if return_dict is not None else self.config.use_return_dict

						__a      :	Tuple    =  self.base_model(
						    input_ids=__a						,    attention_mask=__a						,    token_type_ids=__a						,    position_ids=__a						,    head_mask=__a						,    inputs_embeds=__a						,    encoder_hidden_states=__a						,    encoder_attention_mask=__a						,    output_attentions=__a						,    output_hidden_states=True if self.has_pre_transformation else output_hidden_states						,    return_dict=__a						,    )

						if self.has_pre_transformation:
											__a      :	Optional[Any]    =  outputs['hidden_states'][-2]
											__a      :	Optional[int]    =  self.pre_LN(__a							)
											__a      :	Union[str, Any]    =  self.transformation_pre(__a							)

											return TransformationModelOutput(
											    projection_state=__a						,    last_hidden_state=outputs.last_hidden_state						,    hidden_states=outputs.hidden_states						,    attentions=outputs.attentions						,    )
						else:
											__a      :	Optional[Any]    =  self.transformation(outputs.last_hidden_state							)
											return TransformationModelOutput(
											    projection_state=__a						,    last_hidden_state=outputs.last_hidden_state						,    hidden_states=outputs.hidden_states						,    attentions=outputs.attentions						,    )



