


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



  """simple docstring"""



  def __init__( self			:  Tuple				,				lowerCamelCase__			:  Dict				,				lowerCamelCase__			:  int=13				,				lowerCamelCase__			:  Dict=7				,				lowerCamelCase__			:  Tuple=True				,				lowerCamelCase__			:  List[str]=True				,				lowerCamelCase__			:  str=True				,				lowerCamelCase__			:  Optional[Any]=True				,				lowerCamelCase__			:  Any=99				,				lowerCamelCase__			:  List[str]=24				,				lowerCamelCase__			:  str=2				,				lowerCamelCase__			:  Dict=6				,				lowerCamelCase__			:  List[str]=37				,				lowerCamelCase__			:  Union[str, Any]="gelu"				,				lowerCamelCase__			:  List[Any]=0.1				,				lowerCamelCase__			:  List[str]=0.1				,				lowerCamelCase__			:  Dict=512				,				lowerCamelCase__			:  Optional[Any]=16				,				lowerCamelCase__			:  Optional[int]=2				,				lowerCamelCase__			:  Any=0.02				,				lowerCamelCase__			:  Optional[Any]=3				,				lowerCamelCase__			:  Optional[int]=None				,				lowerCamelCase__			:  str=1_000				,				):
        a__	:  int     =   parent
        a__	:  Optional[int]     =   batch_size
        a__	:  Optional[Any]     =   seq_length
        a__	:  Tuple     =   is_training
        a__	:  Union[str, Any]     =   use_input_mask
        a__	:  str     =   use_token_type_ids
        a__	:  Union[str, Any]     =   use_labels
        a__	:  Tuple     =   vocab_size
        a__	:  Any     =   hidden_size
        a__	:  Tuple     =   num_hidden_layers
        a__	:  Union[str, Any]     =   num_attention_heads
        a__	:  Union[str, Any]     =   intermediate_size
        a__	:  Any     =   hidden_act
        a__	:  Dict     =   hidden_dropout_prob
        a__	:  Optional[Any]     =   attention_probs_dropout_prob
        a__	:  Optional[int]     =   max_position_embeddings
        a__	:  Any     =   type_vocab_size
        a__	:  str     =   type_sequence_label_size
        a__	:  Union[str, Any]     =   initializer_range
        a__	:  Optional[Any]     =   num_labels
        a__	:  List[Any]     =   scope
        a__	:  List[Any]     =   range_bbox

  def    _UpperCamelCase( self			:  str       ):
        a__	:  Union[str, Any]     =   ids_tensor([self.batch_size, self.seq_length]				,				self.vocab_size       )

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

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

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

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

        a__	:  List[Any]     =   self.get_config()

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

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

        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			:  List[str]				,				lowerCamelCase__			:  Tuple				,				lowerCamelCase__			:  Optional[Any]				,				lowerCamelCase__			:  List[Any]				,				lowerCamelCase__			:  Dict				,				lowerCamelCase__			:  Dict				,				lowerCamelCase__			:  List[str]				,				lowerCamelCase__			:  Any				,				):
        a__	:  Any     =   self.num_labels
        a__	:  Dict     =   LiltForTokenClassification(config=lowerCamelCase__       )
        model.to(lowerCamelCase__       )
        model.eval()
        a__	:  Any     =   model(
            lowerCamelCase__				,				bbox=lowerCamelCase__				,				attention_mask=lowerCamelCase__				,				token_type_ids=lowerCamelCase__				,				labels=lowerCamelCase__       )
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length, self.num_labels)       )

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




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



  """simple docstring"""



  _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			:  str				,				lowerCamelCase__			:  int				,				lowerCamelCase__			:  Dict				,				lowerCamelCase__			:  str				,				lowerCamelCase__			:  str				,				lowerCamelCase__			:  Tuple       ):
        return True

  def    _UpperCamelCase( self			:  Dict       ):
        a__	:  Any     =   LiltModelTester(self       )
        a__	:  Tuple     =   ConfigTester(self				,				config_class=lowerCamelCase__				,				hidden_size=37       )

  def    _UpperCamelCase( self			:  Any       ):
        self.config_tester.run_common_tests()

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

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

  def    _UpperCamelCase( self			:  Tuple       ):
        a__	:  Dict     =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__       )

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

  @slow
  def    _UpperCamelCase( self			:  Any       ):
        for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              a__	:  Optional[Any]     =   LiltModel.from_pretrained(lowerCamelCase__       )
              self.assertIsNotNone(lowerCamelCase__       )




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



  """simple docstring"""



  def    _UpperCamelCase( self			:  int       ):
        a__	:  int     =   LiltModel.from_pretrained("SCUT-DLVCLab/lilt-roberta-en-base"       ).to(lowerCamelCase__       )

        a__	:  Union[str, Any]     =   torch.tensor([[1, 2]]				,				device=lowerCamelCase__       )
        a__	:  Optional[Any]     =   torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]				,				device=lowerCamelCase__       )

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

        a__	:  Any     =   torch.Size([1, 2, 768]       )
        a__	:  Any     =   torch.tensor(
            [[-0.0653, 0.0950, -0.0061], [-0.0545, 0.0926, -0.0324]]				,				device=lowerCamelCase__				,				)

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





'''simple docstring'''


import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


_UpperCamelCase 							=  logging.getLogger(__name__)







class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase=None			,     lowerCamelCase=None		):


          '''simple docstring'''





          __A	:       List[Any]      =							self.layer[current_layer](lowerCamelCase			,     lowerCamelCase			,     head_mask[current_layer]		)

          __A	:       Tuple      =							layer_outputs[0]

          return hidden_states







@add_start_docstrings(
    'The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top.'    ,		lowerCAmelCase    ,		)
class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__(lowerCamelCase		)

          __A	:       Any      =							BertEncoderWithPabee(lowerCamelCase		)

          self.init_weights()
          __A	:       Optional[Any]      =							0
          __A	:       Union[str, Any]      =							0
          __A	:       Dict      =							0

          __A	:       List[Any]      =							0


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Any      =							threshold


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       List[Any]      =							patience


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       List[str]      =							0
          __A	:       int      =							0


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[Any]      =							self.inference_layers_num / self.inference_instances_num
          __A	:       int      =							(
              f"*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ="
              f" {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***"
          )
          print(lowerCamelCase		)


     @add_start_docstrings_to_model_forward(lowerCamelCase		)
     def lowerCAmelCase__  (   self			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=False			,     ):


          '''simple docstring'''





          if input_ids is not None and inputs_embeds is not None:
               raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time"		)
          elif input_ids is not None:
               __A	:       Optional[int]      =							input_ids.size()
          elif inputs_embeds is not None:
               __A	:       Any      =							inputs_embeds.size()[:-1]
          else:
               raise ValueError("You have to specify either input_ids or inputs_embeds"		)

          __A	:       List[str]      =							input_ids.device if input_ids is not None else inputs_embeds.device

          if attention_mask is None:
               __A	:       Any      =							torch.ones(lowerCamelCase			,     device=lowerCamelCase		)
          if token_type_ids is None:
               __A	:       List[Any]      =							torch.zeros(lowerCamelCase			,     dtype=torch.long			,     device=lowerCamelCase		)

          # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
          # ourselves in which case we just need to make it broadcastable to all heads.
          __A	:       torch.Tensor      =							self.get_extended_attention_mask(lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		)

          # If a 2D ou 3D attention mask is provided for the cross-attention
          # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
          if self.config.is_decoder and encoder_hidden_states is not None:
               __A  ,__A  ,__A	:       Tuple      =							encoder_hidden_states.size()
               __A	:       Any      =							(encoder_batch_size, encoder_sequence_length)
               if encoder_attention_mask is None:
                    __A	:       List[str]      =							torch.ones(lowerCamelCase			,     device=lowerCamelCase		)
               __A	:       Dict      =							self.invert_attention_mask(lowerCamelCase		)
          else:
               __A	:       Union[str, Any]      =							None

          # Prepare head mask if needed
          # 1.0 in head_mask indicate we keep the head
          # attention_probs has shape bsz x n_heads x N x N
          # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
          # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
          __A	:       Any      =							self.get_head_mask(lowerCamelCase			,     self.config.num_hidden_layers		)

          __A	:       Tuple      =							self.embeddings(
              input_ids=lowerCamelCase			,     position_ids=lowerCamelCase			,     token_type_ids=lowerCamelCase			,     inputs_embeds=lowerCamelCase		)
          __A	:       str      =							embedding_output

          if self.training:
               __A	:       Dict      =							[]
               for i in range(self.config.num_hidden_layers		):
                    __A	:       str      =							self.encoder.adaptive_forward(
                        lowerCamelCase			,     current_layer=lowerCamelCase			,     attention_mask=lowerCamelCase			,     head_mask=lowerCamelCase		)

                    __A	:       Optional[Any]      =							self.pooler(lowerCamelCase		)
                    __A	:       Optional[Any]      =							output_layers[i](output_dropout(lowerCamelCase		)		)
                    res.append(lowerCamelCase		)
          elif self.patience == 0:  # Use all layers for inference
               __A	:       int      =							self.encoder(
                   lowerCamelCase			,     attention_mask=lowerCamelCase			,     head_mask=lowerCamelCase			,     encoder_hidden_states=lowerCamelCase			,     encoder_attention_mask=lowerCamelCase			,     )
               __A	:       Optional[Any]      =							self.pooler(encoder_outputs[0]		)
               __A	:       List[str]      =							[output_layers[self.config.num_hidden_layers - 1](lowerCamelCase		)]
          else:
               __A	:       int      =							0
               __A	:       List[str]      =							None
               __A	:       Optional[int]      =							0
               for i in range(self.config.num_hidden_layers		):
                    calculated_layer_num += 1
                    __A	:       List[str]      =							self.encoder.adaptive_forward(
                        lowerCamelCase			,     current_layer=lowerCamelCase			,     attention_mask=lowerCamelCase			,     head_mask=lowerCamelCase		)

                    __A	:       List[str]      =							self.pooler(lowerCamelCase		)
                    __A	:       str      =							output_layers[i](lowerCamelCase		)
                    if regression:
                         __A	:       int      =							logits.detach()
                         if patient_result is not None:
                              __A	:       Union[str, Any]      =							patient_result.detach()
                         if (patient_result is not None) and torch.abs(patient_result - labels		) < self.regression_threshold:
                              patient_counter += 1
                         else:
                              __A	:       Any      =							0
                    else:
                         __A	:       str      =							logits.detach().argmax(dim=1		)
                         if patient_result is not None:
                              __A	:       List[str]      =							patient_result.detach().argmax(dim=1		)
                         if (patient_result is not None) and torch.all(labels.eq(lowerCamelCase		)		):
                              patient_counter += 1
                         else:
                              __A	:       Union[str, Any]      =							0

                    __A	:       int      =							logits
                    if patient_counter == self.patience:
                         break
               __A	:       List[str]      =							[patient_result]
               self.inference_layers_num += calculated_layer_num
               self.inference_instances_num += 1

          return res







@add_start_docstrings(
    'Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of\n    the pooled output) e.g. for GLUE tasks. '    ,		lowerCAmelCase    ,		)
class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__(lowerCamelCase		)
          __A	:       Union[str, Any]      =							config.num_labels

          __A	:       Union[str, Any]      =							BertModelWithPabee(lowerCamelCase		)
          __A	:       Union[str, Any]      =							nn.Dropout(config.hidden_dropout_prob		)
          __A	:       int      =							nn.ModuleList(
              [nn.Linear(config.hidden_size			,     self.config.num_labels		) for _ in range(config.num_hidden_layers		)]		)

          self.init_weights()


     @add_start_docstrings_to_model_forward(lowerCamelCase		)
     def lowerCAmelCase__  (   self			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     ):


          '''simple docstring'''





          __A	:       Optional[Any]      =							self.bert(
              input_ids=lowerCamelCase			,     attention_mask=lowerCamelCase			,     token_type_ids=lowerCamelCase			,     position_ids=lowerCamelCase			,     head_mask=lowerCamelCase			,     inputs_embeds=lowerCamelCase			,     output_dropout=self.dropout			,     output_layers=self.classifiers			,     regression=self.num_labels == 1			,     )

          __A	:       Any      =							(logits[-1],)

          if labels is not None:
               __A	:       Optional[Any]      =							None
               __A	:       int      =							0
               for ix, logits_item in enumerate(lowerCamelCase		):
                    if self.num_labels == 1:
                         #  We are doing regression
                         __A	:       Union[str, Any]      =							MSELoss()
                         __A	:       str      =							loss_fct(logits_item.view(-1		)			,     labels.view(-1		)		)
                    else:
                         __A	:       Optional[int]      =							CrossEntropyLoss()
                         __A	:       Optional[int]      =							loss_fct(logits_item.view(-1			,     self.num_labels		)			,     labels.view(-1		)		)
                    if total_loss is None:
                         __A	:       str      =							loss
                    else:
                         total_loss += loss * (ix + 1)
                    total_weights += ix + 1
               __A	:       Tuple      =							(total_loss / total_weights,) + outputs

          return outputs

