




import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, 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.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
    import torch







class __snake_case			:







   def __init__(      self		:						int					,						_snake_case		:						Tuple					,						_snake_case		:						Optional[int]=14					,						_snake_case		:						List[str]=7					,						_snake_case		:						int=True					,						_snake_case		:						str=True					,						_snake_case		:						str=False					,						_snake_case		:						List[str]=True					,						_snake_case		:						List[str]=99					,						_snake_case		:						Optional[Any]=32					,						_snake_case		:						Optional[Any]=4					,						_snake_case		:						Optional[Any]=4					,						_snake_case		:						Dict=4					,						_snake_case		:						Optional[Any]=37					,						_snake_case		:						Union[str, Any]="gelu"					,						_snake_case		:						Optional[int]=0.1					,						_snake_case		:						str=0.1					,						_snake_case		:						Optional[int]=512					,						_snake_case		:						int=0.0_2					,						):

        """simple docstring"""


        UpperCAmelCase_          =							parent
        UpperCAmelCase_          =							batch_size
        UpperCAmelCase_          =							seq_length
        UpperCAmelCase_          =							is_training
        UpperCAmelCase_          =							use_input_mask
        UpperCAmelCase_          =							use_token_type_ids
        UpperCAmelCase_          =							use_labels
        UpperCAmelCase_          =							vocab_size
        UpperCAmelCase_          =							hidden_size
        UpperCAmelCase_          =							rotary_dim
        UpperCAmelCase_          =							num_hidden_layers
        UpperCAmelCase_          =							num_attention_heads
        UpperCAmelCase_          =							intermediate_size
        UpperCAmelCase_          =							hidden_act
        UpperCAmelCase_          =							hidden_dropout_prob
        UpperCAmelCase_          =							attention_probs_dropout_prob
        UpperCAmelCase_          =							max_position_embeddings
        UpperCAmelCase_          =							initializer_range
        UpperCAmelCase_          =							None
        UpperCAmelCase_          =							vocab_size - 1
        UpperCAmelCase_          =							vocab_size - 1
        UpperCAmelCase_          =							vocab_size - 1







   def lowerCamelCase						(      self		:						Tuple):

        """simple docstring"""


        UpperCAmelCase_          =							ids_tensor([self.batch_size, self.seq_length]					,						self.vocab_size)

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

        UpperCAmelCase_          =							GPTJConfig(
            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=_snake_case					,						bos_token_id=self.bos_token_id					,						eos_token_id=self.eos_token_id					,						pad_token_id=self.pad_token_id					,						rotary_dim=self.rotary_dim					,						)

        return (config, input_ids, input_mask)







   def lowerCamelCase						(      self		:						List[Any]):

        """simple docstring"""


        UpperCAmelCase_          =							self.prepare_config_and_inputs()
        UpperCAmelCase_     ,       UpperCAmelCase_     ,       UpperCAmelCase_          =							config_and_inputs
        UpperCAmelCase_          =							{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
        return config, inputs_dict







   def lowerCamelCase						(      self		:						List[Any]					,						_snake_case		:						str					,						_snake_case		:						str					,						_snake_case		:						Dict					,						_snake_case		:						List[Any]):

        """simple docstring"""


        UpperCAmelCase_          =							20
        UpperCAmelCase_          =							model_class_name(_snake_case)

        UpperCAmelCase_          =							model.init_cache(input_ids.shape[0]					,						_snake_case)
        UpperCAmelCase_          =							jnp.ones((input_ids.shape[0], max_decoder_length)					,						dtype='''i4''')

        UpperCAmelCase_          =							jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1)[None, :]					,						(input_ids.shape[0], input_ids.shape[-1] - 1))
        UpperCAmelCase_          =							model(
            input_ids[:, :-1]					,						attention_mask=_snake_case					,						past_key_values=_snake_case					,						position_ids=_snake_case					,						)

        UpperCAmelCase_          =							jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,						dtype='''i4''')
        UpperCAmelCase_          =							model(
            input_ids[:, -1:]					,						attention_mask=_snake_case					,						past_key_values=outputs_cache.past_key_values					,						position_ids=_snake_case					,						)

        UpperCAmelCase_          =							model(_snake_case)

        UpperCAmelCase_          =							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 lowerCamelCase						(      self		:						Tuple					,						_snake_case		:						Any					,						_snake_case		:						List[Any]					,						_snake_case		:						Optional[Any]					,						_snake_case		:						List[Any]):

        """simple docstring"""


        UpperCAmelCase_          =							20
        UpperCAmelCase_          =							model_class_name(_snake_case)

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

        UpperCAmelCase_          =							model.init_cache(input_ids.shape[0]					,						_snake_case)
        UpperCAmelCase_          =							jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1)[None, :]					,						(input_ids.shape[0], input_ids.shape[-1] - 1))

        UpperCAmelCase_          =							model(
            input_ids[:, :-1]					,						attention_mask=_snake_case					,						past_key_values=_snake_case					,						position_ids=_snake_case					,						)
        UpperCAmelCase_          =							jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,						dtype='''i4''')
        UpperCAmelCase_          =							model(
            input_ids[:, -1:]					,						past_key_values=outputs_cache.past_key_values					,						attention_mask=_snake_case					,						position_ids=_snake_case					,						)

        UpperCAmelCase_          =							model(_snake_case					,						attention_mask=_snake_case)

        UpperCAmelCase_          =							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			(			a ,  a ,  unittest.TestCase		):
   UpperCAmelCase__  :    str          =  (FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
   UpperCAmelCase__  :    List[Any]          =  (FlaxGPTJForCausalLM,) if is_flax_available() else ()







   def lowerCamelCase						(      self		:						str):

        """simple docstring"""


        UpperCAmelCase_          =							FlaxGPTJModelTester(self)







   def lowerCamelCase						(      self		:						Tuple):

        """simple docstring"""


        for model_class_name in self.all_model_classes:
             UpperCAmelCase_     ,       UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward(_snake_case					,						_snake_case					,						_snake_case					,						_snake_case)







   def lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


        for model_class_name in self.all_model_classes:
             UpperCAmelCase_     ,       UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward_with_attn_mask(
                 _snake_case					,						_snake_case					,						_snake_case					,						_snake_case)







   @tooslow
   def lowerCamelCase						(      self		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							GPTaTokenizer.from_pretrained('''gpt2'''					,						pad_token='''<|endoftext|>'''					,						padding_side='''left''')
        UpperCAmelCase_          =							tokenizer(['''Hello this is a long string''', '''Hey''']					,						return_tensors='''np'''					,						padding=_snake_case					,						truncation=_snake_case)

        UpperCAmelCase_          =							FlaxGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B''')
        UpperCAmelCase_          =							False
        UpperCAmelCase_          =							model.config.eos_token_id

        UpperCAmelCase_          =							jax.jit(model.generate)

        UpperCAmelCase_          =							jit_generate(
            inputs['''input_ids''']					,						attention_mask=inputs['''attention_mask''']					,						pad_token_id=tokenizer.pad_token_id).sequences

        UpperCAmelCase_          =							tokenizer.batch_decode(_snake_case					,						skip_special_tokens=_snake_case)

        UpperCAmelCase_          =							[
            '''Hello this is a long string of text.\n\nI\'m trying to get the text of the''',
            '''Hey, I\'m a little late to the party. I\'m going to''',
        ]

        self.assertListEqual(_snake_case					,						_snake_case)







   @is_pt_flax_cross_test
   def lowerCamelCase						(      self		:						Union[str, Any]):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							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_          =							self._prepare_for_class(_snake_case					,						_snake_case)
                  UpperCAmelCase_          =							{k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

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

                  UpperCAmelCase_     ,       UpperCAmelCase_          =							pt_inputs['''input_ids'''].shape
                  UpperCAmelCase_          =							np.random.randint(0					,						seq_length - 1					,						size=(batch_size,))
                  for batch_idx, start_index in enumerate(_snake_case):
                       UpperCAmelCase_          =							0
                       UpperCAmelCase_          =							1
                       UpperCAmelCase_          =							0
                       UpperCAmelCase_          =							1
                  UpperCAmelCase_          =							pt_model_class(_snake_case).eval()
                  UpperCAmelCase_          =							model_class(_snake_case					,						dtype=jnp.floataa)

                  UpperCAmelCase_          =							convert_pytorch_state_dict_to_flax(pt_model.state_dict()					,						_snake_case)
                  UpperCAmelCase_          =							fx_state

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

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

                  with tempfile.TemporaryDirectory() as tmpdirname:
                       pt_model.save_pretrained(_snake_case)
                       UpperCAmelCase_          =							model_class.from_pretrained(_snake_case					,						from_pt=_snake_case)

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







   @is_pt_flax_cross_test
   def lowerCamelCase						(      self		:						List[Any]):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							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_          =							self._prepare_for_class(_snake_case					,						_snake_case)
                  UpperCAmelCase_          =							{k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

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

                  UpperCAmelCase_          =							pt_model_class(_snake_case).eval()
                  UpperCAmelCase_          =							model_class(_snake_case					,						dtype=jnp.floataa)

                  UpperCAmelCase_          =							load_flax_weights_in_pytorch_model(_snake_case					,						fx_model.params)
                  UpperCAmelCase_     ,       UpperCAmelCase_          =							pt_inputs['''input_ids'''].shape
                  UpperCAmelCase_          =							np.random.randint(0					,						seq_length - 1					,						size=(batch_size,))
                  for batch_idx, start_index in enumerate(_snake_case):
                       UpperCAmelCase_          =							0
                       UpperCAmelCase_          =							1
                       UpperCAmelCase_          =							0
                       UpperCAmelCase_          =							1

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

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

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

                  with tempfile.TemporaryDirectory() as tmpdirname:
                       fx_model.save_pretrained(_snake_case)
                       UpperCAmelCase_          =							pt_model_class.from_pretrained(_snake_case					,						from_flax=_snake_case)

                  with torch.no_grad():
                       UpperCAmelCase_          =							pt_model_loaded(**_snake_case).to_tuple()

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







   @tooslow
   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        for model_class_name in self.all_model_classes:
             UpperCAmelCase_          =							model_class_name.from_pretrained('''EleutherAI/gpt-j-6B''')
             UpperCAmelCase_          =							model(np.ones((1, 1)))
             self.assertIsNotNone(_snake_case)












import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
snake_case_ : List[Any]    		=       data_utils.TransfoXLTokenizer
snake_case_ : int    		=       data_utils.TransfoXLCorpus
snake_case_ : List[Any]    		=       data_utils
snake_case_ : int    		=       data_utils


def    A						(__A   :    Dict   ,				__A   :    List[Any]   ,				__A   :    Union[str, Any]   ,				__A   :    Tuple       )   ->      Union[str, Any]:







     """simple docstring"""


     if transfo_xl_dataset_file:
          # Convert a pre-processed corpus (see original TensorFlow repo)
          with open(__A   ,				'''rb'''       ) as fp:
               UpperCAmelCase_          =							pickle.load(__A   ,				encoding='''latin1'''       )
          # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
          UpperCAmelCase_          =							pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
          print(F"""Save vocabulary to {pytorch_vocab_dump_path}"""       )
          UpperCAmelCase_          =							corpus.vocab.__dict__
          torch.save(__A   ,				__A       )

          UpperCAmelCase_          =							corpus.__dict__
          corpus_dict_no_vocab.pop('''vocab'''   ,				__A       )
          UpperCAmelCase_          =							pytorch_dump_folder_path + '''/''' + CORPUS_NAME
          print(F"""Save dataset to {pytorch_dataset_dump_path}"""       )
          torch.save(__A   ,				__A       )

     if tf_checkpoint_path:
          # Convert a pre-trained TensorFlow model
          UpperCAmelCase_          =							os.path.abspath(__A       )
          UpperCAmelCase_          =							os.path.abspath(__A       )

          print(F"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""       )
          # Initialise PyTorch model
          if transfo_xl_config_file == "":
               UpperCAmelCase_          =							TransfoXLConfig()
          else:
               UpperCAmelCase_          =							TransfoXLConfig.from_json_file(__A       )
          print(F"""Building PyTorch model from configuration: {config}"""       )
          UpperCAmelCase_          =							TransfoXLLMHeadModel(__A       )

          UpperCAmelCase_          =							load_tf_weights_in_transfo_xl(__A   ,				__A   ,				__A       )
          # Save pytorch-model
          UpperCAmelCase_          =							os.path.join(__A   ,				__A       )
          UpperCAmelCase_          =							os.path.join(__A   ,				__A       )
          print(F"""Save PyTorch model to {os.path.abspath(__A       )}"""       )
          torch.save(model.state_dict()   ,				__A       )
          print(F"""Save configuration file to {os.path.abspath(__A       )}"""       )
          with open(__A   ,				'''w'''   ,				encoding='''utf-8'''       ) as f:
               f.write(config.to_json_string()       )


if __name__ == "__main__":
    snake_case_ : List[str]    		=       argparse.ArgumentParser()
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        required=True,
        help="Path to the folder to store the PyTorch model or dataset/vocab.",
    )
    parser.add_argument(
        "--tf_checkpoint_path",
        default="",
        type=str,
        help="An optional path to a TensorFlow checkpoint path to be converted.",
    )
    parser.add_argument(
        "--transfo_xl_config_file",
        default="",
        type=str,
        help=(
            "An optional config json file corresponding to the pre-trained BERT model. \n"
            "This specifies the model architecture."
        ),
    )
    parser.add_argument(
        "--transfo_xl_dataset_file",
        default="",
        type=str,
        help="An optional dataset file to be converted in a vocabulary.",
    )
    snake_case_ : int    		=       parser.parse_args()
    convert_transfo_xl_checkpoint_to_pytorch(
        args.tf_checkpoint_path,
        args.transfo_xl_config_file,
        args.pytorch_dump_folder_path,
        args.transfo_xl_dataset_file,
    )






