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






     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    ,					):
       _lowerCAmelCase						:	List[str]        =					parent
       _lowerCAmelCase						:	Union[str, Any]        =					batch_size
       _lowerCAmelCase						:	List[str]        =					seq_length
       _lowerCAmelCase						:	List[str]        =					is_training
       _lowerCAmelCase						:	Tuple        =					use_input_mask
       _lowerCAmelCase						:	List[Any]        =					use_token_type_ids
       _lowerCAmelCase						:	Tuple        =					use_labels
       _lowerCAmelCase						:	int        =					vocab_size
       _lowerCAmelCase						:	Optional[Any]        =					hidden_size
       _lowerCAmelCase						:	List[Any]        =					num_hidden_layers
       _lowerCAmelCase						:	int        =					num_attention_heads
       _lowerCAmelCase						:	Any        =					intermediate_size
       _lowerCAmelCase						:	int        =					hidden_act
       _lowerCAmelCase						:	List[str]        =					hidden_dropout_prob
       _lowerCAmelCase						:	Dict        =					attention_probs_dropout_prob
       _lowerCAmelCase						:	List[str]        =					max_position_embeddings
       _lowerCAmelCase						:	Optional[Any]        =					type_vocab_size
       _lowerCAmelCase						:	Optional[int]        =					type_sequence_label_size
       _lowerCAmelCase						:	List[Any]        =					initializer_range
       _lowerCAmelCase						:	Optional[Any]        =					num_labels
       _lowerCAmelCase						:	int        =					scope
       _lowerCAmelCase						:	Tuple        =					range_bbox






     def    __A  (		self  ):
       _lowerCAmelCase						:	Dict        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.vocab_size  )

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

       _lowerCAmelCase						:	str        =					None
       if self.use_input_mask:
         _lowerCAmelCase						:	Optional[int]        =					ids_tensor([self.batch_size, self.seq_length]    ,					vocab_size=2  )

       _lowerCAmelCase						:	Dict        =					None
       if self.use_token_type_ids:
         _lowerCAmelCase						:	Optional[Any]        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.type_vocab_size  )

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

       _lowerCAmelCase						:	int        =					self.get_config()

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






     def    __A  (		self  ):
       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    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	Optional[Any]        =					LiltModel(config=a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	str        =					model(a__    ,					bbox=a__    ,					attention_mask=a__    ,					token_type_ids=a__  )
       _lowerCAmelCase						:	Optional[Any]        =					model(a__    ,					bbox=a__    ,					token_type_ids=a__  )
       _lowerCAmelCase						:	int        =					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    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	List[Any]        =					self.num_labels
       _lowerCAmelCase						:	List[str]        =					LiltForTokenClassification(config=a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	Dict        =					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    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	List[Any]        =					LiltForQuestionAnswering(config=a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	Any        =					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    __A  (		self  ):
       _lowerCAmelCase						:	Optional[int]        =					self.prepare_config_and_inputs()
       (
           (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     
       )						:	Optional[Any]        =					config_and_inputs
       _lowerCAmelCase						:	Dict        =					{
           """input_ids""": input_ids,
           """bbox""": bbox,
           """token_type_ids""": token_type_ids,
           """attention_mask""": input_mask,
       }
       return config, inputs_dict







@require_torch
class 						__A							(			SCREAMING_SNAKE_CASE_   ,     SCREAMING_SNAKE_CASE_   ,     SCREAMING_SNAKE_CASE_   ,     unittest.TestCase		):
     _UpperCamelCase							:      Tuple		 =		(
         (
             LiltModel,
             LiltForSequenceClassification,
             LiltForTokenClassification,
             LiltForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     _UpperCamelCase							:      Optional[int]		 =		(
         {
             "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							:      List[str]		 =		False






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__  ):
       return True






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[int]        =					LiltModelTester(self  )
       _lowerCAmelCase						:	List[Any]        =					ConfigTester(self    ,					config_class=a__    ,					hidden_size=37  )






     def    __A  (		self  ):
       self.config_tester.run_common_tests()






     def    __A  (		self  ):
       _lowerCAmelCase						:	List[str]        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	str        =					self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
         _lowerCAmelCase						:	Optional[Any]        =					type
         self.model_tester.create_and_check_model(*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	List[Any]        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[Any]        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*a__  )






     @slow
     def    __A  (		self  ):
       for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _lowerCAmelCase						:	List[Any]        =					LiltModel.from_pretrained(a__  )
         self.assertIsNotNone(a__  )







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






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[Any]        =					LiltModel.from_pretrained("""SCUT-DLVCLab/lilt-roberta-en-base"""  ).to(a__  )

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

       # forward pass
       with torch.no_grad():
         _lowerCAmelCase						:	Any        =					model(input_ids=a__    ,					bbox=a__  )

       _lowerCAmelCase						:	Dict        =					torch.Size([1, 2, 768]  )
       _lowerCAmelCase						:	Any        =					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  )  )





"""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 						__A							(			SCREAMING_SNAKE_CASE_		):
     _UpperCamelCase							:      Optional[torch.FloatTensor]		 =		None
     _UpperCamelCase							:      torch.FloatTensor		 =		None
     _UpperCamelCase							:      Optional[Tuple[torch.FloatTensor]]		 =		None
     _UpperCamelCase							:      Optional[Tuple[torch.FloatTensor]]		 =		None







class 						__A							(			SCREAMING_SNAKE_CASE_		):






     def __init__(		self    ,					a__=1    ,					a__=0    ,					a__=2    ,					a__=512    ,					a__="cls"    ,					a__=False    ,					a__=True    ,					**a__    ,					):
       super().__init__(pad_token_id=a__    ,					bos_token_id=a__    ,					eos_token_id=a__    ,					**a__  )
       _lowerCAmelCase						:	Optional[Any]        =					project_dim
       _lowerCAmelCase						:	List[str]        =					pooler_fn
       _lowerCAmelCase						:	Any        =					learn_encoder
       _lowerCAmelCase						:	Optional[int]        =					use_attention_mask







class 						__A							(			SCREAMING_SNAKE_CASE_		):
     _UpperCamelCase							:      Optional[int]		 =		[R"pooler", R"logit_scale"]
     _UpperCamelCase							:      List[Any]		 =		[R"position_ids", R"predictions.decoder.bias"]
     _UpperCamelCase							:      List[Any]		 =		"roberta"
     _UpperCamelCase							:      Optional[int]		 =		RobertaSeriesConfig






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






     def    __A  (		self    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					a__ = None    ,					):
       _lowerCAmelCase						:	List[str]        =					return_dict if return_dict is not None else self.config.use_return_dict

       _lowerCAmelCase						:	Optional[int]        =					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:
         _lowerCAmelCase						:	Optional[Any]        =					outputs["""hidden_states"""][-2]
         _lowerCAmelCase						:	Optional[Any]        =					self.pre_LN(a__  )
         _lowerCAmelCase						:	int        =					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:
         _lowerCAmelCase						:	Union[str, 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    ,					)




