

import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
    import os

    import jax.numpy as jnp
    from jax import jit

    from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
    from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

    __A			:					Optional[Any]						=			'0.12'  # assumed parallelism: 8


if is_torch_available():
    import torch





def 	__a     (      A__			:  Any			,  A__			:  Tuple			,  A__			:  str=None       ):
    if rng is None:
        SCREAMING_SNAKE_CASE										=   random.Random()

    SCREAMING_SNAKE_CASE										=   1
    for dim in shape:
        total_dims *= dim

    SCREAMING_SNAKE_CASE										=   []
    for _ in range(A__       ):
        values.append(rng.randint(0			,  vocab_size - 1       )       )

    SCREAMING_SNAKE_CASE										=   np.array(A__			,  dtype=jnp.intaa       ).reshape(A__       )

    return output





def 	__a     (      A__			:  Union[str, Any]			,  A__			:  Union[str, Any]=None       ):
    SCREAMING_SNAKE_CASE										=   ids_tensor(A__			,  vocab_size=2			,  rng=A__       )
    # make sure that at least one token is attended to for each batch
    SCREAMING_SNAKE_CASE										=   1
    return attn_mask





@require_flax
class 			_SCREAMING_SNAKE_CASE :







   '''simple docstring'''




   lowerCamelCase__				=      None
   lowerCamelCase__				=      ()





   def      _snake_case		(       self					:			List[Any]			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs_for_common()

       # cut to half length & take max batch_size 3
       SCREAMING_SNAKE_CASE										=   2
       SCREAMING_SNAKE_CASE										=   inputs["input_ids"].shape[-1] // 2
       SCREAMING_SNAKE_CASE										=   inputs["input_ids"][:max_batch_size, :sequence_length]

       SCREAMING_SNAKE_CASE										=   jnp.ones_like(__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   attention_mask[:max_batch_size, :sequence_length]

       # generate max 5 tokens
       SCREAMING_SNAKE_CASE										=   input_ids.shape[-1] + 5
       if config.eos_token_id is not None and config.pad_token_id is None:
           # hack to allow generate for models such as GPT2 as is done in `generate()`
           SCREAMING_SNAKE_CASE										=   config.eos_token_id
       return config, input_ids, attention_mask, max_length





   @is_pt_flax_cross_test
   def      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()
       SCREAMING_SNAKE_CASE										=   False
       SCREAMING_SNAKE_CASE										=   max_length
       SCREAMING_SNAKE_CASE										=   0

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
           SCREAMING_SNAKE_CASE										=   getattr(__lowerCamelCase     , __lowerCamelCase			)
           SCREAMING_SNAKE_CASE										=   pt_model_class(__lowerCamelCase			).eval()
           SCREAMING_SNAKE_CASE										=   load_flax_weights_in_pytorch_model(__lowerCamelCase     , flax_model.params			)

           SCREAMING_SNAKE_CASE										=   flax_model.generate(__lowerCamelCase			).sequences
           SCREAMING_SNAKE_CASE										=   pt_model.generate(torch.tensor(__lowerCamelCase     , dtype=torch.long			)			)

           if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
               SCREAMING_SNAKE_CASE										=   flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

           self.assertListEqual(pt_generation_outputs.numpy().tolist()     , flax_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			List[Any]			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()
       SCREAMING_SNAKE_CASE										=   False
       SCREAMING_SNAKE_CASE										=   max_length

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			Tuple			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()
       SCREAMING_SNAKE_CASE										=   True
       SCREAMING_SNAKE_CASE										=   max_length

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			List[Any]			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()
       SCREAMING_SNAKE_CASE										=   False
       SCREAMING_SNAKE_CASE										=   max_length
       SCREAMING_SNAKE_CASE										=   2

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			Any			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()
       SCREAMING_SNAKE_CASE										=   False
       SCREAMING_SNAKE_CASE										=   max_length
       SCREAMING_SNAKE_CASE										=   2
       SCREAMING_SNAKE_CASE										=   2

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[0]     , input_ids.shape[0] * config.num_return_sequences			)





   def      _snake_case		(       self					:			Dict			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()
       SCREAMING_SNAKE_CASE										=   True
       SCREAMING_SNAKE_CASE										=   max_length
       SCREAMING_SNAKE_CASE										=   0.8
       SCREAMING_SNAKE_CASE										=   10
       SCREAMING_SNAKE_CASE										=   0.3
       SCREAMING_SNAKE_CASE										=   1
       SCREAMING_SNAKE_CASE										=   8
       SCREAMING_SNAKE_CASE										=   9

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			str			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()
       SCREAMING_SNAKE_CASE										=   max_length
       SCREAMING_SNAKE_CASE										=   1
       SCREAMING_SNAKE_CASE										=   8
       SCREAMING_SNAKE_CASE										=   9

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			int			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()
       SCREAMING_SNAKE_CASE										=   max_length
       SCREAMING_SNAKE_CASE										=   2
       SCREAMING_SNAKE_CASE										=   1
       SCREAMING_SNAKE_CASE										=   8
       SCREAMING_SNAKE_CASE										=   9

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			Any			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()

       # pad attention mask on the left
       SCREAMING_SNAKE_CASE										=   attention_mask.at[(0, 0)].set(0			)

       SCREAMING_SNAKE_CASE										=   False
       SCREAMING_SNAKE_CASE										=   max_length

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase     , attention_mask=__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase     , attention_mask=__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			str			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()

       # pad attention mask on the left
       SCREAMING_SNAKE_CASE										=   attention_mask.at[(0, 0)].set(0			)

       SCREAMING_SNAKE_CASE										=   True
       SCREAMING_SNAKE_CASE										=   max_length

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase     , attention_mask=__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase     , attention_mask=__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





   def      _snake_case		(       self					:			int			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self._get_input_ids_and_config()

       # pad attention mask on the left
       SCREAMING_SNAKE_CASE										=   attention_mask.at[(0, 0)].set(0			)

       SCREAMING_SNAKE_CASE										=   2
       SCREAMING_SNAKE_CASE										=   max_length

       for model_class in self.all_generative_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   model.generate(__lowerCamelCase     , attention_mask=__lowerCamelCase			).sequences
           self.assertEqual(generation_outputs.shape[-1]     , __lowerCamelCase			)

           SCREAMING_SNAKE_CASE										=   jit(model.generate			)
           SCREAMING_SNAKE_CASE										=   jit_generate(__lowerCamelCase     , attention_mask=__lowerCamelCase			).sequences

           self.assertListEqual(generation_outputs.tolist()     , jit_generation_outputs.tolist()			)





@require_flax
class 			_SCREAMING_SNAKE_CASE (   unittest.TestCase		):







   '''simple docstring'''




   def      _snake_case		(       self					:			Dict			):
       SCREAMING_SNAKE_CASE										=   AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"			)
       SCREAMING_SNAKE_CASE										=   FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"			)

       SCREAMING_SNAKE_CASE										=   "Hello world"
       SCREAMING_SNAKE_CASE										=   tokenizer(__lowerCamelCase     , return_tensors="np"			).input_ids

       # typos are quickly detected (the correct argument is `do_sample`)
       with self.assertRaisesRegex(__lowerCamelCase     , "do_samples"			):
           model.generate(__lowerCamelCase     , do_samples=__lowerCamelCase			)

       # arbitrary arguments that will not be used anywhere are also not accepted
       with self.assertRaisesRegex(__lowerCamelCase     , "foo"			):
           SCREAMING_SNAKE_CASE										=   {"foo": "bar"}
           model.generate(__lowerCamelCase     , **__lowerCamelCase			)

import string

import numpy
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     int			,							UpperCAmelCase_	:     int				)			->						int:
       return b if a == 0 else greatest_common_divisor(b % a			,							UpperCAmelCase_				)


class 			lowercase_		:
       __lowerCamelCase				       =						string.ascii_uppercase + string.digits
       # This cipher takes alphanumerics into account
       # i.e. a total of 36 characters

       # take x and return x % len(key_string)
       __lowerCamelCase				       =						numpy.vectorize(lambda A       : x % 3_6       )

       __lowerCamelCase				       =						numpy.vectorize(A       )


       def __init__(    self						,       __A				)						->    None:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=self.modulus(__A				)  # mod36 calc's on the encrypt key
              self.check_determinant()  # validate the determinant of the encryption key
              SCREAMING_SNAKE_CASE_							:				Dict		=encrypt_key.shape[0]


       def 	_snake_case				(    self						,       __A				)						->    int:
              return self.key_string.index(__A				)


       def 	_snake_case				(    self						,       __A				)						->    str:
              return self.key_string[round(__A				)]


       def 	_snake_case				(    self				)						->    None:
              SCREAMING_SNAKE_CASE_							:				List[str]		=round(numpy.linalg.det(self.encrypt_key				)				)

              if det < 0:
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=det % len(self.key_string				)

              SCREAMING_SNAKE_CASE_							:				List[Any]		=len(self.key_string				)
              if greatest_common_divisor(__A						,       len(self.key_string				)				) != 1:
                     SCREAMING_SNAKE_CASE_							:				Any		=(
                         F'determinant modular {req_l} of encryption key({det}) '
                         F'is not co prime w.r.t {req_l}.\nTry another key.'
                     )
                     raise ValueError(__A				)


       def 	_snake_case				(    self						,       __A				)						->    str:
              SCREAMING_SNAKE_CASE_							:				List[Any]		=[char for char in text.upper() if char in self.key_string]

              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=chars[-1]
              while len(__A				) % self.break_key != 0:
                     chars.append(__A				)

              return "".join(__A				)


       def 	_snake_case				(    self						,       __A				)						->    str:
              SCREAMING_SNAKE_CASE_							:				Dict		=self.process_text(text.upper()				)
              SCREAMING_SNAKE_CASE_							:				Dict		=''''''

              for i in range(0						,       len(__A				) - self.break_key + 1						,       self.break_key				):
                     SCREAMING_SNAKE_CASE_							:				Any		=text[i : i + self.break_key]
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=[self.replace_letters(__A				) for char in batch]
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=numpy.array([vec]				).T
                     SCREAMING_SNAKE_CASE_							:				List[str]		=self.modulus(self.encrypt_key.dot(__A				)				).T.tolist()[
                         0
                     ]
                     SCREAMING_SNAKE_CASE_							:				List[Any]		=''''''.join(
                         self.replace_digits(__A				) for num in batch_encrypted				)
                     encrypted += encrypted_batch

              return encrypted


       def 	_snake_case				(    self				)						->    numpy.ndarray:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=round(numpy.linalg.det(self.encrypt_key				)				)

              if det < 0:
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=det % len(self.key_string				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=None
              for i in range(len(self.key_string				)				):
                     if (det * i) % len(self.key_string				) == 1:
                            SCREAMING_SNAKE_CASE_							:				List[str]		=i
                            break

              SCREAMING_SNAKE_CASE_							:				List[str]		=(
                  det_inv
                  * numpy.linalg.det(self.encrypt_key				)
                  * numpy.linalg.inv(self.encrypt_key				)
              )

              return self.to_int(self.modulus(__A				)				)





       def 	_snake_case				(    self						,       __A				)						->    str:
              SCREAMING_SNAKE_CASE_							:				int		=self.make_decrypt_key()
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.process_text(text.upper()				)
              SCREAMING_SNAKE_CASE_							:				Any		=''''''

              for i in range(0						,       len(__A				) - self.break_key + 1						,       self.break_key				):
                     SCREAMING_SNAKE_CASE_							:				List[str]		=text[i : i + self.break_key]
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=[self.replace_letters(__A				) for char in batch]
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=numpy.array([vec]				).T
                     SCREAMING_SNAKE_CASE_							:				str		=self.modulus(decrypt_key.dot(__A				)				).T.tolist()[0]
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=''''''.join(
                         self.replace_digits(__A				) for num in batch_decrypted				)
                     decrypted += decrypted_batch

              return decrypted






def 							SCREAMING_SNAKE_CASE_	(					)			->						None:
       SCREAMING_SNAKE_CASE_							:				List[str]		=int(input('''Enter the order of the encryption key: '''				)				)
       SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=[]

       print('''Enter each row of the encryption key with space separated integers'''				)
       for _ in range(UpperCAmelCase_				):
              SCREAMING_SNAKE_CASE_							:				Tuple		=[int(UpperCAmelCase_				) for x in input().split()]
              hill_matrix.append(UpperCAmelCase_				)

       SCREAMING_SNAKE_CASE_							:				Optional[int]		=HillCipher(numpy.array(UpperCAmelCase_				)				)

       print('''Would you like to encrypt or decrypt some text? (1 or 2)'''				)
       SCREAMING_SNAKE_CASE_							:				Dict		=input('''\n1. Encrypt\n2. Decrypt\n'''				)
       if option == "1":
              SCREAMING_SNAKE_CASE_							:				List[Any]		=input('''What text would you like to encrypt?: '''				)
              print('''Your encrypted text is:'''				)
              print(hc.encrypt(UpperCAmelCase_				)				)
       elif option == "2":
              SCREAMING_SNAKE_CASE_							:				Any		=input('''What text would you like to decrypt?: '''				)
              print('''Your decrypted text is:'''				)
              print(hc.decrypt(UpperCAmelCase_				)				)


if __name__ == "__main__":
 import doctest

 doctest.testmod()

 main()
