




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				:	Optional[Any]			,      __lowercase				:	int			,      __lowercase				:	str=13			,      __lowercase				:	Tuple=7			,      __lowercase				:	Any=True			,      __lowercase				:	Optional[int]=True			,      __lowercase				:	Any=True			,      __lowercase				:	Tuple=True			,      __lowercase				:	List[Any]=99			,      __lowercase				:	List[Any]=32			,      __lowercase				:	List[Any]=5			,      __lowercase				:	str=4			,      __lowercase				:	Optional[int]=37			,      __lowercase				:	Any="gelu"			,      __lowercase				:	int=0.1			,      __lowercase				:	Dict=0.1			,      __lowercase				:	List[Any]=5_12			,      __lowercase				:	str=16			,      __lowercase				:	Any=2			,      __lowercase				:	List[str]=0.02			,      __lowercase				:	List[str]=3			,      __lowercase				:	List[str]=4			,      __lowercase				:	List[str]=None			,      ):






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	int		):






							"""simple docstring"""






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

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

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

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

							snake_case_						       =	self.get_config()

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






						def 			snake_case__		(		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 			snake_case__		(		self				:	str			,      __lowercase				:	Union[str, Any]			,      __lowercase				:	List[str]			,      __lowercase				:	str			,      __lowercase				:	Optional[Any]			,      __lowercase				:	Union[str, Any]			,      __lowercase				:	Dict			,      __lowercase				:	Tuple		):






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	Any			,      __lowercase				:	Union[str, Any]			,      __lowercase				:	List[str]			,      __lowercase				:	Any			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	Optional[Any]			,      __lowercase				:	int		):






							"""simple docstring"""






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






							"""simple docstring"""






							snake_case_						       =	NystromformerForQuestionAnswering(config=__lowercase		)
							model.to(__lowercase		)
							model.eval()
							snake_case_						       =	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 			snake_case__		(		self				:	int			,      __lowercase				:	Dict			,      __lowercase				:	str			,      __lowercase				:	int			,      __lowercase				:	str			,      __lowercase				:	Optional[Any]			,      __lowercase				:	List[Any]			,      __lowercase				:	List[Any]		):






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	Dict			,      __lowercase				:	List[Any]			,      __lowercase				:	List[Any]			,      __lowercase				:	List[Any]			,      __lowercase				:	Union[str, Any]			,      __lowercase				:	Optional[int]			,      __lowercase				:	List[Any]			,      __lowercase				:	Dict		):






							"""simple docstring"""






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






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	Tuple		):






							"""simple docstring"""






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







@require_torch
class 						UpperCAmelCase  (							UpperCAmelCase__			,  UpperCAmelCase__			,  unittest.TestCase     ):







						'''simple docstring'''





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






						def 			snake_case__		(		self				:	Dict		):






							"""simple docstring"""






							snake_case_						       =	NystromformerModelTester(self		)
							snake_case_						       =	ConfigTester(self			,      config_class=__lowercase			,      hidden_size=37		)






						def 			snake_case__		(		self				:	Any		):






							"""simple docstring"""






							self.config_tester.run_common_tests()






						def 			snake_case__		(		self				:	Dict		):






							"""simple docstring"""






							snake_case_						       =	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__lowercase		)






						def 			snake_case__		(		self				:	Optional[int]		):






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	Tuple		):






							"""simple docstring"""






							snake_case_						       =	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__lowercase		)






						def 			snake_case__		(		self				:	Dict		):






							"""simple docstring"""






							snake_case_						       =	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*__lowercase		)






						def 			snake_case__		(		self				:	Tuple		):






							"""simple docstring"""






							snake_case_						       =	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*__lowercase		)






						def 			snake_case__		(		self				:	Optional[Any]		):






							"""simple docstring"""






							snake_case_						       =	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*__lowercase		)






						def 			snake_case__		(		self				:	List[Any]		):






							"""simple docstring"""






							snake_case_						       =	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*__lowercase		)






						@slow
						def 			snake_case__		(		self				:	Optional[int]		):






							"""simple docstring"""






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







@require_torch
class 						UpperCAmelCase  (							unittest.TestCase     ):







						'''simple docstring'''





						@slow
						def 			snake_case__		(		self				:	Union[str, Any]		):






							"""simple docstring"""






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

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

							snake_case_						       =	torch.Size((1, 6, 7_68)		)
							self.assertEqual(output.shape			,      __lowercase		)

							snake_case_						       =	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 			snake_case__		(		self				:	int		):






							"""simple docstring"""






							snake_case_						       =	"the [MASK] of Belgium is Brussels"

							snake_case_						       =	AutoTokenizer.from_pretrained("uw-madison/nystromformer-512"		)
							snake_case_						       =	NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512"		)

							snake_case_						       =	tokenizer(__lowercase			,      return_tensors="pt"		)

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

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

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






import numpy as np
import qiskit


def 		lowerCamelCase__       (						_A = 8	,      _A = None				):







	'''simple docstring'''





	snake_case_						       =	np.random.default_rng(seed=_A				)

	# Roughly 25% of the qubits will contribute to the key.
	# So we take more than we need.
	snake_case_						       =	6 * key_len
	# Measurement basis for Alice's qubits.
	snake_case_						       =	rng.integers(2	,      size=_A				)
	# The set of states Alice will prepare.
	snake_case_						       =	rng.integers(2	,      size=_A				)
	# Measurement basis for Bob's qubits.
	snake_case_						       =	rng.integers(2	,      size=_A				)

	# Quantum Circuit to simulate BB84
	snake_case_						       =	qiskit.QuantumCircuit(_A	,      name="BB84"				)

	# Alice prepares her qubits according to rules above.
	for index, _ in enumerate(_A				):
		if alice_state[index] == 1:
			bbaa_circ.x(_A				)
		if alice_basis[index] == 1:
			bbaa_circ.h(_A				)
	bbaa_circ.barrier()

	# Bob measures the received qubits according to rules above.
	for index, _ in enumerate(_A				):
		if bob_basis[index] == 1:
			bbaa_circ.h(_A				)

	bbaa_circ.barrier()
	bbaa_circ.measure_all()

	# Simulate the quantum circuit.
	snake_case_						       =	qiskit.Aer.get_backend("aer_simulator"				)
	# We only need to run one shot because the key is unique.
	# Multiple shots will produce the same key.
	snake_case_						       =	qiskit.execute(_A	,      _A	,      shots=1	,      seed_simulator=_A				)
	# Returns the result of measurement.
	snake_case_						       =	job.result().get_counts(_A				).most_frequent()

	# Extracting the generated key from the simulation results.
	# Only keep measurement results where Alice and Bob chose the same basis.
	snake_case_						       =	"".join(
	    [
	        result_bit
	        for alice_basis_bit, bob_basis_bit, result_bit in zip(
	            _A	,      _A	,      _A				)
	        if alice_basis_bit == bob_basis_bit
	    ]				)

	# Get final key. Pad with 0 if too short, otherwise truncate.
	snake_case_						       =	gen_key[:key_len] if len(_A				) >= key_len else gen_key.ljust(_A	,      "0"				)
	return key


if __name__ == "__main__":
		print(f'''The generated key is : {bbaa(8, seed=0)}''')
		from doctest import testmod

		testmod()
