



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       _snake_case				:






   '''simple docstring'''






   def __init__(    self:     int  ,lowerCamelCase_:     Any  ,lowerCamelCase_:     Tuple=14  ,lowerCamelCase_:     List[str]=7  ,lowerCamelCase_:     Optional[Any]=True  ,lowerCamelCase_:     Optional[Any]=True  ,lowerCamelCase_:     str=False  ,lowerCamelCase_:     str=True  ,lowerCamelCase_:     List[Any]=99  ,lowerCamelCase_:     List[Any]=32  ,lowerCamelCase_:     Optional[int]=5  ,lowerCamelCase_:     Optional[Any]=4  ,lowerCamelCase_:     List[Any]=37  ,lowerCamelCase_:     Optional[Any]="gelu"  ,lowerCamelCase_:     List[str]=0.1  ,lowerCamelCase_:     str=0.1  ,lowerCamelCase_:     str=512  ,lowerCamelCase_:     Any=0.0_2  ,)	->      Any:
       UpperCAmelCase_     :       List[Any]  		=					parent
       UpperCAmelCase_     :       str  		=					batch_size
       UpperCAmelCase_     :       Union[str, Any]  		=					seq_length
       UpperCAmelCase_     :       Tuple  		=					is_training
       UpperCAmelCase_     :       Tuple  		=					use_input_mask
       UpperCAmelCase_     :       Any  		=					use_token_type_ids
       UpperCAmelCase_     :       Optional[Any]  		=					use_labels
       UpperCAmelCase_     :       List[str]  		=					vocab_size
       UpperCAmelCase_     :       Optional[int]  		=					hidden_size
       UpperCAmelCase_     :       Optional[Any]  		=					num_hidden_layers
       UpperCAmelCase_     :       Dict  		=					num_attention_heads
       UpperCAmelCase_     :       int  		=					intermediate_size
       UpperCAmelCase_     :       List[str]  		=					hidden_act
       UpperCAmelCase_     :       Tuple  		=					hidden_dropout_prob
       UpperCAmelCase_     :       str  		=					attention_probs_dropout_prob
       UpperCAmelCase_     :       Tuple  		=					max_position_embeddings
       UpperCAmelCase_     :       str  		=					initializer_range
       UpperCAmelCase_     :       Optional[int]  		=					None
       UpperCAmelCase_     :       str  		=					vocab_size - 1
       UpperCAmelCase_     :       Any  		=					vocab_size - 1
       UpperCAmelCase_     :       Dict  		=					vocab_size - 1


   def       A__    (    self:     str							)	->      Any:
       UpperCAmelCase_     :       str  		=					ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size							)

       UpperCAmelCase_     :       List[str]  		=					None
       if self.use_input_mask:
           UpperCAmelCase_     :       Any  		=					random_attention_mask([self.batch_size, self.seq_length]							)

       UpperCAmelCase_     :       Tuple  		=					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:     Any							)	->      Tuple:
       UpperCAmelCase_     :       Any  		=					self.prepare_config_and_inputs()
       UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_     :       Dict  		=					config_and_inputs
       UpperCAmelCase_     :       Union[str, Any]  		=					{"""input_ids""": input_ids, """attention_mask""": attention_mask}
       return config, inputs_dict


   def       A__    (    self:     Tuple							)	->      Any:
       UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_     :       Optional[Any]  		=					self.prepare_config_and_inputs()

       UpperCAmelCase_     :       List[Any]  		=					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
       UpperCAmelCase_     :       Optional[Any]  		=					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:     Tuple  ,lowerCamelCase_:     Union[str, Any]  ,lowerCamelCase_:     Any  ,lowerCamelCase_:     Any  ,lowerCamelCase_:     List[str]							)	->      Tuple:
       UpperCAmelCase_     :       List[Any]  		=					20
       UpperCAmelCase_     :       Union[str, Any]  		=					model_class_name(lowerCamelCase_							)

       UpperCAmelCase_     :       Optional[int]  		=					model.init_cache(input_ids.shape[0]  ,lowerCamelCase_							)
       UpperCAmelCase_     :       Tuple  		=					jnp.ones((input_ids.shape[0], max_decoder_length)  ,dtype="""i4"""							)

       UpperCAmelCase_     :       Tuple  		=					jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1							)[None, :]  ,(input_ids.shape[0], input_ids.shape[-1] - 1)							)
       UpperCAmelCase_     :       List[str]  		=					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_     :       Tuple  		=					model(
           input_ids[:, -1:]  ,attention_mask=lowerCamelCase_  ,past_key_values=outputs_cache.past_key_values  ,position_ids=lowerCamelCase_  ,)

       UpperCAmelCase_     :       Tuple  		=					model(lowerCamelCase_							)

       UpperCAmelCase_     :       List[str]  		=					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:     List[Any]  ,lowerCamelCase_:     Dict  ,lowerCamelCase_:     Tuple  ,lowerCamelCase_:     List[Any]  ,lowerCamelCase_:     List[Any]							)	->      Union[str, Any]:
       UpperCAmelCase_     :       List[Any]  		=					20
       UpperCAmelCase_     :       Union[str, Any]  		=					model_class_name(lowerCamelCase_							)

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

       UpperCAmelCase_     :       Optional[int]  		=					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_     :       Optional[int]  		=					jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]  ,dtype="""i4"""							)
       UpperCAmelCase_     :       Optional[Any]  		=					model(
           input_ids[:, -1:]  ,past_key_values=outputs_cache.past_key_values  ,attention_mask=lowerCamelCase_  ,position_ids=lowerCamelCase_  ,)

       UpperCAmelCase_     :       Any  		=					model(lowerCamelCase_  ,attention_mask=lowerCamelCase_							)

       UpperCAmelCase_     :       str  		=					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       _snake_case				(     __snake_case ,							__snake_case ,							unittest.TestCase ):






   '''simple docstring'''






   A__      : Tuple       =						(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
   A__      : Optional[Any]       =						(FlaxGPTaLMHeadModel,) if is_flax_available() else ()


   def       A__    (    self:     str							)	->      List[Any]:
       UpperCAmelCase_     :       List[Any]  		=					FlaxGPTaModelTester(self							)


   def       A__    (    self:     List[Any]							)	->      Union[str, Any]:
       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:     List[str]							)	->      List[Any]:
       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:     List[Any]							)	->      Optional[Any]:
       UpperCAmelCase_     :       List[str]  		=					GPTaTokenizer.from_pretrained("""gpt2"""  ,pad_token="""</s>"""  ,padding_side="""left"""							)
       UpperCAmelCase_     :       List[str]  		=					tokenizer(["""Hello this is a long string""", """Hey"""]  ,return_tensors="""np"""  ,padding=lowerCamelCase_  ,truncation=lowerCamelCase_							)

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

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

       UpperCAmelCase_     :       Union[str, Any]  		=					jit_generate(inputs["""input_ids"""]  ,attention_mask=inputs["""attention_mask"""]							).sequences

       UpperCAmelCase_     :       Optional[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:     Optional[int]							)	->      Union[str, Any]:
       UpperCAmelCase_		, UpperCAmelCase_     :       Tuple  		=					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_     :       Optional[Any]  		=					self._prepare_for_class(lowerCamelCase_  ,lowerCamelCase_							)
               UpperCAmelCase_     :       List[str]  		=					{k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

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

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

               UpperCAmelCase_     :       Union[str, Any]  		=					convert_pytorch_state_dict_to_flax(pt_model.state_dict()  ,lowerCamelCase_							)
               UpperCAmelCase_     :       Any  		=					fx_state

               with torch.no_grad():
                   UpperCAmelCase_     :       Optional[int]  		=					pt_model(**lowerCamelCase_							).to_tuple()

               UpperCAmelCase_     :       List[str]  		=					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_     :       Union[str, Any]  		=					model_class.from_pretrained(lowerCamelCase_  ,from_pt=lowerCamelCase_							)

               UpperCAmelCase_     :       Any  		=					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:     Optional[int]							)	->      Any:
       UpperCAmelCase_		, UpperCAmelCase_     :       Union[str, Any]  		=					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_     :       Any  		=					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_     :       Dict  		=					model_class.__name__[4:]  # Skip the "Flax" at the beginning
               UpperCAmelCase_     :       Dict  		=					getattr(lowerCamelCase_  ,lowerCamelCase_							)

               UpperCAmelCase_     :       Union[str, Any]  		=					pt_model_class(lowerCamelCase_							).eval()
               UpperCAmelCase_     :       Dict  		=					model_class(lowerCamelCase_  ,dtype=jnp.floataa							)

               UpperCAmelCase_     :       List[Any]  		=					load_flax_weights_in_pytorch_model(lowerCamelCase_  ,fx_model.params							)
               UpperCAmelCase_		, UpperCAmelCase_     :       int  		=					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_     :       Tuple  		=					0
                   UpperCAmelCase_     :       Optional[Any]  		=					1
                   UpperCAmelCase_     :       Dict  		=					0
                   UpperCAmelCase_     :       Any  		=					1

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

               with torch.no_grad():
                   UpperCAmelCase_     :       str  		=					pt_model(**lowerCamelCase_							).to_tuple()

               UpperCAmelCase_     :       Optional[int]  		=					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[Any]  		=					pt_model_class.from_pretrained(lowerCamelCase_  ,from_flax=lowerCamelCase_							)

               with torch.no_grad():
                   UpperCAmelCase_     :       Any  		=					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:     Optional[Any]							)	->      List[Any]:
       for model_class_name in self.all_model_classes:
           UpperCAmelCase_     :       int  		=					model_class_name.from_pretrained("""gpt2"""  ,from_pt=lowerCamelCase_							)
           UpperCAmelCase_     :       Dict  		=					model(np.ones((1, 1)							)							)
           self.assertIsNotNone(lowerCamelCase_							)





import itertools
from pathlib import Path







def       lowerCamelCase_  (    _a       :						list[str]    ):



    '''simple docstring'''




    UpperCAmelCase_     :       Union[str, Any]  		=					[tuple(_a    ) for login in logins]

    UpperCAmelCase_     :       str  		=					{char for login in split_logins for char in login}

    for permutation in itertools.permutations(_a    ):
        UpperCAmelCase_     :       str  		=					True
        for login in logins:
            if not (
                permutation.index(login[0]    )
                < permutation.index(login[1]    )
                < permutation.index(login[2]    )
            ):
                UpperCAmelCase_     :       Optional[int]  		=					False
                break

        if satisfied:
            return int("""""".join(_a    )    )

    raise Exception("""Unable to find the secret passcode"""    )







def       lowerCamelCase_  (    _a       :						str = "keylog.txt"    ):



    '''simple docstring'''




    UpperCAmelCase_     :       Union[str, Any]  		=					Path(_a    ).parent.joinpath(_a    ).read_text().splitlines()

    return find_secret_passcode(_a    )


if __name__ == "__main__":
      print(F"{solution() = }")
