




'''simple docstring'''


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






      def __init__( self   :  Optional[int]						,						lowercase__   :  Tuple						,						lowercase__   :  List[Any]=13						,						lowercase__   :  Dict=7						,						lowercase__   :  Tuple=True						,						lowercase__   :  Optional[int]=True						,						lowercase__   :  int=True						,						lowercase__   :  Tuple=True						,						lowercase__   :  Optional[int]=99						,						lowercase__   :  Any=32						,						lowercase__   :  Dict=5						,						lowercase__   :  Any=4						,						lowercase__   :  str=37						,						lowercase__   :  Tuple="gelu"						,						lowercase__   :  str=0.1						,						lowercase__   :  Optional[Any]=0.1						,						lowercase__   :  Tuple=512						,						lowercase__   :  Optional[Any]=16						,						lowercase__   :  Tuple=2						,						lowercase__   :  Any=0.02						,						lowercase__   :  int=3						,						lowercase__   :  Dict=4						,						lowercase__   :  Union[str, Any]=None						,						):


           '''simple docstring'''

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






      def 							lowercase_						( self   :  int       ):


           '''simple docstring'''

           a_     :							Dict     						=  ids_tensor([self.batch_size, self.seq_length]						,						self.vocab_size       )

           a_     :							Optional[Any]     						=  None
           if self.use_input_mask:
                a_     :							Any     						=  random_attention_mask([self.batch_size, self.seq_length]       )

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

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

           a_     :							Tuple     						=  self.get_config()

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






      def 							lowercase_						( self   :  Optional[Any]       ):


           '''simple docstring'''

           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=lowercase__						,						initializer_range=self.initializer_range						,						)






      def 							lowercase_						( self   :  List[Any]						,						lowercase__   :  Tuple						,						lowercase__   :  Dict						,						lowercase__   :  str						,						lowercase__   :  str						,						lowercase__   :  Tuple						,						lowercase__   :  Any						,						lowercase__   :  Any       ):


           '''simple docstring'''

           a_     :							str     						=  NystromformerModel(config=lowercase__       )
           model.to(lowercase__       )
           model.eval()
           a_     :							int     						=  model(lowercase__						,						attention_mask=lowercase__						,						token_type_ids=lowercase__       )
           a_     :							Dict     						=  model(lowercase__						,						token_type_ids=lowercase__       )
           a_     :							List[str]     						=  model(lowercase__       )
           self.parent.assertEqual(result.last_hidden_state.shape						,						(self.batch_size, self.seq_length, self.hidden_size)       )






      def 							lowercase_						( self   :  Optional[int]						,						lowercase__   :  Optional[Any]						,						lowercase__   :  int						,						lowercase__   :  Optional[int]						,						lowercase__   :  Any						,						lowercase__   :  List[str]						,						lowercase__   :  int						,						lowercase__   :  List[str]       ):


           '''simple docstring'''

           a_     :							Union[str, Any]     						=  NystromformerForMaskedLM(config=lowercase__       )
           model.to(lowercase__       )
           model.eval()
           a_     :							Union[str, Any]     						=  model(lowercase__						,						attention_mask=lowercase__						,						token_type_ids=lowercase__						,						labels=lowercase__       )
           self.parent.assertEqual(result.logits.shape						,						(self.batch_size, self.seq_length, self.vocab_size)       )






      def 							lowercase_						( self   :  Any						,						lowercase__   :  Any						,						lowercase__   :  Union[str, Any]						,						lowercase__   :  Optional[int]						,						lowercase__   :  Dict						,						lowercase__   :  Dict						,						lowercase__   :  List[Any]						,						lowercase__   :  str       ):


           '''simple docstring'''

           a_     :							Dict     						=  NystromformerForQuestionAnswering(config=lowercase__       )
           model.to(lowercase__       )
           model.eval()
           a_     :							Dict     						=  model(
               lowercase__						,						attention_mask=lowercase__						,						token_type_ids=lowercase__						,						start_positions=lowercase__						,						end_positions=lowercase__						,						)
           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 							lowercase_						( self   :  Optional[Any]						,						lowercase__   :  Optional[Any]						,						lowercase__   :  Dict						,						lowercase__   :  Union[str, Any]						,						lowercase__   :  Dict						,						lowercase__   :  List[Any]						,						lowercase__   :  Union[str, Any]						,						lowercase__   :  str       ):


           '''simple docstring'''

           a_     :							Optional[Any]     						=  self.num_labels
           a_     :							Any     						=  NystromformerForSequenceClassification(lowercase__       )
           model.to(lowercase__       )
           model.eval()
           a_     :							Union[str, Any]     						=  model(lowercase__						,						attention_mask=lowercase__						,						token_type_ids=lowercase__						,						labels=lowercase__       )
           self.parent.assertEqual(result.logits.shape						,						(self.batch_size, self.num_labels)       )






      def 							lowercase_						( self   :  Union[str, Any]						,						lowercase__   :  Dict						,						lowercase__   :  List[Any]						,						lowercase__   :  List[Any]						,						lowercase__   :  Tuple						,						lowercase__   :  Optional[Any]						,						lowercase__   :  List[Any]						,						lowercase__   :  str       ):


           '''simple docstring'''

           a_     :							List[Any]     						=  self.num_labels
           a_     :							str     						=  NystromformerForTokenClassification(config=lowercase__       )
           model.to(lowercase__       )
           model.eval()
           a_     :							Union[str, Any]     						=  model(lowercase__						,						attention_mask=lowercase__						,						token_type_ids=lowercase__						,						labels=lowercase__       )
           self.parent.assertEqual(result.logits.shape						,						(self.batch_size, self.seq_length, self.num_labels)       )






      def 							lowercase_						( self   :  Optional[int]						,						lowercase__   :  List[Any]						,						lowercase__   :  Optional[int]						,						lowercase__   :  Union[str, Any]						,						lowercase__   :  List[Any]						,						lowercase__   :  int						,						lowercase__   :  Optional[int]						,						lowercase__   :  str       ):


           '''simple docstring'''

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






      def 							lowercase_						( self   :  Optional[Any]       ):


           '''simple docstring'''

           a_     :							Any     						=  self.prepare_config_and_inputs()
           (
               (
               a_
           )						,     (
               a_
           )						,     (
               a_
           )						,     (
               a_
           )						,     (
               a_
           )						,     (
               a_
           )						,     (
               a_
           )						,     
           )     :							int     						=  config_and_inputs
           a_     :							List[Any]     						=  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
           return config, inputs_dict







@require_torch
class      SCREAMING_SNAKE_CASE    (						snake_case_							,      snake_case_							,      unittest.TestCase				):
      __magic_name__  :       Optional[int]   							=					(
          (
              NystromformerModel,
              NystromformerForMaskedLM,
              NystromformerForMultipleChoice,
              NystromformerForQuestionAnswering,
              NystromformerForSequenceClassification,
              NystromformerForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      __magic_name__  :       Union[str, Any]   							=					(
          {
              '''feature-extraction''': NystromformerModel,
              '''fill-mask''': NystromformerForMaskedLM,
              '''question-answering''': NystromformerForQuestionAnswering,
              '''text-classification''': NystromformerForSequenceClassification,
              '''token-classification''': NystromformerForTokenClassification,
              '''zero-shot''': NystromformerForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __magic_name__  :       Any   							=					False
      __magic_name__  :       str   							=					False






      def 							lowercase_						( self   :  List[Any]       ):


           '''simple docstring'''

           a_     :							str     						=  NystromformerModelTester(self       )
           a_     :							Optional[Any]     						=  ConfigTester(self						,						config_class=lowercase__						,						hidden_size=37       )






      def 							lowercase_						( self   :  int       ):


           '''simple docstring'''

           self.config_tester.run_common_tests()






      def 							lowercase_						( self   :  Any       ):


           '''simple docstring'''

           a_     :							int     						=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowercase__       )






      def 							lowercase_						( self   :  str       ):


           '''simple docstring'''

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






      def 							lowercase_						( self   :  Any       ):


           '''simple docstring'''

           a_     :							List[str]     						=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*lowercase__       )






      def 							lowercase_						( self   :  Optional[Any]       ):


           '''simple docstring'''

           a_     :							List[str]     						=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*lowercase__       )






      def 							lowercase_						( self   :  List[str]       ):


           '''simple docstring'''

           a_     :							Dict     						=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*lowercase__       )






      def 							lowercase_						( self   :  Tuple       ):


           '''simple docstring'''

           a_     :							Optional[int]     						=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*lowercase__       )






      def 							lowercase_						( self   :  Union[str, Any]       ):


           '''simple docstring'''

           a_     :							List[str]     						=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*lowercase__       )






      @slow
      def 							lowercase_						( self   :  List[Any]       ):


           '''simple docstring'''

           for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                a_     :							Tuple     						=  NystromformerModel.from_pretrained(lowercase__       )
                self.assertIsNotNone(lowercase__       )







@require_torch
class      SCREAMING_SNAKE_CASE    (						unittest.TestCase				):






      @slow
      def 							lowercase_						( self   :  Dict       ):


           '''simple docstring'''

           a_     :							Tuple     						=  NystromformerModel.from_pretrained("""uw-madison/nystromformer-512"""       )
           a_     :							Optional[int]     						=  torch.tensor([[0, 1, 2, 3, 4, 5]]       )

           with torch.no_grad():
                a_     :							List[Any]     						=  model(lowercase__       )[0]

           a_     :							Dict     						=  torch.Size((1, 6, 768)       )
           self.assertEqual(output.shape						,						lowercase__       )

           a_     :							Dict     						=  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]						,						lowercase__						,						atol=1e-4       )       )






      @slow
      def 							lowercase_						( self   :  str       ):


           '''simple docstring'''

           a_     :							List[str]     						=  """the [MASK] of Belgium is Brussels"""

           a_     :							Dict     						=  AutoTokenizer.from_pretrained("""uw-madison/nystromformer-512"""       )
           a_     :							Optional[int]     						=  NystromformerForMaskedLM.from_pretrained("""uw-madison/nystromformer-512"""       )

           a_     :							List[str]     						=  tokenizer(lowercase__						,						return_tensors="""pt"""       )

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

           a_     :							Union[str, Any]     						=  token_logits[:, 2, :].argmax(-1       )[0]

           self.assertEqual(tokenizer.decode(lowercase__       )						,						"""capital"""       )






'''simple docstring'''


from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase_      :  str				     =  logging.get_logger(__name__)

lowerCAmelCase_      :  Any				     =  {
    'transfo-xl-wt103': 'https://huggingface.co/transfo-xl-wt103/resolve/main/config.json',
}







class      SCREAMING_SNAKE_CASE    (						snake_case_				):
      __magic_name__  :       str   							=					'''transfo-xl'''
      __magic_name__  :       List[str]   							=					['''mems''']
      __magic_name__  :       Dict   							=					{
          '''n_token''': '''vocab_size''',
          '''hidden_size''': '''d_model''',
          '''num_attention_heads''': '''n_head''',
          '''num_hidden_layers''': '''n_layer''',
      }






      def __init__( self   :  Union[str, Any]						,						lowercase__   :  List[Any]=26_7735						,						lowercase__   :  Optional[Any]=[2_0000, 4_0000, 20_0000]						,						lowercase__   :  Optional[Any]=1024						,						lowercase__   :  str=1024						,						lowercase__   :  Any=16						,						lowercase__   :  int=64						,						lowercase__   :  str=4096						,						lowercase__   :  Union[str, Any]=4						,						lowercase__   :  List[Any]=False						,						lowercase__   :  List[Any]=18						,						lowercase__   :  str=1600						,						lowercase__   :  str=1000						,						lowercase__   :  Any=True						,						lowercase__   :  Optional[Any]=True						,						lowercase__   :  Union[str, Any]=0						,						lowercase__   :  str=-1						,						lowercase__   :  int=True						,						lowercase__   :  str=0.1						,						lowercase__   :  Optional[Any]=0.0						,						lowercase__   :  Tuple=True						,						lowercase__   :  Optional[int]="normal"						,						lowercase__   :  str=0.01						,						lowercase__   :  List[str]=0.01						,						lowercase__   :  Union[str, Any]=0.02						,						lowercase__   :  str=1e-5						,						lowercase__   :  Any=0						,						**lowercase__   :  List[str]						,						):


           '''simple docstring'''

           a_     :							Optional[Any]     						=  vocab_size
           a_     :							Optional[int]     						=  []
           self.cutoffs.extend(lowercase__       )
           if proj_share_all_but_first:
                a_     :							Any     						=  [False] + [True] * len(self.cutoffs       )
           else:
                a_     :							Tuple     						=  [False] + [False] * len(self.cutoffs       )
           a_     :							Tuple     						=  d_model
           a_     :							Optional[int]     						=  d_embed
           a_     :							List[Any]     						=  d_head
           a_     :							List[str]     						=  d_inner
           a_     :							Tuple     						=  div_val
           a_     :							Dict     						=  pre_lnorm
           a_     :							Optional[Any]     						=  n_layer
           a_     :							Dict     						=  n_head
           a_     :							Any     						=  mem_len
           a_     :							Union[str, Any]     						=  same_length
           a_     :							Dict     						=  attn_type
           a_     :							List[str]     						=  clamp_len
           a_     :							str     						=  sample_softmax
           a_     :							Any     						=  adaptive
           a_     :							List[Any]     						=  dropout
           a_     :							str     						=  dropatt
           a_     :							Dict     						=  untie_r
           a_     :							Tuple     						=  init
           a_     :							Optional[int]     						=  init_range
           a_     :							List[Any]     						=  proj_init_std
           a_     :							Optional[int]     						=  init_std
           a_     :							int     						=  layer_norm_epsilon
           super().__init__(eos_token_id=lowercase__						,						**lowercase__       )






      @property
      def 							lowercase_						( self   :  int       ):


           '''simple docstring'''

           logger.info(F"The model {self.model_type} is one of the few models that has no sequence length limit."       )
           return -1






      @max_position_embeddings.setter
      def 							lowercase_						( self   :  Optional[int]						,						lowercase__   :  Tuple       ):


           '''simple docstring'''

           raise NotImplementedError(
               F"The model {self.model_type} is one of the few models that has no sequence length limit."       )
