
import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
							import jax.numpy as jnp

							from transformers.models.albert.modeling_flax_albert import (
							    FlaxAlbertForMaskedLM,
							    FlaxAlbertForMultipleChoice,
							    FlaxAlbertForPreTraining,
							    FlaxAlbertForQuestionAnswering,
							    FlaxAlbertForSequenceClassification,
							    FlaxAlbertForTokenClassification,
							    FlaxAlbertModel,
							)




class   __lowercase						(						unittest.TestCase       ):






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






	def       _a(self				: Tuple    ) ->			List[Any]:
							_lowercase       :	Union[str, Any]           =   ids_tensor([self.batch_size, self.seq_length]      ,	self.vocab_size    )

							_lowercase       :	Union[str, Any]           =   None
							if self.use_attention_mask:
													_lowercase       :	Optional[Any]           =   random_attention_mask([self.batch_size, self.seq_length]    )

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

							_lowercase       :	Union[str, Any]           =   AlbertConfig(
							    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      ,	)

							return config, input_ids, token_type_ids, attention_mask






	def       _a(self				: str    ) ->			List[Any]:
							_lowercase       :	Dict           =   self.prepare_config_and_inputs()
							_lowercase      ,    _lowercase      ,    _lowercase      ,    _lowercase       :	List[str]           =   config_and_inputs
							_lowercase       :	Dict           =   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
							return config, inputs_dict




@require_flax
class   __lowercase						(						__snake_case   , unittest.TestCase       ):
	_A       =   (
	    (
	        FlaxAlbertModel,
	        FlaxAlbertForPreTraining,
	        FlaxAlbertForMaskedLM,
	        FlaxAlbertForMultipleChoice,
	        FlaxAlbertForQuestionAnswering,
	        FlaxAlbertForSequenceClassification,
	        FlaxAlbertForTokenClassification,
	        FlaxAlbertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)






	def       _a(self				: Tuple    ) ->			Dict:
							_lowercase       :	Optional[Any]           =   FlaxAlbertModelTester(self    )






	@slow
	def       _a(self				: Dict    ) ->			Optional[Any]:
							for model_class_name in self.all_model_classes:
													_lowercase       :	List[Any]           =   model_class_name.from_pretrained("albert-base-v2"    )
													_lowercase       :	Tuple           =   model(np.ones((1, 1)    )    )
													self.assertIsNotNone(snake_case    )




@require_flax
class   __lowercase						(						unittest.TestCase       ):






	@slow
	def       _a(self				: Optional[Any]    ) ->			List[Any]:
							_lowercase       :	Dict           =   FlaxAlbertModel.from_pretrained("albert-base-v2"    )
							_lowercase       :	Optional[int]           =   np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
							_lowercase       :	Union[str, Any]           =   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
							_lowercase       :	str           =   model(snake_case      ,	attention_mask=snake_case    )[0]
							_lowercase       :	Union[str, Any]           =   (1, 11, 768)
							self.assertEqual(output.shape      ,	snake_case    )
							_lowercase       :	List[str]           =   np.array(
							    [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]    )

							self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]      ,	snake_case      ,	atol=1e-4    )    )






import os
import random
import sys

from . import cryptomath_module as cryptomath
from . import rabin_miller

UpperCamelCase_						:     List[Any]	       =							3



def 		UpperCamelCase     (       _UpperCAmelCase  :   int  )   ->			int:






						'''simple docstring'''

						print("Generating primitive root of p"  )
						while True:
												_lowercase       :	Any           =   random.randrange(3 ,			_UpperCAmelCase  )
												if pow(_UpperCAmelCase ,			2 ,			_UpperCAmelCase  ) == 1:
																		continue
												if pow(_UpperCAmelCase ,			_UpperCAmelCase ,			_UpperCAmelCase  ) == 1:
																		continue
												return g



def 		UpperCamelCase     (       _UpperCAmelCase  :   int  )   ->			tuple[tuple[int, int, int, int], tuple[int, int]]:






						'''simple docstring'''

						print("Generating prime p..."  )
						_lowercase       :	Optional[int]           =   rabin_miller.generate_large_prime(_UpperCAmelCase  )  # select large prime number.
						_lowercase       :	Dict           =   primitive_root(_UpperCAmelCase  )  # one primitive root on modulo p.
						_lowercase       :	Any           =   random.randrange(3 ,			_UpperCAmelCase  )  # private_key -> have to be greater than 2 for safety.
						_lowercase       :	Tuple           =   cryptomath.find_mod_inverse(pow(_UpperCAmelCase ,			_UpperCAmelCase ,			_UpperCAmelCase  ) ,			_UpperCAmelCase  )

						_lowercase       :	List[Any]           =   (key_size, e_a, e_a, p)
						_lowercase       :	Union[str, Any]           =   (key_size, d)

						return public_key, private_key



def 		UpperCamelCase     (       _UpperCAmelCase  :   str ,			_UpperCAmelCase  :   int  )   ->			None:






						'''simple docstring'''

						if os.path.exists(f"""{name}_pubkey.txt"""  ) or os.path.exists(f"""{name}_privkey.txt"""  ):
												print("\nWARNING:"  )
												print(
												    f"""\"{name}_pubkey.txt\" or \"{name}_privkey.txt\" already exists. \n"""
												    "Use a different name or delete these files and re-run this program."  )
												sys.exit()

						_lowercase      ,    _lowercase       :	Optional[Any]           =   generate_key(_UpperCAmelCase  )
						print(f"""\nWriting public key to file {name}_pubkey.txt..."""  )
						with open(f"""{name}_pubkey.txt""" ,			"w"  ) as fo:
												fo.write(f"""{public_key[0]},{public_key[1]},{public_key[2]},{public_key[3]}"""  )

						print(f"""Writing private key to file {name}_privkey.txt..."""  )
						with open(f"""{name}_privkey.txt""" ,			"w"  ) as fo:
												fo.write(f"""{private_key[0]},{private_key[1]}"""  )



def 		UpperCamelCase     (       )   ->			None:






						'''simple docstring'''

						print("Making key files..."  )
						make_key_files("elgamal" ,			2048  )
						print("Key files generation successful"  )


if __name__ == "__main__":
							main()




