



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







						"""simple docstring"""
						lowerCAmelCase__		  =  parent
						lowerCAmelCase__		  =  batch_size
						lowerCAmelCase__		  =  seq_length
						lowerCAmelCase__		  =  is_training
						lowerCAmelCase__		  =  use_attention_mask
						lowerCAmelCase__		  =  use_token_type_ids
						lowerCAmelCase__		  =  use_labels
						lowerCAmelCase__		  =  vocab_size
						lowerCAmelCase__		  =  hidden_size
						lowerCAmelCase__		  =  num_hidden_layers
						lowerCAmelCase__		  =  num_attention_heads
						lowerCAmelCase__		  =  intermediate_size
						lowerCAmelCase__		  =  hidden_act
						lowerCAmelCase__		  =  hidden_dropout_prob
						lowerCAmelCase__		  =  attention_probs_dropout_prob
						lowerCAmelCase__		  =  max_position_embeddings
						lowerCAmelCase__		  =  type_vocab_size
						lowerCAmelCase__		  =  type_sequence_label_size
						lowerCAmelCase__		  =  initializer_range
						lowerCAmelCase__		  =  num_choices
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  ids_tensor([self.batch_size, self.seq_length]    ,self.vocab_size						)

						lowerCAmelCase__		  =  None
						if self.use_attention_mask:
										lowerCAmelCase__		  =  random_attention_mask([self.batch_size, self.seq_length]						)

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

						lowerCAmelCase__		  =  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=lowerCAmelCase__    ,initializer_range=self.initializer_range    ,)

						return config, input_ids, token_type_ids, attention_mask



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.prepare_config_and_inputs()
						lowerCAmelCase__		  =  config_and_inputs
						lowerCAmelCase__		  =  {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
						return config, inputs_dict



@require_flax
class 	__snake_case   ( __SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		SCREAMING_SNAKE_CASE__								=	(
		    (
		        FlaxAlbertModel,
		        FlaxAlbertForPreTraining,
		        FlaxAlbertForMaskedLM,
		        FlaxAlbertForMultipleChoice,
		        FlaxAlbertForQuestionAnswering,
		        FlaxAlbertForSequenceClassification,
		        FlaxAlbertForTokenClassification,
		        FlaxAlbertForQuestionAnswering,
		    )
		    if is_flax_available()
		    else ()
		)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  FlaxAlbertModelTester(self						)



		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""




						for model_class_name in self.all_model_classes:
										lowerCAmelCase__		  =  model_class_name.from_pretrained('albert-base-v2'						)
										lowerCAmelCase__		  =  model(np.ones((1, 1)						)						)
										self.assertIsNotNone(lowerCAmelCase__						)



@require_flax
class 	__snake_case   ( unittest.TestCase ):
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  FlaxAlbertModel.from_pretrained('albert-base-v2'						)
						lowerCAmelCase__		  =  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]						)
						lowerCAmelCase__		  =  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]						)
						lowerCAmelCase__		  =  model(lowerCAmelCase__    ,attention_mask=lowerCAmelCase__						)[0]
						lowerCAmelCase__		  =  (1, 11, 768)
						self.assertEqual(output.shape    ,lowerCAmelCase__						)
						lowerCAmelCase__		  =  np.array(
						    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]						)

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




def 						a__			(				A__,      A__				):
			SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					''
			for i in table:
						res += inp[i - 1]
			return res






def 						a__			(				A__				):
			return data[1:] + data[0]






def 						a__			(				A__,      A__				):
			SCREAMING_SNAKE_CASE_		:     Optional[int]					=					''
			for i in range(len(A__				)				):
						if a[i] == b[i]:
									res += "0"
						else:
									res += "1"
			return res






def 						a__			(				A__,      A__				):
			SCREAMING_SNAKE_CASE_		:     Tuple					=					int('0b' + data[0] + data[-1],      2				)
			SCREAMING_SNAKE_CASE_		:     Tuple					=					int('0b' + data[1:3],      2				)
			return bin(s[row][col]				)[2:]






def 						a__			(				A__,      A__,      A__,      A__,      A__				):
			SCREAMING_SNAKE_CASE_		:     List[Any]					=					message[:4]
			SCREAMING_SNAKE_CASE_		:     int					=					message[4:]
			SCREAMING_SNAKE_CASE_		:     List[str]					=					apply_table(A__,      A__				)
			SCREAMING_SNAKE_CASE_		:     Dict					=					xor(A__,      A__				)
			SCREAMING_SNAKE_CASE_		:     Any					=					apply_sbox(A__,      temp[:4]				)  # noqa: E741
			SCREAMING_SNAKE_CASE_		:     List[Any]					=					apply_sbox(A__,      temp[4:]				)
			SCREAMING_SNAKE_CASE_		:     Any					=					'0' * (2 - len(A__				)) + l  # noqa: E741
			SCREAMING_SNAKE_CASE_		:     List[Any]					=					'0' * (2 - len(A__				)) + r
			SCREAMING_SNAKE_CASE_		:     Optional[int]					=					apply_table(l + r,      A__				)
			SCREAMING_SNAKE_CASE_		:     Dict					=					xor(A__,      A__				)
			return temp + right


if __name__ == "__main__":
				lowerCAmelCase__			:      Optional[Any]								=input('Enter 10 bit key: ')
				lowerCAmelCase__			:      int								=input('Enter 8 bit message: ')

				lowerCAmelCase__			:      Any								=[6, 3, 7, 4, 8, 5, 10, 9]
				lowerCAmelCase__			:      List[str]								=[3, 5, 2, 7, 4, 10, 1, 9, 8, 6]
				lowerCAmelCase__			:      Dict								=[2, 4, 3, 1]
				lowerCAmelCase__			:      Dict								=[2, 6, 3, 1, 4, 8, 5, 7]
				lowerCAmelCase__			:      Union[str, Any]								=[4, 1, 3, 5, 7, 2, 8, 6]
				lowerCAmelCase__			:      Dict								=[4, 1, 2, 3, 2, 3, 4, 1]
				lowerCAmelCase__			:      str								=[[1, 0, 3, 2], [3, 2, 1, 0], [0, 2, 1, 3], [3, 1, 3, 2]]
				lowerCAmelCase__			:      Any								=[[0, 1, 2, 3], [2, 0, 1, 3], [3, 0, 1, 0], [2, 1, 0, 3]]

				# key generation
				lowerCAmelCase__			:      Dict								=apply_table(key, paa_table)
				lowerCAmelCase__			:      List[str]								=temp[:5]
				lowerCAmelCase__			:      List[Any]								=temp[5:]
				lowerCAmelCase__			:      Tuple								=left_shift(left)
				lowerCAmelCase__			:      Optional[Any]								=left_shift(right)
				lowerCAmelCase__			:      Optional[int]								=apply_table(left + right, pa_table)
				lowerCAmelCase__			:      Optional[Any]								=left_shift(left)
				lowerCAmelCase__			:      Optional[Any]								=left_shift(right)
				lowerCAmelCase__			:      List[str]								=left_shift(left)
				lowerCAmelCase__			:      Optional[int]								=left_shift(right)
				lowerCAmelCase__			:      Tuple								=apply_table(left + right, pa_table)

				# encryption
				lowerCAmelCase__			:      Tuple								=apply_table(message, IP)
				lowerCAmelCase__			:      List[Any]								=function(expansion, sa, sa, keya, temp)
				lowerCAmelCase__			:      str								=temp[4:] + temp[:4]
				lowerCAmelCase__			:      Optional[Any]								=function(expansion, sa, sa, keya, temp)
				lowerCAmelCase__			:      Optional[int]								=apply_table(temp, IP_inv)
				print('Cipher text is:', CT)

				# decryption
				lowerCAmelCase__			:      Union[str, Any]								=apply_table(CT, IP)
				lowerCAmelCase__			:      Tuple								=function(expansion, sa, sa, keya, temp)
				lowerCAmelCase__			:      Union[str, Any]								=temp[4:] + temp[:4]
				lowerCAmelCase__			:      Dict								=function(expansion, sa, sa, keya, temp)
				lowerCAmelCase__			:      int								=apply_table(temp, IP_inv)
				print('Plain text after decypting is:', PT)





