




import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
    import os

    from flax.core.frozen_dict import unfreeze
    from flax.traverse_util import flatten_dict

    from transformers import FlaxBertModel

    snake_case_ : int    		=       "0.12"  # assumed parallelism: 8







@require_flax
@is_staging_test
class __snake_case			(			unittest.TestCase		):







   @classmethod
   def lowerCamelCase						(      cls		:						Union[str, Any]):

        """simple docstring"""


        UpperCAmelCase_          =							TOKEN
        HfFolder.save_token(_snake_case)







   @classmethod
   def lowerCamelCase						(      cls		:						Union[str, Any]):

        """simple docstring"""


        try:
             delete_repo(token=cls._token					,						repo_id='''test-model-flax''')
        except HTTPError:
             pass

        try:
             delete_repo(token=cls._token					,						repo_id='''valid_org/test-model-flax-org''')
        except HTTPError:
             pass







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_          =							BertConfig(
            vocab_size=99					,						hidden_size=32					,						num_hidden_layers=5					,						num_attention_heads=4					,						intermediate_size=37)
        UpperCAmelCase_          =							FlaxBertModel(_snake_case)
        model.push_to_hub('''test-model-flax'''					,						use_auth_token=self._token)

        UpperCAmelCase_          =							FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax""")

        UpperCAmelCase_          =							flatten_dict(unfreeze(model.params))
        UpperCAmelCase_          =							flatten_dict(unfreeze(new_model.params))

        for key in base_params.keys():
             UpperCAmelCase_          =							(base_params[key] - new_params[key]).sum().item()
             self.assertLessEqual(_snake_case					,						1e-3					,						msg=F"""{key} not identical""")

        # Reset repo
        delete_repo(token=self._token					,						repo_id='''test-model-flax''')

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
             model.save_pretrained(_snake_case					,						repo_id='''test-model-flax'''					,						push_to_hub=_snake_case					,						use_auth_token=self._token)

        UpperCAmelCase_          =							FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax""")

        UpperCAmelCase_          =							flatten_dict(unfreeze(model.params))
        UpperCAmelCase_          =							flatten_dict(unfreeze(new_model.params))

        for key in base_params.keys():
             UpperCAmelCase_          =							(base_params[key] - new_params[key]).sum().item()
             self.assertLessEqual(_snake_case					,						1e-3					,						msg=F"""{key} not identical""")







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

        """simple docstring"""


        UpperCAmelCase_          =							BertConfig(
            vocab_size=99					,						hidden_size=32					,						num_hidden_layers=5					,						num_attention_heads=4					,						intermediate_size=37)
        UpperCAmelCase_          =							FlaxBertModel(_snake_case)
        model.push_to_hub('''valid_org/test-model-flax-org'''					,						use_auth_token=self._token)

        UpperCAmelCase_          =							FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org''')

        UpperCAmelCase_          =							flatten_dict(unfreeze(model.params))
        UpperCAmelCase_          =							flatten_dict(unfreeze(new_model.params))

        for key in base_params.keys():
             UpperCAmelCase_          =							(base_params[key] - new_params[key]).sum().item()
             self.assertLessEqual(_snake_case					,						1e-3					,						msg=F"""{key} not identical""")

        # Reset repo
        delete_repo(token=self._token					,						repo_id='''valid_org/test-model-flax-org''')

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
             model.save_pretrained(
                 _snake_case					,						repo_id='''valid_org/test-model-flax-org'''					,						push_to_hub=_snake_case					,						use_auth_token=self._token)

        UpperCAmelCase_          =							FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org''')

        UpperCAmelCase_          =							flatten_dict(unfreeze(model.params))
        UpperCAmelCase_          =							flatten_dict(unfreeze(new_model.params))

        for key in base_params.keys():
             UpperCAmelCase_          =							(base_params[key] - new_params[key]).sum().item()
             self.assertLessEqual(_snake_case					,						1e-3					,						msg=F"""{key} not identical""")







def    A						(__A   :    Dict   ,				__A   :    Dict       )   ->      int:







     """simple docstring"""


     UpperCAmelCase_          =							True
     UpperCAmelCase_          =							flatten_dict(modela.params       )
     UpperCAmelCase_          =							flatten_dict(modela.params       )
     for key in flat_params_a.keys():
          if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]       )       ) > 1E-4:
               UpperCAmelCase_          =							False

     return models_are_equal







@require_flax
class __snake_case			(			unittest.TestCase		):







   def lowerCamelCase						(      self		:						Optional[int]):

        """simple docstring"""


        UpperCAmelCase_          =							BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''')
        UpperCAmelCase_          =							FlaxBertModel(_snake_case)

        UpperCAmelCase_          =							'''bert'''
        with tempfile.TemporaryDirectory() as tmp_dir:
             model.save_pretrained(os.path.join(_snake_case					,						_snake_case))

             with self.assertRaises(_snake_case):
                  UpperCAmelCase_          =							FlaxBertModel.from_pretrained(_snake_case)

             UpperCAmelCase_          =							FlaxBertModel.from_pretrained(_snake_case					,						subfolder=_snake_case)

        self.assertTrue(check_models_equal(_snake_case					,						_snake_case))







   def lowerCamelCase						(      self		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''')
        UpperCAmelCase_          =							FlaxBertModel(_snake_case)

        UpperCAmelCase_          =							'''bert'''
        with tempfile.TemporaryDirectory() as tmp_dir:
             model.save_pretrained(os.path.join(_snake_case					,						_snake_case)					,						max_shard_size='''10KB''')

             with self.assertRaises(_snake_case):
                  UpperCAmelCase_          =							FlaxBertModel.from_pretrained(_snake_case)

             UpperCAmelCase_          =							FlaxBertModel.from_pretrained(_snake_case					,						subfolder=_snake_case)

        self.assertTrue(check_models_equal(_snake_case					,						_snake_case))







   def lowerCamelCase						(      self		:						List[Any]):

        """simple docstring"""


        UpperCAmelCase_          =							'''bert'''
        UpperCAmelCase_          =							'''hf-internal-testing/tiny-random-bert-subfolder'''

        with self.assertRaises(_snake_case):
             UpperCAmelCase_          =							FlaxBertModel.from_pretrained(_snake_case)

        UpperCAmelCase_          =							FlaxBertModel.from_pretrained(_snake_case					,						subfolder=_snake_case)

        self.assertIsNotNone(_snake_case)







   def lowerCamelCase						(      self		:						Optional[int]):

        """simple docstring"""


        UpperCAmelCase_          =							'''bert'''
        UpperCAmelCase_          =							'''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
        with self.assertRaises(_snake_case):
             UpperCAmelCase_          =							FlaxBertModel.from_pretrained(_snake_case)

        UpperCAmelCase_          =							FlaxBertModel.from_pretrained(_snake_case					,						subfolder=_snake_case)

        self.assertIsNotNone(_snake_case)












import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()


def    A						(__A   :    Optional[int]   ,				__A   :    int   ,				__A   :    str=None       )   ->      List[Any]:







     """simple docstring"""


     assert torch_layer.weight.shape == weight.shape, F"""{torch_layer} layer.weight does not match"""
     UpperCAmelCase_          =							nn.Parameter(__A       )
     if bias is not None:
          assert torch_layer.bias.shape == bias.shape, F"""{torch_layer} layer.bias does not match"""
          UpperCAmelCase_          =							nn.Parameter(__A       )


def    A						(__A   :    Tuple   ,				__A   :    Dict   ,				__A   :    str       )   ->      Tuple:







     """simple docstring"""


     UpperCAmelCase_          =							np.asarray(weights[0]       )
     UpperCAmelCase_          =							np.asarray(weights[1]       )
     UpperCAmelCase_          =							np.asarray(weights[2]       )

     set_param(
         torch_layer.self_attention.query_key   ,				torch.tensor(__A       ).transpose(1   ,				2       ).contiguous().view(-1   ,				__A       )   ,				)
     set_param(
         torch_layer.self_attention.value   ,				torch.tensor(__A       ).transpose(1   ,				2       ).contiguous().view(-1   ,				__A       )   ,				)
     set_param(
         torch_layer.output.dense   ,				torch.tensor(__A       ).view(-1   ,				__A       ).contiguous().transpose(0   ,				1       )   ,				)


def    A						(__A   :    Optional[Any]   ,				__A   :    Any   ,				__A   :    List[Any]       )   ->      int:







     """simple docstring"""


     UpperCAmelCase_          =							np.asarray(weights[0]       )
     UpperCAmelCase_          =							np.asarray(weights[1]       )
     UpperCAmelCase_          =							np.asarray(weights[2]       )
     UpperCAmelCase_          =							np.asarray(weights[3]       )

     set_param(
         torch_layer.self_attention.query   ,				torch.tensor(__A       ).transpose(1   ,				2       ).contiguous().view(-1   ,				__A       )   ,				)
     set_param(
         torch_layer.self_attention.key   ,				torch.tensor(__A       ).transpose(1   ,				2       ).contiguous().view(-1   ,				__A       )   ,				)
     set_param(
         torch_layer.self_attention.value   ,				torch.tensor(__A       ).transpose(1   ,				2       ).contiguous().view(-1   ,				__A       )   ,				)
     set_param(
         torch_layer.output.dense   ,				torch.tensor(__A       ).view(-1   ,				__A       ).contiguous().transpose(0   ,				1       )   ,				)


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







     """simple docstring"""


     UpperCAmelCase_          =							weights[0][0][0]
     UpperCAmelCase_          =							np.asarray(layer_norm_a[0]       )
     UpperCAmelCase_          =							np.asarray(layer_norm_a[1]       )
     set_param(
         torch_block.attention.layer_norm   ,				torch.tensor(__A       )   ,				torch.tensor(__A       )   ,				)

     # lsh weights + output
     UpperCAmelCase_          =							weights[0][1]
     if len(__A       ) < 4:
          set_layer_weights_in_torch_lsh(__A   ,				torch_block.attention   ,				__A       )
     else:
          set_layer_weights_in_torch_local(__A   ,				torch_block.attention   ,				__A       )

     # intermediate weighs
     UpperCAmelCase_          =							weights[2][0][1][2]

     # Chunked Feed Forward
     if len(__A       ) == 4:
          UpperCAmelCase_          =							intermediate_weights[2]

     # layernorm 2
     UpperCAmelCase_          =							np.asarray(intermediate_weights[0][0]       )
     UpperCAmelCase_          =							np.asarray(intermediate_weights[0][1]       )
     set_param(
         torch_block.feed_forward.layer_norm   ,				torch.tensor(__A       )   ,				torch.tensor(__A       )   ,				)

     # intermediate dense
     UpperCAmelCase_          =							np.asarray(intermediate_weights[1][0]       )
     UpperCAmelCase_          =							np.asarray(intermediate_weights[1][1]       )
     set_param(
         torch_block.feed_forward.dense.dense   ,				torch.tensor(__A       ).transpose(0   ,				1       ).contiguous()   ,				torch.tensor(__A       )   ,				)

     # intermediate out
     UpperCAmelCase_          =							np.asarray(intermediate_weights[4][0]       )
     UpperCAmelCase_          =							np.asarray(intermediate_weights[4][1]       )
     set_param(
         torch_block.feed_forward.output.dense   ,				torch.tensor(__A       ).transpose(0   ,				1       ).contiguous()   ,				torch.tensor(__A       )   ,				)


def    A						(__A   :    Optional[int]   ,				__A   :    Tuple   ,				__A   :    Any       )   ->      Tuple:







     """simple docstring"""


     UpperCAmelCase_          =							torch_model.reformer

     # word embeds
     UpperCAmelCase_          =							np.asarray(weights[1]       )
     set_param(
         torch_model_reformer.embeddings.word_embeddings   ,				torch.tensor(__A       )   ,				)

     if isinstance(weights[3]   ,				__A       ):
          UpperCAmelCase_          =							torch_model_reformer.embeddings.position_embeddings
          for emb_idx in range(len(position_embeddings.weights       )       ):
               UpperCAmelCase_          =							np.asarray(weights[3][emb_idx][0]       )
               assert (
                   position_embeddings.weights[emb_idx].shape == emb_weights.shape
               ), F"""{position_embeddings[emb_idx]} emb does not match"""
               UpperCAmelCase_          =							nn.Parameter(torch.tensor(__A       )       )

     UpperCAmelCase_          =							weights[5]
     assert len(torch_model_reformer.encoder.layers       ) * 4 == len(
         __A       ), "HF and trax model do not have the same number of layers"
     for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers       ):
          UpperCAmelCase_          =							trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
          set_block_weights_in_torch(__A   ,				__A   ,				__A       )

     # output layer norm
     UpperCAmelCase_          =							np.asarray(weights[7][0]       )
     UpperCAmelCase_          =							np.asarray(weights[7][1]       )
     set_param(
         torch_model_reformer.encoder.layer_norm   ,				torch.tensor(__A       )   ,				torch.tensor(__A       )   ,				)

     # output embeddings
     UpperCAmelCase_          =							np.asarray(weights[9][0]       )
     UpperCAmelCase_          =							np.asarray(weights[9][1]       )
     set_param(
         torch_model.lm_head.decoder   ,				torch.tensor(__A       ).transpose(0   ,				1       ).contiguous()   ,				torch.tensor(__A       )   ,				)


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







     """simple docstring"""


     UpperCAmelCase_          =							ReformerConfig.from_json_file(__A       )
     print(F"""Building PyTorch model from configuration: {config}"""       )
     UpperCAmelCase_          =							ReformerModelWithLMHead(__A       )

     with open(__A   ,				'''rb'''       ) as f:
          UpperCAmelCase_          =							pickle.load(__A       )['''weights''']

     set_model_weights_in_torch(__A   ,				__A   ,				config.hidden_size       )

     # Save pytorch-model
     print(F"""Save PyTorch model to {pytorch_dump_path}"""       )
     torch.save(model.state_dict()   ,				__A       )


if __name__ == "__main__":
    snake_case_ : Optional[Any]    		=       argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--trax_model_pkl_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
    )
    parser.add_argument(
        "--config_file",
        default=None,
        type=str,
        required=True,
        help=(
            "The config json file corresponding to the pre-trained Reformer model. \n"
            "This specifies the model architecture."
        ),
    )
    parser.add_argument(
        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    snake_case_ : List[Any]    		=       parser.parse_args()
    convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)






