





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 	__lowerCAmelCase      :





   def __init__(			self       ,      snake_case       ,      snake_case=13       ,      snake_case=7       ,      snake_case=True       ,      snake_case=True       ,      snake_case=True       ,      snake_case=True       ,      snake_case=99       ,      snake_case=24       ,      snake_case=2       ,      snake_case=6       ,      snake_case=37       ,      snake_case="gelu"       ,      snake_case=0.1       ,      snake_case=0.1       ,      snake_case=512       ,      snake_case=16       ,      snake_case=2       ,      snake_case=0.02       ,      snake_case=3       ,      snake_case=None       ,      snake_case=1_000       ,      )					->	int:
      """simple docstring"""





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





   def 						_snake_case			(			self   )					->	List[str]:
      """simple docstring"""





      a__						:    int   		=	ids_tensor([self.batch_size, self.seq_length]       ,      self.vocab_size   )

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

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

      a__						:    str   		=	None
      if self.use_token_type_ids:
         a__						:    Optional[int]   		=	ids_tensor([self.batch_size, self.seq_length]       ,      self.type_vocab_size   )

      a__						:    List[str]   		=	None
      a__						:    List[str]   		=	None
      if self.use_labels:
         a__						:    List[Any]   		=	ids_tensor([self.batch_size]       ,      self.type_sequence_label_size   )
         a__						:    Dict   		=	ids_tensor([self.batch_size, self.seq_length]       ,      self.num_labels   )

      a__						:    Tuple   		=	self.get_config()

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





   def 						_snake_case			(			self   )					->	List[str]:
      """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 						_snake_case			(			self       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      )					->	Tuple:
      """simple docstring"""





      a__						:    Tuple   		=	LiltModel(config=snake_case   )
      model.to(snake_case   )
      model.eval()
      a__						:    Union[str, Any]   		=	model(snake_case       ,      bbox=snake_case       ,      attention_mask=snake_case       ,      token_type_ids=snake_case   )
      a__						:    List[Any]   		=	model(snake_case       ,      bbox=snake_case       ,      token_type_ids=snake_case   )
      a__						:    Optional[int]   		=	model(snake_case       ,      bbox=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 						_snake_case			(			self       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      )					->	str:
      """simple docstring"""





      a__						:    str   		=	self.num_labels
      a__						:    Union[str, Any]   		=	LiltForTokenClassification(config=snake_case   )
      model.to(snake_case   )
      model.eval()
      a__						:    Union[str, Any]   		=	model(
          snake_case       ,      bbox=snake_case       ,      attention_mask=snake_case       ,      token_type_ids=snake_case       ,      labels=snake_case   )
      self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.seq_length, self.num_labels)   )





   def 						_snake_case			(			self       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      )					->	Any:
      """simple docstring"""





      a__						:    Union[str, Any]   		=	LiltForQuestionAnswering(config=snake_case   )
      model.to(snake_case   )
      model.eval()
      a__						:    List[str]   		=	model(
          snake_case       ,      bbox=snake_case       ,      attention_mask=snake_case       ,      token_type_ids=snake_case       ,      start_positions=snake_case       ,      end_positions=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 						_snake_case			(			self   )					->	str:
      """simple docstring"""





      a__						:    List[Any]   		=	self.prepare_config_and_inputs()
      (
          (
          a__
      )		,  (
          a__
      )		,  (
          a__
      )		,  (
          a__
      )		,  (
          a__
      )		,  (
          a__
      )		,  (
          a__
      )		,  
      )						:    Union[str, Any]   		=	config_and_inputs
      a__						:    Any   		=	{
          "input_ids": input_ids,
          "bbox": bbox,
          "token_type_ids": token_type_ids,
          "attention_mask": input_mask,
      }
      return config, inputs_dict







@require_torch
class 	__lowerCAmelCase      (       _UpperCamelCase		,_UpperCamelCase		,_UpperCamelCase		,unittest.TestCase	):
   _UpperCamelCase			:    Optional[int]  		=					(
       (
           LiltModel,
           LiltForSequenceClassification,
           LiltForTokenClassification,
           LiltForQuestionAnswering,
       )
       if is_torch_available()
       else ()
   )
   _UpperCamelCase			:    Dict  		=					(
       {
           """feature-extraction""": LiltModel,
           """question-answering""": LiltForQuestionAnswering,
           """text-classification""": LiltForSequenceClassification,
           """token-classification""": LiltForTokenClassification,
           """zero-shot""": LiltForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   _UpperCamelCase			:    Tuple  		=					False
   _UpperCamelCase			:    Optional[Any]  		=					False





   def 						_snake_case			(			self       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case   )					->	Optional[Any]:
      """simple docstring"""





      return True





   def 						_snake_case			(			self   )					->	str:
      """simple docstring"""





      a__						:    Tuple   		=	LiltModelTester(self   )
      a__						:    Optional[Any]   		=	ConfigTester(self       ,      config_class=snake_case       ,      hidden_size=37   )





   def 						_snake_case			(			self   )					->	Dict:
      """simple docstring"""





      self.config_tester.run_common_tests()





   def 						_snake_case			(			self   )					->	Optional[Any]:
      """simple docstring"""





      a__						:    List[str]   		=	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*snake_case   )





   def 						_snake_case			(			self   )					->	Tuple:
      """simple docstring"""





      a__						:    Union[str, Any]   		=	self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
         a__						:    Any   		=	type
         self.model_tester.create_and_check_model(*snake_case   )





   def 						_snake_case			(			self   )					->	Tuple:
      """simple docstring"""





      a__						:    str   		=	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*snake_case   )





   def 						_snake_case			(			self   )					->	List[str]:
      """simple docstring"""





      a__						:    Dict   		=	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*snake_case   )





   @slow
   def 						_snake_case			(			self   )					->	List[str]:
      """simple docstring"""





      for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         a__						:    Union[str, Any]   		=	LiltModel.from_pretrained(snake_case   )
         self.assertIsNotNone(snake_case   )







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





   def 						_snake_case			(			self   )					->	Optional[int]:
      """simple docstring"""





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

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

      # forward pass
      with torch.no_grad():
         a__						:    List[str]   		=	model(input_ids=snake_case       ,      bbox=snake_case   )

      a__						:    Optional[int]   		=	torch.Size([1, 2, 768]   )
      a__						:    Tuple   		=	torch.tensor(
          [[-0.0_653, 0.0_950, -0.0_061], [-0.0_545, 0.0_926, -0.0_324]]       ,      device=snake_case       ,      )

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








from packaging import version

from .import_utils import is_accelerate_available


if is_accelerate_available():
  import accelerate






def    _A	(   lowerCamelCase ):
   if not is_accelerate_available():
      return method
   a__						:    List[Any]   		=	version.parse(accelerate.__version__ ).base_version
   if version.parse(lowerCamelCase ) < version.parse("0.17.0" ):
      return method

   def wrapper(self				,			*lowerCamelCase				,			**lowerCamelCase ):
      if hasattr(self				,			"_hf_hook" ) and hasattr(self._hf_hook				,			"pre_forward" ):
         self._hf_hook.pre_forward(self )
      return method(self				,			*lowerCamelCase				,			**lowerCamelCase )

   return wrapper

