



'''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        UpperCAmelCase_			:
    """simple docstring"""







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




    def lowerCamelCase			(       self   :	Dict		):
          snake_case__    :       Any       		= ids_tensor([self.batch_size, self.seq_length]					,	self.vocab_size		)

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

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

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

          snake_case__    :       Union[str, Any]       		= self.get_config()

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




    def lowerCamelCase			(       self   :	str		):
          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=snake_case_					,	initializer_range=self.initializer_range					,	)




    def lowerCamelCase			(       self   :	Any					,	snake_case_   :	List[str]					,	snake_case_   :	Union[str, Any]					,	snake_case_   :	Any					,	snake_case_   :	int					,	snake_case_   :	int					,	snake_case_   :	Optional[Any]					,	snake_case_   :	Dict		):
          snake_case__    :       Tuple       		= NystromformerModel(config=snake_case_		)
          model.to(snake_case_		)
          model.eval()
          snake_case__    :       int       		= model(snake_case_					,	attention_mask=snake_case_					,	token_type_ids=snake_case_		)
          snake_case__    :       List[str]       		= model(snake_case_					,	token_type_ids=snake_case_		)
          snake_case__    :       Tuple       		= model(snake_case_		)
          self.parent.assertEqual(result.last_hidden_state.shape					,	(self.batch_size, self.seq_length, self.hidden_size)		)




    def lowerCamelCase			(       self   :	Any					,	snake_case_   :	Tuple					,	snake_case_   :	Tuple					,	snake_case_   :	Optional[Any]					,	snake_case_   :	str					,	snake_case_   :	int					,	snake_case_   :	List[str]					,	snake_case_   :	Optional[int]		):
          snake_case__    :       str       		= NystromformerForMaskedLM(config=snake_case_		)
          model.to(snake_case_		)
          model.eval()
          snake_case__    :       Tuple       		= model(snake_case_					,	attention_mask=snake_case_					,	token_type_ids=snake_case_					,	labels=snake_case_		)
          self.parent.assertEqual(result.logits.shape					,	(self.batch_size, self.seq_length, self.vocab_size)		)




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




    def lowerCamelCase			(       self   :	Optional[Any]					,	snake_case_   :	Union[str, Any]					,	snake_case_   :	Union[str, Any]					,	snake_case_   :	Optional[int]					,	snake_case_   :	str					,	snake_case_   :	Tuple					,	snake_case_   :	Any					,	snake_case_   :	Optional[int]		):
          snake_case__    :       Union[str, Any]       		= self.num_labels
          snake_case__    :       List[str]       		= NystromformerForTokenClassification(config=snake_case_		)
          model.to(snake_case_		)
          model.eval()
          snake_case__    :       List[str]       		= model(snake_case_					,	attention_mask=snake_case_					,	token_type_ids=snake_case_					,	labels=snake_case_		)
          self.parent.assertEqual(result.logits.shape					,	(self.batch_size, self.seq_length, self.num_labels)		)




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




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







@require_torch
class        UpperCAmelCase_			(			_a						,						_a						,						unittest.TestCase      ):
    """simple docstring"""







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




    def lowerCamelCase			(       self   :	int		):
          snake_case__    :       Dict       		= NystromformerModelTester(self		)
          snake_case__    :       str       		= ConfigTester(self					,	config_class=snake_case_					,	hidden_size=37		)




    def lowerCamelCase			(       self   :	Tuple		):
          self.config_tester.run_common_tests()




    def lowerCamelCase			(       self   :	Tuple		):
          snake_case__    :       List[str]       		= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case_		)




    def lowerCamelCase			(       self   :	Union[str, Any]		):
          snake_case__    :       Optional[int]       		= self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
                snake_case__    :       str       		= type
                self.model_tester.create_and_check_model(*snake_case_		)




    def lowerCamelCase			(       self   :	List[str]		):
          snake_case__    :       int       		= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*snake_case_		)




    def lowerCamelCase			(       self   :	Dict		):
          snake_case__    :       Union[str, Any]       		= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_multiple_choice(*snake_case_		)




    def lowerCamelCase			(       self   :	Union[str, Any]		):
          snake_case__    :       List[str]       		= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*snake_case_		)




    def lowerCamelCase			(       self   :	Optional[Any]		):
          snake_case__    :       Any       		= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*snake_case_		)




    def lowerCamelCase			(       self   :	Any		):
          snake_case__    :       List[str]       		= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*snake_case_		)




    @slow
    def lowerCamelCase			(       self   :	Optional[int]		):
          for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                snake_case__    :       Union[str, Any]       		= NystromformerModel.from_pretrained(snake_case_		)
                self.assertIsNotNone(snake_case_		)







@require_torch
class        UpperCAmelCase_			(			unittest.TestCase      ):
    """simple docstring"""







    @slow
    def lowerCamelCase			(       self   :	List[str]		):
          snake_case__    :       List[Any]       		= NystromformerModel.from_pretrained("""uw-madison/nystromformer-512"""		)
          snake_case__    :       Dict       		= torch.tensor([[0, 1, 2, 3, 4, 5]]		)

          with torch.no_grad():
                snake_case__    :       Any       		= model(snake_case_		)[0]

          snake_case__    :       int       		= torch.Size((1, 6, 768)		)
          self.assertEqual(output.shape					,	snake_case_		)

          snake_case__    :       int       		= 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]					,	snake_case_					,	atol=1E-4		)		)




    @slow
    def lowerCamelCase			(       self   :	Union[str, Any]		):
          snake_case__    :       Optional[int]       		= """the [MASK] of Belgium is Brussels"""

          snake_case__    :       str       		= AutoTokenizer.from_pretrained("""uw-madison/nystromformer-512"""		)
          snake_case__    :       List[str]       		= NystromformerForMaskedLM.from_pretrained("""uw-madison/nystromformer-512"""		)

          snake_case__    :       List[str]       		= tokenizer(snake_case_					,	return_tensors="""pt"""		)

          with torch.no_grad():
                snake_case__    :       List[str]       		= model(encoding.input_ids		).logits

          snake_case__    :       Optional[int]       		= token_logits[:, 2, :].argmax(-1		)[0]

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





'''simple docstring'''





import importlib
import shutil
import threading
import warnings
from typing import List

import fsspec
import fsspec.asyn

from . import compression
from .hffilesystem import HfFileSystem


__a  			=						importlib.util.find_spec("s3fs") is not None

if _has_safs:
       from .safilesystem import SaFileSystem  # noqa: F401

__a  			=						[
    compression.BzaFileSystem,
    compression.GzipFileSystem,
    compression.LzaFileSystem,
    compression.XzFileSystem,
    compression.ZstdFileSystem,
]

# Register custom filesystems
for fs_class in COMPRESSION_FILESYSTEMS + [HfFileSystem]:
       if fs_class.protocol in fsspec.registry and fsspec.registry[fs_class.protocol] is not fs_class:
              warnings.warn(F"A filesystem protocol was already set for {fs_class.protocol} and will be overwritten.")
       fsspec.register_implementation(fs_class.protocol, fs_class, clobber=True)





def 		__snake_case(							_lowerCAmelCase						)						->					str:
      if "://" in dataset_path:
            snake_case__    :       Tuple       		= dataset_path.split("""://"""						)[1]
      return dataset_path





def 		__snake_case(							_lowerCAmelCase						)						->					bool:
      if fs is not None and fs.protocol != "file":
            return True
      else:
            return False





def 		__snake_case(							_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase						)						->					Optional[int]:
      snake_case__    :       int       		= not is_remote_filesystem(_lowerCAmelCase						)
      if is_local:
            # LocalFileSystem.mv does copy + rm, it is more efficient to simply move a local directory
            shutil.move(fs._strip_protocol(_lowerCAmelCase						)	,  fs._strip_protocol(_lowerCAmelCase						)						)
      else:
            fs.mv(_lowerCAmelCase	,  _lowerCAmelCase	,  recursive=_lowerCAmelCase						)





def 		__snake_case(							)						->					None:
      if hasattr(fsspec.asyn	,  """reset_lock"""						):
            # for future fsspec>2022.05.0
            fsspec.asyn.reset_lock()
      else:
            snake_case__    :       Dict       		= None
            snake_case__    :       List[str]       		= None
            snake_case__    :       Union[str, Any]       		= threading.Lock()
