
import unittest

from transformers import LayoutLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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 (
          LayoutLMForMaskedLM,
          LayoutLMForQuestionAnswering,
          LayoutLMForSequenceClassification,
          LayoutLMForTokenClassification,
          LayoutLMModel,
      )






class     a_       :



       """simple docstring"""






       def __init__(		self				,       _lowerCamelCase				,       _lowerCamelCase=13				,       _lowerCamelCase=7				,       _lowerCamelCase=True				,       _lowerCamelCase=True				,       _lowerCamelCase=True				,       _lowerCamelCase=True				,       _lowerCamelCase=99				,       _lowerCamelCase=32				,       _lowerCamelCase=5				,       _lowerCamelCase=4				,       _lowerCamelCase=37				,       _lowerCamelCase="gelu"				,       _lowerCamelCase=0.1				,       _lowerCamelCase=0.1				,       _lowerCamelCase=512				,       _lowerCamelCase=16				,       _lowerCamelCase=2				,       _lowerCamelCase=0.0_2				,       _lowerCamelCase=3				,       _lowerCamelCase=4				,       _lowerCamelCase=None				,       _lowerCamelCase=1000				,       )					->Union[str, Any]:
           SCREAMING_SNAKE_CASE				:				List[Any]				      = parent
           SCREAMING_SNAKE_CASE				:				Any				      = batch_size
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = seq_length
           SCREAMING_SNAKE_CASE				:				Any				      = is_training
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = use_input_mask
           SCREAMING_SNAKE_CASE				:				List[Any]				      = use_token_type_ids
           SCREAMING_SNAKE_CASE				:				str				      = use_labels
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = vocab_size
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = hidden_size
           SCREAMING_SNAKE_CASE				:				str				      = num_hidden_layers
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = num_attention_heads
           SCREAMING_SNAKE_CASE				:				Tuple				      = intermediate_size
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = hidden_act
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = hidden_dropout_prob
           SCREAMING_SNAKE_CASE				:				List[Any]				      = attention_probs_dropout_prob
           SCREAMING_SNAKE_CASE				:				str				      = max_position_embeddings
           SCREAMING_SNAKE_CASE				:				Dict				      = type_vocab_size
           SCREAMING_SNAKE_CASE				:				Tuple				      = type_sequence_label_size
           SCREAMING_SNAKE_CASE				:				int				      = initializer_range
           SCREAMING_SNAKE_CASE				:				Tuple				      = num_labels
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = num_choices
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = scope
           SCREAMING_SNAKE_CASE				:				int				      = range_bbox






       def 	__lowerCAmelCase     (		self   )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size   )

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

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

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

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

           SCREAMING_SNAKE_CASE				:				Dict				      = self.get_config()

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






       def 	__lowerCAmelCase     (		self   )					->int:
           return LayoutLMConfig(
               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 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->Any:
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = LayoutLMModel(config=_lowerCamelCase   )
           model.to(_lowerCamelCase   )
           model.eval()
           SCREAMING_SNAKE_CASE				:				List[Any]				      = model(_lowerCamelCase				,       _lowerCamelCase				,       attention_mask=_lowerCamelCase				,       token_type_ids=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				int				      = model(_lowerCamelCase				,       _lowerCamelCase				,       token_type_ids=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = model(_lowerCamelCase				,       _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 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->Dict:
           SCREAMING_SNAKE_CASE				:				Any				      = LayoutLMForMaskedLM(config=_lowerCamelCase   )
           model.to(_lowerCamelCase   )
           model.eval()
           SCREAMING_SNAKE_CASE				:				Dict				      = model(_lowerCamelCase				,       _lowerCamelCase				,       attention_mask=_lowerCamelCase				,       token_type_ids=_lowerCamelCase				,       labels=_lowerCamelCase   )
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)   )






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->Optional[Any]:
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = self.num_labels
           SCREAMING_SNAKE_CASE				:				Tuple				      = LayoutLMForSequenceClassification(_lowerCamelCase   )
           model.to(_lowerCamelCase   )
           model.eval()
           SCREAMING_SNAKE_CASE				:				Dict				      = model(
               _lowerCamelCase				,       _lowerCamelCase				,       attention_mask=_lowerCamelCase				,       token_type_ids=_lowerCamelCase				,       labels=_lowerCamelCase   )
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)   )






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->str:
           SCREAMING_SNAKE_CASE				:				str				      = self.num_labels
           SCREAMING_SNAKE_CASE				:				Any				      = LayoutLMForTokenClassification(config=_lowerCamelCase   )
           model.to(_lowerCamelCase   )
           model.eval()
           SCREAMING_SNAKE_CASE				:				Dict				      = model(_lowerCamelCase				,       _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 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->Optional[int]:
           SCREAMING_SNAKE_CASE				:				Dict				      = LayoutLMForQuestionAnswering(config=_lowerCamelCase   )
           model.to(_lowerCamelCase   )
           model.eval()
           SCREAMING_SNAKE_CASE				:				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 	__lowerCAmelCase     (		self   )					->Optional[int]:
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = self.prepare_config_and_inputs()
           (
               (
               SCREAMING_SNAKE_CASE
           )      ,							(
               SCREAMING_SNAKE_CASE
           )      ,							(
               SCREAMING_SNAKE_CASE
           )      ,							(
               SCREAMING_SNAKE_CASE
           )      ,							(
               SCREAMING_SNAKE_CASE
           )      ,							(
               SCREAMING_SNAKE_CASE
           )      ,							(
               SCREAMING_SNAKE_CASE
           )      ,							(
               SCREAMING_SNAKE_CASE
           )      ,							
           )				:				Any				      = config_and_inputs
           SCREAMING_SNAKE_CASE				:				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_       (			a__				,    a__				,    unittest.TestCase							):



       """simple docstring"""
       __SCREAMING_SNAKE_CASE       :					Tuple						=  (
           (
               LayoutLMModel,
               LayoutLMForMaskedLM,
               LayoutLMForSequenceClassification,
               LayoutLMForTokenClassification,
               LayoutLMForQuestionAnswering,
           )
           if is_torch_available()
           else None
       )
       __SCREAMING_SNAKE_CASE       :					Tuple						=  (
           {
               'document-question-answering': LayoutLMForQuestionAnswering,
               'feature-extraction': LayoutLMModel,
               'fill-mask': LayoutLMForMaskedLM,
               'text-classification': LayoutLMForSequenceClassification,
               'token-classification': LayoutLMForTokenClassification,
               'zero-shot': LayoutLMForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       __SCREAMING_SNAKE_CASE       :					Optional[int]						=  True






       def 	__lowerCAmelCase     (		self   )					->int:
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = LayoutLMModelTester(self   )
           SCREAMING_SNAKE_CASE				:				int				      = ConfigTester(self				,       config_class=_lowerCamelCase				,       hidden_size=37   )






       def 	__lowerCAmelCase     (		self   )					->Optional[int]:
           self.config_tester.run_common_tests()






       def 	__lowerCAmelCase     (		self   )					->Any:
           SCREAMING_SNAKE_CASE				:				Tuple				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->str:
           SCREAMING_SNAKE_CASE				:				str				      = self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
               SCREAMING_SNAKE_CASE				:				Dict				      = type
               self.model_tester.create_and_check_model(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->Union[str, Any]:
           SCREAMING_SNAKE_CASE				:				List[Any]				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->str:
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->int:
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				Dict				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*_lowerCamelCase   )






       @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''   )
       def 	__lowerCAmelCase     (		self   )					->str:
           pass







def 				UpperCAmelCase_(    ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Dict				      = torch.tensor([[101,1_019,1_014,1_016,1_037,12_849,4_747,1_004,14_246,2_278,5_439,4_524,5_002,2_930,2_193,2_930,4_341,3_208,1_005,1_055,2_171,2_848,11_300,3_531,102],[101,4_070,4_034,7_020,1_024,3_058,1_015,1_013,2_861,1_013,6_070,19_274,2_772,6_205,27_814,16_147,16_147,4_343,2_047,10_283,10_969,14_389,1_012,2_338,102]] ,						device=a__   )  # noqa: E231
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],] ,						device=a__   )  # noqa: E231
    SCREAMING_SNAKE_CASE				:				List[str]				      = torch.tensor([[[0,0,0,0],[423,237,440,251],[427,272,441,287],[419,115,437,129],[961,885,992,912],[256,38,330,58],[256,38,330,58],[336,42,353,57],[360,39,401,56],[360,39,401,56],[411,39,471,59],[479,41,528,59],[533,39,630,60],[67,113,134,131],[141,115,209,132],[68,149,133,166],[141,149,187,164],[195,148,287,165],[195,148,287,165],[195,148,287,165],[295,148,349,165],[441,149,492,166],[497,149,546,164],[64,201,125,218],[1_000,1_000,1_000,1_000]],[[0,0,0,0],[662,150,754,166],[665,199,742,211],[519,213,554,228],[519,213,554,228],[134,433,187,454],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[314,469,376,482],[504,684,582,706],[941,825,973,900],[941,825,973,900],[941,825,973,900],[941,825,973,900],[610,749,652,765],[130,659,168,672],[176,657,237,672],[238,657,312,672],[443,653,628,672],[443,653,628,672],[716,301,825,317],[1_000,1_000,1_000,1_000]]] ,						device=a__   )  # noqa: E231
    SCREAMING_SNAKE_CASE				:				str				      = torch.tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]] ,						device=a__   )  # noqa: E231
    # these are sequence labels (i.e. at the token level)
    SCREAMING_SNAKE_CASE				:				Any				      = torch.tensor([[-100,10,10,10,9,1,-100,7,7,-100,7,7,4,2,5,2,8,8,-100,-100,5,0,3,2,-100],[-100,12,12,12,-100,12,10,-100,-100,-100,-100,10,12,9,-100,-100,-100,10,10,10,9,12,-100,10,-100]] ,						device=a__   )  # noqa: E231
    # fmt: on

    return input_ids, attention_mask, bbox, token_type_ids, labels






@require_torch
class     a_       (			unittest.TestCase							):



       """simple docstring"""






       @slow
       def 	__lowerCAmelCase     (		self   )					->Union[str, Any]:
           SCREAMING_SNAKE_CASE				:				Dict				      = LayoutLMModel.from_pretrained('''microsoft/layoutlm-base-uncased'''   ).to(_lowerCamelCase   )

           SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				str				      = prepare_layoutlm_batch_inputs()

           # forward pass
           SCREAMING_SNAKE_CASE				:				Tuple				      = model(input_ids=_lowerCamelCase				,       bbox=_lowerCamelCase				,       attention_mask=_lowerCamelCase				,       token_type_ids=_lowerCamelCase   )

           # test the sequence output on [0, :3, :3]
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = torch.tensor(
               [[0.1_7_8_5, -0.1_9_4_7, -0.0_4_2_5], [-0.3_2_5_4, -0.2_8_0_7, 0.2_5_5_3], [-0.5_3_9_1, -0.3_3_2_2, 0.3_3_6_4]]				,       device=_lowerCamelCase				,       )

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

           # test the pooled output on [1, :3]
           SCREAMING_SNAKE_CASE				:				Dict				      = torch.tensor([-0.6_5_8_0, -0.0_2_1_4, 0.8_5_5_2]				,       device=_lowerCamelCase   )

           self.assertTrue(torch.allclose(outputs.pooler_output[1, :3]				,       _lowerCamelCase				,       atol=1e-3   )   )






       @slow
       def 	__lowerCAmelCase     (		self   )					->Dict:
           # initialize model with randomly initialized sequence classification head
           SCREAMING_SNAKE_CASE				:				Any				      = LayoutLMForSequenceClassification.from_pretrained('''microsoft/layoutlm-base-uncased'''				,       num_labels=2   ).to(
               _lowerCamelCase   )

           SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				str				      = prepare_layoutlm_batch_inputs()

           # forward pass
           SCREAMING_SNAKE_CASE				:				Tuple				      = model(
               input_ids=_lowerCamelCase				,       bbox=_lowerCamelCase				,       attention_mask=_lowerCamelCase				,       token_type_ids=_lowerCamelCase				,       labels=torch.tensor([1, 1]				,       device=_lowerCamelCase   )				,       )

           # test whether we get a loss as a scalar
           SCREAMING_SNAKE_CASE				:				int				      = outputs.loss
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.Size([]   )
           self.assertEqual(loss.shape				,       _lowerCamelCase   )

           # test the shape of the logits
           SCREAMING_SNAKE_CASE				:				List[Any]				      = outputs.logits
           SCREAMING_SNAKE_CASE				:				Any				      = torch.Size((2, 2)   )
           self.assertEqual(logits.shape				,       _lowerCamelCase   )






       @slow
       def 	__lowerCAmelCase     (		self   )					->Any:
           # initialize model with randomly initialized token classification head
           SCREAMING_SNAKE_CASE				:				str				      = LayoutLMForTokenClassification.from_pretrained('''microsoft/layoutlm-base-uncased'''				,       num_labels=13   ).to(
               _lowerCamelCase   )

           SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				Tuple				      = prepare_layoutlm_batch_inputs()

           # forward pass
           SCREAMING_SNAKE_CASE				:				Dict				      = model(
               input_ids=_lowerCamelCase				,       bbox=_lowerCamelCase				,       attention_mask=_lowerCamelCase				,       token_type_ids=_lowerCamelCase				,       labels=_lowerCamelCase   )

           # test the loss calculation to be around 2.65
           # expected_loss = torch.tensor(2.65, device=torch_device)

           # The loss is currently somewhat random and can vary between 0.1-0.3 atol.
           # self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=0.1))

           # test the shape of the logits
           SCREAMING_SNAKE_CASE				:				List[str]				      = outputs.logits
           SCREAMING_SNAKE_CASE				:				Any				      = torch.Size((2, 25, 13)   )
           self.assertEqual(logits.shape				,       _lowerCamelCase   )






       @slow
       def 	__lowerCAmelCase     (		self   )					->Union[str, Any]:
           # initialize model with randomly initialized token classification head
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = LayoutLMForQuestionAnswering.from_pretrained('''microsoft/layoutlm-base-uncased'''   ).to(_lowerCamelCase   )

           SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				str				      = prepare_layoutlm_batch_inputs()

           # forward pass
           SCREAMING_SNAKE_CASE				:				List[str]				      = model(input_ids=_lowerCamelCase				,       bbox=_lowerCamelCase				,       attention_mask=_lowerCamelCase				,       token_type_ids=_lowerCamelCase   )

           # test the shape of the logits
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.Size((2, 25)   )
           self.assertEqual(outputs.start_logits.shape				,       _lowerCamelCase   )
           self.assertEqual(outputs.end_logits.shape				,       _lowerCamelCase   )


import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


a__				:      Tuple    		=     [nn.LayerNorm]

a__				:      Tuple    		=     logging.get_logger(__name__)

a__				:      int    		=     version.parse(version.parse(torch.__version__).base_version)

a__				:      Dict    		=     parsed_torch_version_base >= version.parse('''2.0''')
a__				:      Any    		=     parsed_torch_version_base >= version.parse('''1.12''')
a__				:      Union[str, Any]    		=     parsed_torch_version_base >= version.parse('''1.11''')
a__				:      Union[str, Any]    		=     parsed_torch_version_base < version.parse('''1.11''')







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						a__ ,						self   ):



    """simple docstring"""



    from torch import _softmax_backward_data

    if is_torch_less_than_1_11:
        return _softmax_backward_data(a__ ,						a__ ,						parent.dim ,						self   )
    else:
        return _softmax_backward_data(a__ ,						a__ ,						parent.dim ,						self.dtype   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = 0   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = index.to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = layer.weight.index_select(a__ ,						a__   ).clone().detach()
    if layer.bias is not None:
        if dim == 1:
            SCREAMING_SNAKE_CASE				:				int				      = layer.bias.clone().detach()
        else:
            SCREAMING_SNAKE_CASE				:				List[Any]				      = layer.bias[index].clone().detach()
    SCREAMING_SNAKE_CASE				:				List[Any]				      = list(layer.weight.size()   )
    SCREAMING_SNAKE_CASE				:				List[Any]				      = len(a__   )
    SCREAMING_SNAKE_CASE				:				str				      = nn.Linear(new_size[1] ,						new_size[0] ,						bias=layer.bias is not None   ).to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = False
    new_layer.weight.copy_(W.contiguous()   )
    SCREAMING_SNAKE_CASE				:				Tuple				      = True
    if layer.bias is not None:
        SCREAMING_SNAKE_CASE				:				Dict				      = False
        new_layer.bias.copy_(b.contiguous()   )
        SCREAMING_SNAKE_CASE				:				Tuple				      = True
    return new_layer






class     a_       (			nn.Module							):



       """simple docstring"""






       def __init__(		self				,       _lowerCamelCase				,       _lowerCamelCase   )					->int:
           super().__init__()
           SCREAMING_SNAKE_CASE				:				List[str]				      = nf
           SCREAMING_SNAKE_CASE				:				List[Any]				      = nn.Parameter(torch.empty(_lowerCamelCase				,       _lowerCamelCase   )   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = nn.Parameter(torch.zeros(_lowerCamelCase   )   )
           nn.init.normal_(self.weight				,       std=0.0_2   )






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase   )					->Tuple:
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = x.size()[:-1] + (self.nf,)
           SCREAMING_SNAKE_CASE				:				int				      = torch.addmm(self.bias				,       x.view(-1				,       x.size(-1   )   )				,       self.weight   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = x.view(_lowerCamelCase   )
           return x







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = 1   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				str				      = index.to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Any				      = layer.weight.index_select(a__ ,						a__   ).clone().detach()
    if dim == 0:
        SCREAMING_SNAKE_CASE				:				Tuple				      = layer.bias.clone().detach()
    else:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = layer.bias[index].clone().detach()
    SCREAMING_SNAKE_CASE				:				int				      = list(layer.weight.size()   )
    SCREAMING_SNAKE_CASE				:				Tuple				      = len(a__   )
    SCREAMING_SNAKE_CASE				:				Any				      = ConvaD(new_size[1] ,						new_size[0]   ).to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = False
    new_layer.weight.copy_(W.contiguous()   )
    SCREAMING_SNAKE_CASE				:				int				      = True
    SCREAMING_SNAKE_CASE				:				Any				      = False
    new_layer.bias.copy_(b.contiguous()   )
    SCREAMING_SNAKE_CASE				:				int				      = True
    return new_layer







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = None   ):



    """simple docstring"""



    if isinstance(a__ ,						nn.Linear   ):
        return prune_linear_layer(a__ ,						a__ ,						dim=0 if dim is None else dim   )
    elif isinstance(a__ ,						a__   ):
        return prune_convad_layer(a__ ,						a__ ,						dim=1 if dim is None else dim   )
    else:
        raise ValueError(F"""Can't prune layer of class {layer.__class__}"""   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						*a__   ):



    """simple docstring"""



    assert len(a__   ) > 0, F"""{input_tensors} has to be a tuple/list of tensors"""

    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    SCREAMING_SNAKE_CASE				:				str				      = len(inspect.signature(a__   ).parameters   )
    if num_args_in_forward_chunk_fn != len(a__   ):
        raise ValueError(
            F"""forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(a__   )} input """
            '''tensors are given'''   )

    if chunk_size > 0:
        SCREAMING_SNAKE_CASE				:				Dict				      = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    F"""All input tenors have to be of the same shape: {tensor_shape}, """
                    F"""found shape {input_tensor.shape[chunk_dim]}"""   )

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                F"""The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk """
                F"""size {chunk_size}"""   )

        SCREAMING_SNAKE_CASE				:				Any				      = input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = tuple(input_tensor.chunk(a__ ,						dim=a__   ) for input_tensor in input_tensors   )
        # apply forward fn to every tuple
        SCREAMING_SNAKE_CASE				:				Tuple				      = tuple(forward_fn(*a__   ) for input_tensors_chunk in zip(*a__   )   )
        # concatenate output at same dimension
        return torch.cat(a__ ,						dim=a__   )

    return forward_fn(*a__   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						a__   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.ones(a__ ,						a__   )
    SCREAMING_SNAKE_CASE				:				str				      = set(a__   ) - already_pruned_heads  # Convert to set and remove already pruned heads
    for head in heads:
        # Compute how many pruned heads are before the head and move the index accordingly
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = head - sum(1 if h < head else 0 for h in already_pruned_heads   )
        SCREAMING_SNAKE_CASE				:				str				      = 0
    SCREAMING_SNAKE_CASE				:				List[str]				      = mask.view(-1   ).contiguous().eq(1   )
    SCREAMING_SNAKE_CASE				:				torch.LongTensor				      = torch.arange(len(a__   )   )[mask].long()
    return heads, index







def 				UpperCAmelCase_(    *a__ ,						a__ = None   ):



    """simple docstring"""



    return torch.meshgrid(*a__ ,						indexing=a__   )







def 				UpperCAmelCase_(    a__   ):



    """simple docstring"""



    return tensor.device, storage_ptr(a__   ), storage_size(a__   )
