

"""simple docstring"""
import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
    import torch

class 			lowerCamelCase  :


  '''simple docstring'''







  def __init__(self       ,					_lowerCamelCase       ,					_lowerCamelCase=14       ,					_lowerCamelCase=7       ,					_lowerCamelCase=True       ,					_lowerCamelCase=True       ,					_lowerCamelCase=False       ,					_lowerCamelCase=True       ,					_lowerCamelCase=99       ,					_lowerCamelCase=32       ,					_lowerCamelCase=5       ,					_lowerCamelCase=4       ,					_lowerCamelCase=37       ,					_lowerCamelCase="gelu"       ,					_lowerCamelCase=0.1       ,					_lowerCamelCase=0.1       ,					_lowerCamelCase=512       ,					_lowerCamelCase=0.02       ,					):


        """simple docstring"""

        UpperCAmelCase__ :       str						=						parent
        UpperCAmelCase__ :       str						=						batch_size
        UpperCAmelCase__ :       Dict						=						seq_length
        UpperCAmelCase__ :       Tuple						=						is_training
        UpperCAmelCase__ :       int						=						use_input_mask
        UpperCAmelCase__ :       Optional[Any]						=						use_token_type_ids
        UpperCAmelCase__ :       Union[str, Any]						=						use_labels
        UpperCAmelCase__ :       List[str]						=						vocab_size
        UpperCAmelCase__ :       Union[str, Any]						=						hidden_size
        UpperCAmelCase__ :       Optional[Any]						=						num_hidden_layers
        UpperCAmelCase__ :       Optional[int]						=						num_attention_heads
        UpperCAmelCase__ :       List[str]						=						intermediate_size
        UpperCAmelCase__ :       List[str]						=						hidden_act
        UpperCAmelCase__ :       int						=						hidden_dropout_prob
        UpperCAmelCase__ :       Dict						=						attention_probs_dropout_prob
        UpperCAmelCase__ :       Tuple						=						max_position_embeddings
        UpperCAmelCase__ :       Optional[int]						=						initializer_range
        UpperCAmelCase__ :       int						=						None
        UpperCAmelCase__ :       Any						=						vocab_size - 1
        UpperCAmelCase__ :       List[Any]						=						vocab_size - 1
        UpperCAmelCase__ :       Any						=						vocab_size - 1



  def    _a			(self  ):


        """simple docstring"""

        UpperCAmelCase__ :       List[str]						=						ids_tensor([self.batch_size, self.seq_length]       ,					self.vocab_size  )

        UpperCAmelCase__ :       Tuple						=						None
        if self.use_input_mask:
              UpperCAmelCase__ :       int						=						random_attention_mask([self.batch_size, self.seq_length]  )

        UpperCAmelCase__ :       List[str]						=						GPTaConfig(
            vocab_size=self.vocab_size       ,					n_embd=self.hidden_size       ,					n_layer=self.num_hidden_layers       ,					n_head=self.num_attention_heads       ,					n_positions=self.max_position_embeddings       ,					use_cache=_lowerCamelCase       ,					bos_token_id=self.bos_token_id       ,					eos_token_id=self.eos_token_id       ,					pad_token_id=self.pad_token_id       ,					)

        return (config, input_ids, input_mask)



  def    _a			(self  ):


        """simple docstring"""

        UpperCAmelCase__ :       Dict						=						self.prepare_config_and_inputs()
        UpperCAmelCase__				,		UpperCAmelCase__				,		UpperCAmelCase__ :       Dict						=						config_and_inputs
        UpperCAmelCase__ :       Tuple						=						{"""input_ids""": input_ids, """attention_mask""": attention_mask}
        return config, inputs_dict



  def    _a			(self  ):


        """simple docstring"""

        UpperCAmelCase__				,		UpperCAmelCase__				,		UpperCAmelCase__ :       Optional[int]						=						self.prepare_config_and_inputs()

        UpperCAmelCase__ :       Optional[int]						=						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
        UpperCAmelCase__ :       List[str]						=						ids_tensor([self.batch_size, self.seq_length]       ,					vocab_size=2  )

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )



  def    _a			(self       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase  ):


        """simple docstring"""

        UpperCAmelCase__ :       List[str]						=						20
        UpperCAmelCase__ :       Any						=						model_class_name(_lowerCamelCase  )

        UpperCAmelCase__ :       int						=						model.init_cache(input_ids.shape[0]       ,					_lowerCamelCase  )
        UpperCAmelCase__ :       Optional[Any]						=						jnp.ones((input_ids.shape[0], max_decoder_length)       ,					dtype="""i4"""  )

        UpperCAmelCase__ :       int						=						jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1  )[None, :]       ,					(input_ids.shape[0], input_ids.shape[-1] - 1)  )
        UpperCAmelCase__ :       Optional[Any]						=						model(
            input_ids[:, :-1]       ,					attention_mask=_lowerCamelCase       ,					past_key_values=_lowerCamelCase       ,					position_ids=_lowerCamelCase       ,					)

        UpperCAmelCase__ :       Union[str, Any]						=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,					dtype="""i4"""  )
        UpperCAmelCase__ :       List[str]						=						model(
            input_ids[:, -1:]       ,					attention_mask=_lowerCamelCase       ,					past_key_values=outputs_cache.past_key_values       ,					position_ids=_lowerCamelCase       ,					)

        UpperCAmelCase__ :       List[str]						=						model(_lowerCamelCase  )

        UpperCAmelCase__ :       Optional[int]						=						np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])  )  )
        self.parent.assertTrue(diff < 1e-3       ,					msg=F"""Max diff is {diff}"""  )




  def    _a			(self       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase  ):


        """simple docstring"""

        UpperCAmelCase__ :       str						=						20
        UpperCAmelCase__ :       Tuple						=						model_class_name(_lowerCamelCase  )

        UpperCAmelCase__ :       Any						=						jnp.concatenate(
            [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])  )]       ,					axis=-1       ,					)

        UpperCAmelCase__ :       Tuple						=						model.init_cache(input_ids.shape[0]       ,					_lowerCamelCase  )
        UpperCAmelCase__ :       List[Any]						=						jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1  )[None, :]       ,					(input_ids.shape[0], input_ids.shape[-1] - 1)  )

        UpperCAmelCase__ :       Union[str, Any]						=						model(
            input_ids[:, :-1]       ,					attention_mask=_lowerCamelCase       ,					past_key_values=_lowerCamelCase       ,					position_ids=_lowerCamelCase       ,					)
        UpperCAmelCase__ :       Any						=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,					dtype="""i4"""  )
        UpperCAmelCase__ :       List[str]						=						model(
            input_ids[:, -1:]       ,					past_key_values=outputs_cache.past_key_values       ,					attention_mask=_lowerCamelCase       ,					position_ids=_lowerCamelCase       ,					)

        UpperCAmelCase__ :       Optional[Any]						=						model(_lowerCamelCase       ,					attention_mask=_lowerCamelCase  )

        UpperCAmelCase__ :       Optional[int]						=						np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])  )  )
        self.parent.assertTrue(diff < 1e-3       ,					msg=F"""Max diff is {diff}"""  )




@require_flax
class 			lowerCamelCase  (					lowerCAmelCase__				,				lowerCAmelCase__				,				unittest.TestCase      ):


  '''simple docstring'''







  SCREAMING_SNAKE_CASE							=     (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
  SCREAMING_SNAKE_CASE							=     (FlaxGPTaLMHeadModel,) if is_flax_available() else ()



  def    _a			(self  ):


        """simple docstring"""

        UpperCAmelCase__ :       int						=						FlaxGPTaModelTester(self  )



  def    _a			(self  ):


        """simple docstring"""

        for model_class_name in self.all_model_classes:
              UpperCAmelCase__				,		UpperCAmelCase__				,		UpperCAmelCase__ :       Optional[Any]						=						self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward(_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase  )



  def    _a			(self  ):


        """simple docstring"""

        for model_class_name in self.all_model_classes:
              UpperCAmelCase__				,		UpperCAmelCase__				,		UpperCAmelCase__ :       Tuple						=						self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward_with_attn_mask(
                  _lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase  )



  @slow
  def    _a			(self  ):


        """simple docstring"""

        UpperCAmelCase__ :       Dict						=						GPTaTokenizer.from_pretrained("""gpt2"""       ,					pad_token="""</s>"""       ,					padding_side="""left"""  )
        UpperCAmelCase__ :       Optional[Any]						=						tokenizer(["""Hello this is a long string""", """Hey"""]       ,					return_tensors="""np"""       ,					padding=_lowerCamelCase       ,					truncation=_lowerCamelCase  )

        UpperCAmelCase__ :       List[str]						=						FlaxGPTaLMHeadModel.from_pretrained("""gpt2"""  )
        UpperCAmelCase__ :       Optional[int]						=						False
        UpperCAmelCase__ :       List[Any]						=						model.config.eos_token_id

        UpperCAmelCase__ :       Optional[int]						=						jax.jit(model.generate  )

        UpperCAmelCase__ :       Any						=						jit_generate(inputs["""input_ids"""]       ,					attention_mask=inputs["""attention_mask"""]  ).sequences

        UpperCAmelCase__ :       int						=						tokenizer.batch_decode(_lowerCamelCase       ,					skip_special_tokens=_lowerCamelCase  )

        UpperCAmelCase__ :       str						=						[
            """Hello this is a long string of words. I'm going to start with the first one.\n""",
            """Hey, I'm not sure if I'm going to be able to do""",
        ]

        self.assertListEqual(_lowerCamelCase       ,					_lowerCamelCase  )



  @is_pt_flax_cross_test
  def    _a			(self  ):


        """simple docstring"""

        UpperCAmelCase__				,		UpperCAmelCase__ :       Optional[int]						=						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__  ):
                    # prepare inputs
                    UpperCAmelCase__ :       Dict						=						self._prepare_for_class(_lowerCamelCase       ,					_lowerCamelCase  )
                    UpperCAmelCase__ :       Tuple						=						{k: torch.tensor(v.tolist()  ) for k, v in prepared_inputs_dict.items()}

                    # load corresponding PyTorch class
                    UpperCAmelCase__ :       Dict						=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
                    UpperCAmelCase__ :       List[str]						=						getattr(_lowerCamelCase       ,					_lowerCamelCase  )

                    UpperCAmelCase__				,		UpperCAmelCase__ :       Union[str, Any]						=						pt_inputs["""input_ids"""].shape
                    UpperCAmelCase__ :       Dict						=						np.random.randint(0       ,					seq_length - 1       ,					size=(batch_size,)  )
                    for batch_idx, start_index in enumerate(_lowerCamelCase  ):
                          UpperCAmelCase__ :       List[str]						=						0
                          UpperCAmelCase__ :       Dict						=						1
                          UpperCAmelCase__ :       Any						=						0
                          UpperCAmelCase__ :       List[Any]						=						1
                    UpperCAmelCase__ :       int						=						pt_model_class(_lowerCamelCase  ).eval()
                    UpperCAmelCase__ :       Union[str, Any]						=						model_class(_lowerCamelCase       ,					dtype=jnp.floataa  )

                    UpperCAmelCase__ :       str						=						convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,					_lowerCamelCase  )
                    UpperCAmelCase__ :       List[str]						=						fx_state

                    with torch.no_grad():
                          UpperCAmelCase__ :       str						=						pt_model(**_lowerCamelCase  ).to_tuple()

                    UpperCAmelCase__ :       Tuple						=						fx_model(**_lowerCamelCase  ).to_tuple()
                    self.assertEqual(len(_lowerCamelCase  )       ,					len(_lowerCamelCase  )       ,					"""Output lengths differ between Flax and PyTorch"""  )
                    for fx_output, pt_output in zip(_lowerCamelCase       ,					_lowerCamelCase  ):
                          self.assert_almost_equals(fx_output[:, -1]       ,					pt_output[:, -1].numpy()       ,					4e-2  )

                    with tempfile.TemporaryDirectory() as tmpdirname:
                          pt_model.save_pretrained(_lowerCamelCase  )
                          UpperCAmelCase__ :       Optional[int]						=						model_class.from_pretrained(_lowerCamelCase       ,					from_pt=_lowerCamelCase  )

                    UpperCAmelCase__ :       Tuple						=						fx_model_loaded(**_lowerCamelCase  ).to_tuple()
                    self.assertEqual(
                        len(_lowerCamelCase  )       ,					len(_lowerCamelCase  )       ,					"""Output lengths differ between Flax and PyTorch"""  )
                    for fx_output_loaded, pt_output in zip(_lowerCamelCase       ,					_lowerCamelCase  ):
                          self.assert_almost_equals(fx_output_loaded[:, -1]       ,					pt_output[:, -1].numpy()       ,					4e-2  )



  @is_pt_flax_cross_test
  def    _a			(self  ):


        """simple docstring"""

        UpperCAmelCase__				,		UpperCAmelCase__ :       int						=						self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__  ):
                    # prepare inputs
                    UpperCAmelCase__ :       str						=						self._prepare_for_class(_lowerCamelCase       ,					_lowerCamelCase  )
                    UpperCAmelCase__ :       Dict						=						{k: torch.tensor(v.tolist()  ) for k, v in prepared_inputs_dict.items()}

                    # load corresponding PyTorch class
                    UpperCAmelCase__ :       Optional[Any]						=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
                    UpperCAmelCase__ :       Dict						=						getattr(_lowerCamelCase       ,					_lowerCamelCase  )

                    UpperCAmelCase__ :       str						=						pt_model_class(_lowerCamelCase  ).eval()
                    UpperCAmelCase__ :       Dict						=						model_class(_lowerCamelCase       ,					dtype=jnp.floataa  )

                    UpperCAmelCase__ :       int						=						load_flax_weights_in_pytorch_model(_lowerCamelCase       ,					fx_model.params  )
                    UpperCAmelCase__				,		UpperCAmelCase__ :       List[str]						=						pt_inputs["""input_ids"""].shape
                    UpperCAmelCase__ :       int						=						np.random.randint(0       ,					seq_length - 1       ,					size=(batch_size,)  )
                    for batch_idx, start_index in enumerate(_lowerCamelCase  ):
                          UpperCAmelCase__ :       Any						=						0
                          UpperCAmelCase__ :       Dict						=						1
                          UpperCAmelCase__ :       Union[str, Any]						=						0
                          UpperCAmelCase__ :       Optional[Any]						=						1

                    # make sure weights are tied in PyTorch
                    pt_model.tie_weights()

                    with torch.no_grad():
                          UpperCAmelCase__ :       Union[str, Any]						=						pt_model(**_lowerCamelCase  ).to_tuple()

                    UpperCAmelCase__ :       Union[str, Any]						=						fx_model(**_lowerCamelCase  ).to_tuple()
                    self.assertEqual(len(_lowerCamelCase  )       ,					len(_lowerCamelCase  )       ,					"""Output lengths differ between Flax and PyTorch"""  )
                    for fx_output, pt_output in zip(_lowerCamelCase       ,					_lowerCamelCase  ):
                          self.assert_almost_equals(fx_output[:, -1]       ,					pt_output[:, -1].numpy()       ,					4e-2  )

                    with tempfile.TemporaryDirectory() as tmpdirname:
                          fx_model.save_pretrained(_lowerCamelCase  )
                          UpperCAmelCase__ :       Optional[int]						=						pt_model_class.from_pretrained(_lowerCamelCase       ,					from_flax=_lowerCamelCase  )

                    with torch.no_grad():
                          UpperCAmelCase__ :       int						=						pt_model_loaded(**_lowerCamelCase  ).to_tuple()

                    self.assertEqual(
                        len(_lowerCamelCase  )       ,					len(_lowerCamelCase  )       ,					"""Output lengths differ between Flax and PyTorch"""  )
                    for fx_output, pt_output in zip(_lowerCamelCase       ,					_lowerCamelCase  ):
                          self.assert_almost_equals(fx_output[:, -1]       ,					pt_output[:, -1].numpy()       ,					4e-2  )




  @slow
  def    _a			(self  ):


        """simple docstring"""

        for model_class_name in self.all_model_classes:
              UpperCAmelCase__ :       Optional[int]						=						model_class_name.from_pretrained("""gpt2"""       ,					from_pt=_lowerCamelCase  )
              UpperCAmelCase__ :       List[Any]						=						model(np.ones((1, 1)  )  )
              self.assertIsNotNone(_lowerCamelCase  )




"""simple docstring"""
def 				a__	(	lowerCAmelCase ,  lowerCAmelCase						)  ->   str:
      if not isinstance(lowerCAmelCase ,  lowerCAmelCase						):
            raise TypeError("""plaintext must be a string"""						)
      if not isinstance(lowerCAmelCase ,  lowerCAmelCase						):
            raise TypeError("""key must be a string"""						)

      if not plaintext:
            raise ValueError("""plaintext is empty"""						)
      if not key:
            raise ValueError("""key is empty"""						)

      key += plaintext
      UpperCAmelCase__ :       Dict						=						plaintext.lower()
      UpperCAmelCase__ :       Tuple						=						key.lower()
      UpperCAmelCase__ :       Optional[int]						=						0
      UpperCAmelCase__ :       Tuple						=						0
      UpperCAmelCase__ :       Optional[int]						=						""""""
      while plaintext_iterator < len(lowerCAmelCase						):
            if (
                ord(plaintext[plaintext_iterator]						) < 97
                or ord(plaintext[plaintext_iterator]						) > 1_22
            ):
                  ciphertext += plaintext[plaintext_iterator]
                  plaintext_iterator += 1
            elif ord(key[key_iterator]						) < 97 or ord(key[key_iterator]						) > 1_22:
                  key_iterator += 1
            else:
                  ciphertext += chr(
                      (
                          (ord(plaintext[plaintext_iterator]						) - 97 + ord(key[key_iterator]						))
                          - 97
                      )
                      % 26
                      + 97						)
                  key_iterator += 1
                  plaintext_iterator += 1
      return ciphertext




def 				a__	(	lowerCAmelCase ,  lowerCAmelCase						)  ->   str:
      if not isinstance(lowerCAmelCase ,  lowerCAmelCase						):
            raise TypeError("""ciphertext must be a string"""						)
      if not isinstance(lowerCAmelCase ,  lowerCAmelCase						):
            raise TypeError("""key must be a string"""						)

      if not ciphertext:
            raise ValueError("""ciphertext is empty"""						)
      if not key:
            raise ValueError("""key is empty"""						)

      UpperCAmelCase__ :       Union[str, Any]						=						key.lower()
      UpperCAmelCase__ :       int						=						0
      UpperCAmelCase__ :       int						=						0
      UpperCAmelCase__ :       int						=						""""""
      while ciphertext_iterator < len(lowerCAmelCase						):
            if (
                ord(ciphertext[ciphertext_iterator]						) < 97
                or ord(ciphertext[ciphertext_iterator]						) > 1_22
            ):
                  plaintext += ciphertext[ciphertext_iterator]
            else:
                  plaintext += chr(
                      (ord(ciphertext[ciphertext_iterator]						) - ord(key[key_iterator]						)) % 26
                      + 97						)
                  key += chr(
                      (ord(ciphertext[ciphertext_iterator]						) - ord(key[key_iterator]						)) % 26
                      + 97						)
                  key_iterator += 1
            ciphertext_iterator += 1
      return plaintext


if __name__ == "__main__":
    import doctest

    doctest.testmod()
    _A						=     int(input("""Type 1 to encrypt or 2 to decrypt:"""))
    if operation == 1:
        _A						=     input("""Typeplaintext to be encrypted:\n""")
        _A						=     input("""Type the key:\n""")
        print(encrypt(plaintext, key))
    elif operation == 2:
        _A						=     input("""Type the ciphertext to be decrypted:\n""")
        _A						=     input("""Type the key:\n""")
        print(decrypt(ciphertext, key))
    decrypt("""jsqqs avvwo""", """coffee""")

