




import unittest

from transformers import AutoTokenizer, NystromformerConfig, 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, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         NystromformerForMaskedLM,
         NystromformerForMultipleChoice,
         NystromformerForQuestionAnswering,
         NystromformerForSequenceClassification,
         NystromformerForTokenClassification,
         NystromformerModel,
     )
     from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST






class  _lowerCAmelCase     :

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

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[str]		       =     ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size		)

        lowerCAmelCase__				:  str		       =     None
        if self.use_input_mask:
          lowerCAmelCase__				:  List[str]		       =     random_attention_mask([self.batch_size, self.seq_length]		)

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

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

        lowerCAmelCase__				:  Optional[int]		       =     self.get_config()

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

      def 	__magic_name__(   self		):
        return NystromformerConfig(
            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      ,				is_decoder=UpperCamelCase__      ,				initializer_range=self.initializer_range      ,				)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  List[str]		       =     NystromformerModel(config=UpperCamelCase__		)
        model.to(UpperCamelCase__		)
        model.eval()
        lowerCAmelCase__				:  Union[str, Any]		       =     model(UpperCamelCase__      ,				attention_mask=UpperCamelCase__      ,				token_type_ids=UpperCamelCase__		)
        lowerCAmelCase__				:  str		       =     model(UpperCamelCase__      ,				token_type_ids=UpperCamelCase__		)
        lowerCAmelCase__				:  Dict		       =     model(UpperCamelCase__		)
        self.parent.assertEqual(result.last_hidden_state.shape      ,				(self.batch_size, self.seq_length, self.hidden_size)		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Tuple		       =     NystromformerForMaskedLM(config=UpperCamelCase__		)
        model.to(UpperCamelCase__		)
        model.eval()
        lowerCAmelCase__				:  str		       =     model(UpperCamelCase__      ,				attention_mask=UpperCamelCase__      ,				token_type_ids=UpperCamelCase__      ,				labels=UpperCamelCase__		)
        self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.vocab_size)		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Any		       =     NystromformerForQuestionAnswering(config=UpperCamelCase__		)
        model.to(UpperCamelCase__		)
        model.eval()
        lowerCAmelCase__				:  Optional[int]		       =     model(
            UpperCamelCase__      ,				attention_mask=UpperCamelCase__      ,				token_type_ids=UpperCamelCase__      ,				start_positions=UpperCamelCase__      ,				end_positions=UpperCamelCase__      ,				)
        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 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Dict		       =     self.num_labels
        lowerCAmelCase__				:  Dict		       =     NystromformerForSequenceClassification(UpperCamelCase__		)
        model.to(UpperCamelCase__		)
        model.eval()
        lowerCAmelCase__				:  List[str]		       =     model(UpperCamelCase__      ,				attention_mask=UpperCamelCase__      ,				token_type_ids=UpperCamelCase__      ,				labels=UpperCamelCase__		)
        self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.num_labels)		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  str		       =     self.num_labels
        lowerCAmelCase__				:  Union[str, Any]		       =     NystromformerForTokenClassification(config=UpperCamelCase__		)
        model.to(UpperCamelCase__		)
        model.eval()
        lowerCAmelCase__				:  int		       =     model(UpperCamelCase__      ,				attention_mask=UpperCamelCase__      ,				token_type_ids=UpperCamelCase__      ,				labels=UpperCamelCase__		)
        self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.num_labels)		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Union[str, Any]		       =     self.num_choices
        lowerCAmelCase__				:  Any		       =     NystromformerForMultipleChoice(config=UpperCamelCase__		)
        model.to(UpperCamelCase__		)
        model.eval()
        lowerCAmelCase__				:  Optional[int]		       =     input_ids.unsqueeze(1		).expand(-1      ,				self.num_choices      ,				-1		).contiguous()
        lowerCAmelCase__				:  Tuple		       =     token_type_ids.unsqueeze(1		).expand(-1      ,				self.num_choices      ,				-1		).contiguous()
        lowerCAmelCase__				:  Tuple		       =     input_mask.unsqueeze(1		).expand(-1      ,				self.num_choices      ,				-1		).contiguous()
        lowerCAmelCase__				:  Union[str, Any]		       =     model(
            UpperCamelCase__      ,				attention_mask=UpperCamelCase__      ,				token_type_ids=UpperCamelCase__      ,				labels=UpperCamelCase__      ,				)
        self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.num_choices)		)







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Dict		       =     self.prepare_config_and_inputs()
        (
            (
            lowerCAmelCase__					
        )       ,			(
            lowerCAmelCase__					
        )       ,			(
            lowerCAmelCase__					
        )       ,			(
            lowerCAmelCase__					
        )       ,			(
            lowerCAmelCase__					
        )       ,			(
            lowerCAmelCase__					
        )       ,			(
            lowerCAmelCase__					
        )       ,								
        )				:  Optional[Any]		       =     config_and_inputs
        lowerCAmelCase__				:  str		       =     {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict






@require_torch
class  _lowerCAmelCase     (  _A    ,			_A    ,			unittest.TestCase	):
      A__		     =		(
          (
              NystromformerModel,
              NystromformerForMaskedLM,
              NystromformerForMultipleChoice,
              NystromformerForQuestionAnswering,
              NystromformerForSequenceClassification,
              NystromformerForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      A__		     =		(
          {
              'feature-extraction': NystromformerModel,
              'fill-mask': NystromformerForMaskedLM,
              'question-answering': NystromformerForQuestionAnswering,
              'text-classification': NystromformerForSequenceClassification,
              'token-classification': NystromformerForTokenClassification,
              'zero-shot': NystromformerForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      A__		     =		False
      A__		     =		False

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Tuple		       =     NystromformerModelTester(self		)
        lowerCAmelCase__				:  Union[str, Any]		       =     ConfigTester(self      ,				config_class=UpperCamelCase__      ,				hidden_size=37		)

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

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*UpperCamelCase__		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  str		       =     self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          lowerCAmelCase__				:  str		       =     type
          self.model_tester.create_and_check_model(*UpperCamelCase__		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  str		       =     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Optional[Any]		       =     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase__		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Tuple		       =     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*UpperCamelCase__		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase__		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__		)







      @slow
      def 	__magic_name__(   self		):






        for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowerCAmelCase__				:  List[Any]		       =     NystromformerModel.from_pretrained(UpperCamelCase__		)
          self.assertIsNotNone(UpperCamelCase__		)






@require_torch
class  _lowerCAmelCase     (  unittest.TestCase	):

      @slow
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     NystromformerModel.from_pretrained('''uw-madison/nystromformer-512'''		)
        lowerCAmelCase__				:  List[str]		       =     torch.tensor([[0, 1, 2, 3, 4, 5]]		)

        with torch.no_grad():
          lowerCAmelCase__				:  Tuple		       =     model(UpperCamelCase__		)[0]

        lowerCAmelCase__				:  Any		       =     torch.Size((1, 6, 768)		)
        self.assertEqual(output.shape      ,				UpperCamelCase__		)

        lowerCAmelCase__				:  List[str]		       =     torch.tensor(
            [[[-0.4532, -0.0936, 0.5137], [-0.2676, 0.0628, 0.6186], [-0.3629, -0.1726, 0.4716]]]		)

        self.assertTrue(torch.allclose(output[:, :3, :3]      ,				UpperCamelCase__      ,				atol=1e-4		)		)







      @slow
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     '''the [MASK] of Belgium is Brussels'''

        lowerCAmelCase__				:  Optional[int]		       =     AutoTokenizer.from_pretrained('''uw-madison/nystromformer-512'''		)
        lowerCAmelCase__				:  Optional[Any]		       =     NystromformerForMaskedLM.from_pretrained('''uw-madison/nystromformer-512'''		)

        lowerCAmelCase__				:  List[Any]		       =     tokenizer(UpperCamelCase__      ,				return_tensors='''pt'''		)

        with torch.no_grad():
          lowerCAmelCase__				:  Optional[int]		       =     model(encoding.input_ids		).logits

        lowerCAmelCase__				:  List[str]		       =     token_logits[:, 2, :].argmax(-1		)[0]

        self.assertEqual(tokenizer.decode(UpperCamelCase__		)      ,				'''capital'''		)







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase_					    = {
    """configuration_lilt""": ["""LILT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """LiltConfig"""],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCAmelCase_					    = [
         """LILT_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """LiltForQuestionAnswering""",
         """LiltForSequenceClassification""",
         """LiltForTokenClassification""",
         """LiltModel""",
         """LiltPreTrainedModel""",
     ]

if TYPE_CHECKING:
     from .configuration_lilt import LILT_PRETRAINED_CONFIG_ARCHIVE_MAP, LiltConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_lilt import (
              LILT_PRETRAINED_MODEL_ARCHIVE_LIST,
              LiltForQuestionAnswering,
              LiltForSequenceClassification,
              LiltForTokenClassification,
              LiltModel,
              LiltPreTrainedModel,
          )

else:
     import sys

     lowerCAmelCase_					    = _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

