
'''simple docstring'''







from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin


@dataclass
class 						__UpperCAmelCase      (					_lowerCamelCase						):
							__lowercase      		=       42


class 						__UpperCAmelCase      (					_lowerCamelCase						,   _lowerCamelCase						):


							@register_to_config
							def __init__(       self    ,   lowerCAmelCase_ = 16    ,   lowerCAmelCase_ = 88    ,   lowerCAmelCase_ = None    ,   lowerCAmelCase_ = None    ,   lowerCAmelCase_ = 1    ,   lowerCAmelCase_ = 0.0    ,   lowerCAmelCase_ = 32    ,   lowerCAmelCase_ = None    ,   lowerCAmelCase_ = False    ,   lowerCAmelCase_ = None    ,   lowerCAmelCase_ = "geglu"    ,   lowerCAmelCase_ = True    ,   lowerCAmelCase_ = True    ,   ):





								"""simple docstring"""






								super().__init__()
								_snake_case   						=		num_attention_heads
								_snake_case   						=		attention_head_dim
								_snake_case   						=		num_attention_heads * attention_head_dim

								_snake_case   						=		in_channels

								_snake_case   						=		torch.nn.GroupNorm(num_groups=lowerCAmelCase_    ,   num_channels=lowerCAmelCase_    ,   eps=1E-6    ,   affine=lowerCAmelCase_  )
								_snake_case   						=		nn.Linear(lowerCAmelCase_    ,   lowerCAmelCase_  )

								# 3. Define transformers blocks
								_snake_case   						=		nn.ModuleList(
								    [
								        BasicTransformerBlock(
								            lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   dropout=lowerCAmelCase_    ,   cross_attention_dim=lowerCAmelCase_    ,   activation_fn=lowerCAmelCase_    ,   attention_bias=lowerCAmelCase_    ,   double_self_attention=lowerCAmelCase_    ,   norm_elementwise_affine=lowerCAmelCase_    ,   )
								        for d in range(lowerCAmelCase_  )
								    ]  )

								_snake_case   						=		nn.Linear(lowerCAmelCase_    ,   lowerCAmelCase_  )




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_=None    ,   lowerCAmelCase_=None    ,   lowerCAmelCase_=None    ,   lowerCAmelCase_=1    ,   lowerCAmelCase_=None    ,   lowerCAmelCase_ = True    ,   ):





								"""simple docstring"""






								_snake_case						, _snake_case						, _snake_case						, _snake_case   						=		hidden_states.shape
								_snake_case   						=		batch_frames // num_frames

								_snake_case   						=		hidden_states

								_snake_case   						=		hidden_states[None, :].reshape(lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_  )
								_snake_case   						=		hidden_states.permute(0    ,   2    ,   1    ,   3    ,   4  )

								_snake_case   						=		self.norm(lowerCAmelCase_  )
								_snake_case   						=		hidden_states.permute(0    ,   3    ,   4    ,   2    ,   1  ).reshape(batch_size * height * width    ,   lowerCAmelCase_    ,   lowerCAmelCase_  )

								_snake_case   						=		self.proj_in(lowerCAmelCase_  )

								# 2. Blocks
								for block in self.transformer_blocks:
									_snake_case   						=		block(
									    lowerCAmelCase_    ,   encoder_hidden_states=lowerCAmelCase_    ,   timestep=lowerCAmelCase_    ,   cross_attention_kwargs=lowerCAmelCase_    ,   class_labels=lowerCAmelCase_    ,   )

								# 3. Output
								_snake_case   						=		self.proj_out(lowerCAmelCase_  )
								_snake_case   						=		(
								    hidden_states[None, None, :]
								    .reshape(lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_  )
								    .permute(0    ,   3    ,   4    ,   1    ,   2  )
								    .contiguous()
								)
								_snake_case   						=		hidden_states.reshape(lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_  )

								_snake_case   						=		hidden_states + residual

								if not return_dict:
									return (output,)

								return TransformerTemporalModelOutput(sample=lowerCAmelCase_  )








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.02       ,							A_=3       ,							A_=None       ,							A_=1000       ,							):
       '''simple docstring'''




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





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




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

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

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

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

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

       UpperCamelCase	:					List[Any]       	=   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'''




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




       UpperCamelCase	:					Any       	=   self.num_labels
       UpperCamelCase	:					Dict       	=   LiltForTokenClassification(config=A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					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        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   LiltForQuestionAnswering(config=A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					List[str]       	=   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'''




       UpperCamelCase	:					Any       	=   self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    
       )	:					Tuple       	=   config_and_inputs
       UpperCamelCase	:					Tuple       	=   {
           "input_ids": input_ids,
           "bbox": bbox,
           "token_type_ids": token_type_ids,
           "attention_mask": input_mask,
       }
       return config, inputs_dict



@require_torch
class 					A__   (       __snake_case       ,   __snake_case       ,   __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:Union[str, Any]						    =   (
          (
              LiltModel,
              LiltForSequenceClassification,
              LiltForTokenClassification,
              LiltForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      _UpperCAmelCase	:Optional[Any]						    =   (
          {
              'feature-extraction': LiltModel,
              'question-answering': LiltForQuestionAnswering,
              'text-classification': LiltForSequenceClassification,
              'token-classification': LiltForTokenClassification,
              'zero-shot': LiltForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _UpperCAmelCase	:Dict						    =   False
      _UpperCAmelCase	:Union[str, Any]						    =   False





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       return True





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   LiltModelTester(self   )
       UpperCamelCase	:					Optional[int]       	=   ConfigTester(self       ,							config_class=A_       ,							hidden_size=37   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




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





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   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]:
        UpperCamelCase	:					Dict       	=   LiltModel.from_pretrained(A_   )
        self.assertIsNotNone(A_   )



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





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   LiltModel.from_pretrained("SCUT-DLVCLab/lilt-roberta-en-base"   ).to(A_   )

       UpperCamelCase	:					Tuple       	=   torch.tensor([[1, 2]]       ,							device=A_   )
       UpperCamelCase	:					List[str]       	=   torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]       ,							device=A_   )

       # forward pass
       with torch.no_grad():
        UpperCamelCase	:					Optional[int]       	=   model(input_ids=A_       ,							bbox=A_   )

       UpperCamelCase	:					List[str]       	=   torch.Size([1, 2, 768]   )
       UpperCamelCase	:					Any       	=   torch.tensor(
           [[-0.06_53, 0.09_50, -0.00_61], [-0.05_45, 0.09_26, -0.03_24]]       ,							device=A_       ,							)

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