



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   :







					'''simple docstring'''







					def __init__(      self     ,						lowerCAmelCase     ,						lowerCAmelCase=13     ,						lowerCAmelCase=7     ,						lowerCAmelCase=True     ,						lowerCAmelCase=True     ,						lowerCAmelCase=True     ,						lowerCAmelCase=True     ,						lowerCAmelCase=99     ,						lowerCAmelCase=32     ,						lowerCAmelCase=5     ,						lowerCAmelCase=4     ,						lowerCAmelCase=37     ,						lowerCAmelCase="gelu"     ,						lowerCAmelCase=0.1     ,						lowerCAmelCase=0.1     ,						lowerCAmelCase=512     ,						lowerCAmelCase=16     ,						lowerCAmelCase=2     ,						lowerCAmelCase=0.02     ,						lowerCAmelCase=3     ,						lowerCAmelCase=4     ,						lowerCAmelCase=None     ,						):
							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  ):
							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  ):
							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     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase  ):
							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     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase  ):
							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     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase  ):
							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     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase  ):
							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     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase  ):
							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     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase     ,						lowerCAmelCase  ):
							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  ):
							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 						lowerCamelCase   (    __snake_case,  __snake_case,  unittest.TestCase	):







					'''simple docstring'''







					lowerCAmelCase_					:     int					       =	(
					    (
					        NystromformerModel,
					        NystromformerForMaskedLM,
					        NystromformerForMultipleChoice,
					        NystromformerForQuestionAnswering,
					        NystromformerForSequenceClassification,
					        NystromformerForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					lowerCAmelCase_					:     Any					       =	(
					    {
					        'feature-extraction': NystromformerModel,
					        'fill-mask': NystromformerForMaskedLM,
					        'question-answering': NystromformerForQuestionAnswering,
					        'text-classification': NystromformerForSequenceClassification,
					        'token-classification': NystromformerForTokenClassification,
					        'zero-shot': NystromformerForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					lowerCAmelCase_					:     Optional[Any]					       =	False
					lowerCAmelCase_					:     Optional[Any]					       =	False



					def  A__		(      self  ):
							UpperCAmelCase_						       =					NystromformerModelTester(self  )
							UpperCAmelCase_						       =					ConfigTester(self     ,						config_class=__UpperCamelCase     ,						hidden_size=37  )



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



					def  A__		(      self  ):
							UpperCAmelCase_						       =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__UpperCamelCase  )



					def  A__		(      self  ):
							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  ):
							UpperCAmelCase_						       =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase  )



					def  A__		(      self  ):
							UpperCAmelCase_						       =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*__UpperCamelCase  )



					def  A__		(      self  ):
							UpperCAmelCase_						       =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase  )



					def  A__		(      self  ):
							UpperCAmelCase_						       =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase  )



					def  A__		(      self  ):
							UpperCAmelCase_						       =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase  )



					@slow
					def  A__		(      self  ):


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





@require_torch
class 						lowerCamelCase   (    unittest.TestCase	):







					'''simple docstring'''







					@slow
					def  A__		(      self  ):
							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, 768)  )
							self.assertEqual(output.shape     ,						__UpperCamelCase  )

							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]     ,						__UpperCamelCase     ,						atol=1e-4  )  )



					@slow
					def  A__		(      self  ):
							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"  )




import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
A									=  logging.getLogger()

A									=  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)


class SCREAMING_SNAKE_CASE	(		__snake_case				):






			"""simple docstring"""




			def       __lowerCAmelCase			(     self			,					__UpperCamelCase				):

				"""simple docstring"""


				os.makedirs(__UpperCamelCase			,					exist_ok=__UpperCamelCase				)
				snake_case_   		=			{'source': 'What is love ?', 'target': 'life'}
				snake_case_   		=			{'train': 12, 'val': 2, 'test': 2}
				for split in ["train", "test", "val"]:
					for field in ["source", "target"]:
						snake_case_   		=			'\n'.join([contents[field]] * n_lines[split]				)
						with open(os.path.join(__UpperCamelCase			,					f"""{split}.{field}"""				)			,					'w'				) as f:
							f.write(__UpperCamelCase				)


			def       __lowerCAmelCase			(     self			,					__UpperCamelCase			,					__UpperCamelCase = "pytorch"				):

				"""simple docstring"""


				snake_case_   		=			self.get_auto_remove_tmp_dir()
				snake_case_   		=			os.path.join(__UpperCamelCase			,					'output'				)
				snake_case_   		=			os.path.join(__UpperCamelCase			,					'data'				)
				self._create_dummy_data(data_dir=__UpperCamelCase				)

				snake_case_   		=			f"""
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            """.split()

				if gpus > 0:
					testargs.append(f"""--gpus={gpus}"""				)
					if is_apex_available():
						testargs.append('--fp16'				)
				else:
					testargs.append('--gpus=0'				)
					testargs.append('--distributed_backend=ddp_cpu'				)
					testargs.append('--num_processes=2'				)

				snake_case_   		=			[sys.executable, str(Path(finetune_rag.__file__				).resolve()				)] + testargs
				execute_subprocess_async(__UpperCamelCase			,					env=self.get_env()				)

				snake_case_   		=			os.path.join(__UpperCamelCase			,					'metrics.json'				)
				with open(__UpperCamelCase				) as f:
					snake_case_   		=			json.load(__UpperCamelCase				)
				return result


			@require_torch_gpu
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self._run_finetune(gpus=1				)
				self.assertGreaterEqual(result['test'][0]['test_avg_em']			,					0.2				)


			@require_torch_multi_gpu
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self._run_finetune(gpus=2				)
				self.assertGreaterEqual(result['test'][0]['test_avg_em']			,					0.2				)


			@require_torch_gpu
			@require_ray
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self._run_finetune(gpus=1			,					distributed_retriever='ray'				)
				self.assertGreaterEqual(result['test'][0]['test_avg_em']			,					0.2				)







			@require_torch_multi_gpu
			@require_ray
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self._run_finetune(gpus=1			,					distributed_retriever='ray'				)
				self.assertGreaterEqual(result['test'][0]['test_avg_em']			,					0.2				)
