



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






    """simple docstring"""





    def __init__(  self      :					Dict  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					Optional[Any]=1_3  ,						UpperCamelCase__      :					str=7  ,						UpperCamelCase__      :					str=True  ,						UpperCamelCase__      :					Union[str, Any]=True  ,						UpperCamelCase__      :					Any=True  ,						UpperCamelCase__      :					List[str]=True  ,						UpperCamelCase__      :					Tuple=9_9  ,						UpperCamelCase__      :					List[str]=3_2  ,						UpperCamelCase__      :					Any=5  ,						UpperCamelCase__      :					List[Any]=4  ,						UpperCamelCase__      :					List[str]=3_7  ,						UpperCamelCase__      :					int="gelu"  ,						UpperCamelCase__      :					Union[str, Any]=0.1  ,						UpperCamelCase__      :					str=0.1  ,						UpperCamelCase__      :					List[str]=5_1_2  ,						UpperCamelCase__      :					Optional[Any]=1_6  ,						UpperCamelCase__      :					List[Any]=2  ,						UpperCamelCase__      :					int=0.0_2  ,						UpperCamelCase__      :					List[Any]=3  ,						UpperCamelCase__      :					int=4  ,						UpperCamelCase__      :					Any=None  ,						):


       """simple docstring"""







       UpperCamelCase					=  parent
       UpperCamelCase					=  batch_size
       UpperCamelCase					=  seq_length
       UpperCamelCase					=  is_training
       UpperCamelCase					=  use_input_mask
       UpperCamelCase					=  use_token_type_ids
       UpperCamelCase					=  use_labels
       UpperCamelCase					=  vocab_size
       UpperCamelCase					=  hidden_size
       UpperCamelCase					=  num_hidden_layers
       UpperCamelCase					=  num_attention_heads
       UpperCamelCase					=  intermediate_size
       UpperCamelCase					=  hidden_act
       UpperCamelCase					=  hidden_dropout_prob
       UpperCamelCase					=  attention_probs_dropout_prob
       UpperCamelCase					=  max_position_embeddings
       UpperCamelCase					=  type_vocab_size
       UpperCamelCase					=  type_sequence_label_size
       UpperCamelCase					=  initializer_range
       UpperCamelCase					=  num_labels
       UpperCamelCase					=  num_choices
       UpperCamelCase					=  scope






    def  A       (  self      :					Dict ):


       """simple docstring"""







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

       UpperCamelCase					=  None
       if self.use_input_mask:
          UpperCamelCase					=  random_attention_mask([self.batch_size, self.seq_length] )

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

       UpperCamelCase					=  None
       UpperCamelCase					=  None
       UpperCamelCase					=  None
       if self.use_labels:
          UpperCamelCase					=  ids_tensor([self.batch_size]  ,						self.type_sequence_label_size )
          UpperCamelCase					=  ids_tensor([self.batch_size, self.seq_length]  ,						self.num_labels )
          UpperCamelCase					=  ids_tensor([self.batch_size]  ,						self.num_choices )

       UpperCamelCase					=  self.get_config()

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






    def  A       (  self      :					List[str] ):


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






    def  A       (  self      :					Union[str, Any]  ,						UpperCamelCase__      :					int  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					Optional[Any]  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					Tuple ):


       """simple docstring"""







       UpperCamelCase					=  NystromformerModel(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						token_type_ids=UpperCamelCase__ )
       UpperCamelCase					=  model(UpperCamelCase__  ,						token_type_ids=UpperCamelCase__ )
       UpperCamelCase					=  model(UpperCamelCase__ )
       self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.hidden_size) )






    def  A       (  self      :					Dict  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					Optional[int]  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					int  ,						UpperCamelCase__      :					int  ,						UpperCamelCase__      :					Any ):


       """simple docstring"""







       UpperCamelCase					=  NystromformerForMaskedLM(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  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  A       (  self      :					List[str]  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					int  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					Any  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					Optional[int] ):


       """simple docstring"""







       UpperCamelCase					=  NystromformerForQuestionAnswering(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  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  A       (  self      :					Dict  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					int  ,						UpperCamelCase__      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.num_labels
       UpperCamelCase					=  NystromformerForSequenceClassification(UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  ,						labels=UpperCamelCase__ )
       self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_labels) )






    def  A       (  self      :					Optional[int]  ,						UpperCamelCase__      :					Optional[Any]  ,						UpperCamelCase__      :					Optional[int]  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					Optional[int]  ,						UpperCamelCase__      :					str  ,						UpperCamelCase__      :					Optional[Any]  ,						UpperCamelCase__      :					int ):


       """simple docstring"""







       UpperCamelCase					=  self.num_labels
       UpperCamelCase					=  NystromformerForTokenClassification(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  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  A       (  self      :					str  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					Optional[Any]  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					str ):


       """simple docstring"""







       UpperCamelCase					=  self.num_choices
       UpperCamelCase					=  NystromformerForMultipleChoice(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  input_ids.unsqueeze(1 ).expand(-1  ,						self.num_choices  ,						-1 ).contiguous()
       UpperCamelCase					=  token_type_ids.unsqueeze(1 ).expand(-1  ,						self.num_choices  ,						-1 ).contiguous()
       UpperCamelCase					=  input_mask.unsqueeze(1 ).expand(-1  ,						self.num_choices  ,						-1 ).contiguous()
       UpperCamelCase					=  model(
           UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  ,						labels=UpperCamelCase__  ,						)
       self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_choices) )






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      
       )					=  config_and_inputs
       UpperCamelCase					=  {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
       return config, inputs_dict

@require_torch
class  SCREAMING_SNAKE_CASE (  _a   ,						_a   ,						unittest.TestCase		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  (
        (
            NystromformerModel,
            NystromformerForMaskedLM,
            NystromformerForMultipleChoice,
            NystromformerForQuestionAnswering,
            NystromformerForSequenceClassification,
            NystromformerForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    _SCREAMING_SNAKE_CASE	  =  (
        {
            """feature-extraction""": NystromformerModel,
            """fill-mask""": NystromformerForMaskedLM,
            """question-answering""": NystromformerForQuestionAnswering,
            """text-classification""": NystromformerForSequenceClassification,
            """token-classification""": NystromformerForTokenClassification,
            """zero-shot""": NystromformerForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    _SCREAMING_SNAKE_CASE	  =  False
    _SCREAMING_SNAKE_CASE	  =  False






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


       """simple docstring"""







       UpperCamelCase					=  NystromformerModelTester(self )
       UpperCamelCase					=  ConfigTester(self  ,						config_class=UpperCamelCase__  ,						hidden_size=3_7 )






    def  A       (  self      :					Tuple ):


       """simple docstring"""







       self.config_tester.run_common_tests()






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCamelCase__ )






    def  A       (  self      :					Any ):


       """simple docstring"""







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






    def  A       (  self      :					int ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__ )






    def  A       (  self      :					int ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase__ )






    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*UpperCamelCase__ )






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase__ )






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__ )






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


       """simple docstring"""







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

@require_torch
class  SCREAMING_SNAKE_CASE (  unittest.TestCase		):






    """simple docstring"""





    @slow
    def  A       (  self      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  NystromformerModel.from_pretrained('uw-madison/nystromformer-512' )
       UpperCamelCase					=  torch.tensor([[0, 1, 2, 3, 4, 5]] )

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

       UpperCamelCase					=  torch.Size((1, 6, 7_6_8) )
       self.assertEqual(output.shape  ,						UpperCamelCase__ )

       UpperCamelCase					=  torch.tensor(
           [[[-0.4_5_3_2, -0.0_9_3_6, 0.5_1_3_7], [-0.2_6_7_6, 0.0_6_2_8, 0.6_1_8_6], [-0.3_6_2_9, -0.1_7_2_6, 0.4_7_1_6]]] )

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






    @slow
    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       UpperCamelCase					=  'the [MASK] of Belgium is Brussels'

       UpperCamelCase					=  AutoTokenizer.from_pretrained('uw-madison/nystromformer-512' )
       UpperCamelCase					=  NystromformerForMaskedLM.from_pretrained('uw-madison/nystromformer-512' )

       UpperCamelCase					=  tokenizer(UpperCamelCase__  ,						return_tensors='pt' )

       with torch.no_grad():
          UpperCamelCase					=  model(encoding.input_ids ).logits

       UpperCamelCase					=  token_logits[:, 2, :].argmax(-1 )[0]

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






'''simple docstring'''


import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


_lowerCamelCase	:       List[str]      	=	5_0000
_lowerCamelCase	:       Optional[int]      	=	5000

_lowerCamelCase	,_lowerCamelCase	:       int      	=	os.path.split(__file__)
_lowerCamelCase	:       str      	=	os.path.join(RESULTS_BASEPATH, "results", RESULTS_FILENAME.replace(".py", ".json"))







@get_duration
def 				__lowerCamelCase				(   A__    ,						A__    )       ->	Any:






   """simple docstring"""
   for i in range(A__    ):
      UpperCamelCase					=  dataset[i]







@get_duration
def 				__lowerCamelCase				(   A__    ,						A__    ,						A__    )       ->	int:






   """simple docstring"""
   for i in range(0    ,						len(A__    )    ,						A__    ):
      UpperCamelCase					=  dataset[i : i + batch_size]







@get_duration
def 				__lowerCamelCase				(   A__    ,						A__    ,						A__    )       ->	List[Any]:






   """simple docstring"""
   with dataset.formatted_as(type=A__    ):
      for i in range(A__    ):
         UpperCamelCase					=  dataset[i]







@get_duration
def 				__lowerCamelCase				(   A__    ,						A__    ,						A__    ,						A__    )       ->	int:






   """simple docstring"""
   with dataset.formatted_as(type=A__    ):
      for i in range(0    ,						A__    ,						A__    ):
         UpperCamelCase					=  dataset[i : i + batch_size]







def 				__lowerCamelCase				(   )       ->	List[str]:






   """simple docstring"""
   UpperCamelCase					=  {'num examples': SPEED_TEST_N_EXAMPLES}
   UpperCamelCase					=  [
       (read, {'length': SMALL_TEST}),
       (read, {'length': SPEED_TEST_N_EXAMPLES}),
       (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 10}),
       (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 100}),
       (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 1_000}),
       (read_formatted, {'type': 'numpy', 'length': SMALL_TEST}),
       (read_formatted, {'type': 'pandas', 'length': SMALL_TEST}),
       (read_formatted, {'type': 'torch', 'length': SMALL_TEST}),
       (read_formatted, {'type': 'tensorflow', 'length': SMALL_TEST}),
       (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 10}),
       (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 1_000}),
   ]

   UpperCamelCase					=  [
       (read, {'length': SMALL_TEST}),
       (read, {'length': SPEED_TEST_N_EXAMPLES}),
       (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 10}),
       (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 100}),
       (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 1_000}),
       (read_formatted, {'type': 'numpy', 'length': SMALL_TEST}),
       (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 10}),
       (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 1_000}),
   ]
   with tempfile.TemporaryDirectory() as tmp_dir:
      print('generating dataset'    )
      UpperCamelCase					=  datasets.Features(
          {'list': datasets.Sequence(datasets.Value('float32'    )    ), 'numbers': datasets.Value('float32'    )}    )
      UpperCamelCase					=  generate_example_dataset(
          os.path.join(A__    ,						'dataset.arrow'    )    ,						A__    ,						num_examples=A__    ,						seq_shapes={'list': (100,)}    ,						)
      print('first set of iterations'    )
      for func, kwargs in functions:
         print(func.__name__    ,						str(A__    )    )
         UpperCamelCase					=  func(A__    ,						**A__    )

      print('shuffling dataset'    )
      UpperCamelCase					=  dataset.shuffle()
      print('Second set of iterations (after shuffling'    )
      for func, kwargs in functions_shuffled:
         print('shuffled '    ,						func.__name__    ,						str(A__    )    )
         UpperCamelCase					=  func(
             A__    ,						**A__    )

   with open(A__    ,						'wb'    ) as f:
      f.write(json.dumps(A__    ).encode('utf-8'    )    )


if __name__ == "__main__":  # useful to run the profiler
   benchmark_iterating()

