




'''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 			_a  :
       def __init__(	self						:	Optional[Any]   ,				lowercase						:	int   ,				lowercase						:	str=13   ,				lowercase						:	Any=7   ,				lowercase						:	str=True   ,				lowercase						:	int=True   ,				lowercase						:	int=True   ,				lowercase						:	Any=True   ,				lowercase						:	Any=99   ,				lowercase						:	Any=32   ,				lowercase						:	Dict=5   ,				lowercase						:	Optional[int]=4   ,				lowercase						:	Dict=37   ,				lowercase						:	int="gelu"   ,				lowercase						:	Union[str, Any]=0.1   ,				lowercase						:	Union[str, Any]=0.1   ,				lowercase						:	str=512   ,				lowercase						:	Tuple=16   ,				lowercase						:	List[str]=2   ,				lowercase						:	str=0.02   ,				lowercase						:	str=3   ,				lowercase						:	Dict=4   ,				lowercase						:	int=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						:	Optional[Any]				):



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



           '''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     A      (	self						:	Optional[Any]   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Optional[Any]   ,				lowercase						:	Tuple   ,				lowercase						:	Optional[int]   ,				lowercase						:	Optional[int]   ,				lowercase						:	List[Any]   ,				lowercase						:	Optional[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		NystromformerModel(config=lowercase				)
           model.to(lowercase				)
           model.eval()
           UpperCAmelCase			       =		model(lowercase   ,				attention_mask=lowercase   ,				token_type_ids=lowercase				)
           UpperCAmelCase			       =		model(lowercase   ,				token_type_ids=lowercase				)
           UpperCAmelCase			       =		model(lowercase				)
           self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.hidden_size)				)
       def     A      (	self						:	Union[str, Any]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Any   ,				lowercase						:	str   ,				lowercase						:	int   ,				lowercase						:	int   ,				lowercase						:	Dict   ,				lowercase						:	int				):



           '''simple docstring'''

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



           '''simple docstring'''

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



           '''simple docstring'''

           UpperCAmelCase			       =		self.num_labels
           UpperCAmelCase			       =		NystromformerForSequenceClassification(lowercase				)
           model.to(lowercase				)
           model.eval()
           UpperCAmelCase			       =		model(lowercase   ,				attention_mask=lowercase   ,				token_type_ids=lowercase   ,				labels=lowercase				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_labels)				)
       def     A      (	self						:	Any   ,				lowercase						:	str   ,				lowercase						:	List[Any]   ,				lowercase						:	str   ,				lowercase						:	Optional[int]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Tuple   ,				lowercase						:	Optional[Any]				):



           '''simple docstring'''

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



           '''simple docstring'''

           UpperCAmelCase			       =		self.num_choices
           UpperCAmelCase			       =		NystromformerForMultipleChoice(config=lowercase				)
           model.to(lowercase				)
           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(
               lowercase   ,				attention_mask=lowercase   ,				token_type_ids=lowercase   ,				labels=lowercase   ,				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_choices)				)







       def     A      (	self						:	Dict				):



           '''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 			_a  (	__a							,							__a							,							unittest.TestCase ):
       __a				:     Optional[int]   				=					(
           (
               NystromformerModel,
               NystromformerForMaskedLM,
               NystromformerForMultipleChoice,
               NystromformerForQuestionAnswering,
               NystromformerForSequenceClassification,
               NystromformerForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       __a				:     List[Any]   				=					(
           {
               """feature-extraction""": NystromformerModel,
               """fill-mask""": NystromformerForMaskedLM,
               """question-answering""": NystromformerForQuestionAnswering,
               """text-classification""": NystromformerForSequenceClassification,
               """token-classification""": NystromformerForTokenClassification,
               """zero-shot""": NystromformerForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       __a				:     Optional[int]   				=					False
       __a				:     int   				=					False
       def     A      (	self						:	int				):



           '''simple docstring'''

           UpperCAmelCase			       =		NystromformerModelTester(self				)
           UpperCAmelCase			       =		ConfigTester(self   ,				config_class=lowercase   ,				hidden_size=37				)
       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(*lowercase				)
       def     A      (	self						:	str				):



           '''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(*lowercase				)
       def     A      (	self						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*lowercase				)
       def     A      (	self						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*lowercase				)
       def     A      (	self						:	int				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*lowercase				)
       def     A      (	self						:	Tuple				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*lowercase				)
       def     A      (	self						:	Union[str, Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*lowercase				)







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



           '''simple docstring'''

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






@require_torch
class 			_a  (	unittest.TestCase ):
       @slow
       def     A      (	self						:	Dict				):



           '''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(lowercase				)[0]

           UpperCAmelCase			       =		torch.Size((1, 6, 768)				)
           self.assertEqual(output.shape   ,				lowercase				)

           UpperCAmelCase			       =		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     A      (	self						:	List[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(lowercase   ,				return_tensors='''pt'''				)

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

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

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











import pytest


__UpperCAmelCase       :   Optional[Any]				       =							"__dummy_dataset1__"

__UpperCAmelCase       :   List[str]				       =							"\nimport json\nimport os\n\nimport datasets\n\n\nREPO_URL = \"https://huggingface.co/datasets/albertvillanova/tests-raw-jsonl/resolve/main/\"\nURLS = {\"train\": REPO_URL + \"wikiann-bn-train.jsonl\", \"validation\": REPO_URL + \"wikiann-bn-validation.jsonl\"}\n\n\nclass __DummyDataset1__(datasets.GeneratorBasedBuilder):\n\n    def _info(self):\n        features = datasets.Features(\n            {\n                \"tokens\": datasets.Sequence(datasets.Value(\"string\")),\n                \"ner_tags\": datasets.Sequence(\n                    datasets.features.ClassLabel(\n                        names=[\n                            \"O\",\n                            \"B-PER\",\n                            \"I-PER\",\n                            \"B-ORG\",\n                            \"I-ORG\",\n                            \"B-LOC\",\n                            \"I-LOC\",\n                        ]\n                    )\n                ),\n                \"langs\": datasets.Sequence(datasets.Value(\"string\")),\n                \"spans\": datasets.Sequence(datasets.Value(\"string\")),\n            }\n        )\n        return datasets.DatasetInfo(features=features)\n\n    def _split_generators(self, dl_manager):\n        dl_path = dl_manager.download(URLS)\n        return [\n            datasets.SplitGenerator(datasets.Split.TRAIN, gen_kwargs={\"filepath\": dl_path[\"train\"]}),\n            datasets.SplitGenerator(datasets.Split.VALIDATION, gen_kwargs={\"filepath\": dl_path[\"validation\"]}),\n        ]\n\n    def _generate_examples(self, filepath):\n        with open(filepath, \"r\", encoding=\"utf-8\") as f:\n            for i, line in enumerate(f):\n                yield i, json.loads(line)\n"







@pytest.fixture
def 			A__							(    )   ->    Optional[int]:
  return DATASET_LOADING_SCRIPT_NAME







@pytest.fixture
def 			A__							(    )   ->    Tuple:
  return DATASET_LOADING_SCRIPT_CODE







@pytest.fixture
def 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)   ->    Tuple:
  __snake_case:						List[Any]   =						dataset_loading_script_name
  __snake_case:						Any   =						tmp_path / """datasets""" / script_name
  script_dir.mkdir(parents=SCREAMING_SNAKE_CASE__)
  __snake_case:						int   =						script_dir / F'''{script_name}.py'''
  with open(SCREAMING_SNAKE_CASE__					,					"""w""") as f:
    f.write(SCREAMING_SNAKE_CASE__)
  return str(SCREAMING_SNAKE_CASE__)






