

'''simple docstring'''







import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


__lowerCamelCase									=     logging.get_logger(__name__)
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__=False					)      ->      Tuple:
       try:
              import torch  # noqa: F401
       except ImportError:
              logger.error(
                  """Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see"""
                  """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
                  """ instructions."""					)
              raise

       if not is_sharded:
              A_						       =   os.path.abspath(UpperCAmelCase__					)
              logger.info(F'''Loading PyTorch weights from {pt_path}'''					)

              A_						       =   torch.load(UpperCAmelCase__,							map_location="""cpu"""					)
              logger.info(F'''PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()					):,} parameters.'''					)

              A_						       =   convert_pytorch_state_dict_to_flax(UpperCAmelCase__,							UpperCAmelCase__					)
       else:
              # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
              A_						       =   convert_pytorch_sharded_state_dict_to_flax(UpperCAmelCase__,							UpperCAmelCase__					)
       return flax_state_dict
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							)      ->      (Tuple[str], np.ndarray):

       def is_key_or_prefix_key_in_dict(UpperCAmelCase__					) -> bool:
              return len(set(UpperCAmelCase__					) & {key, (model_prefix,) + key}					) > 0

       # layer norm
       A_						       =   pt_tuple_key[:-1] + ("""scale""",)
       if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(UpperCAmelCase__					):
              return renamed_pt_tuple_key, pt_tensor

       # batch norm layer mean
       A_						       =   pt_tuple_key[:-1] + ("""mean""",)
       if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(UpperCAmelCase__					):
              return renamed_pt_tuple_key, pt_tensor

       # batch norm layer var
       A_						       =   pt_tuple_key[:-1] + ("""var""",)
       if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(UpperCAmelCase__					):
              return renamed_pt_tuple_key, pt_tensor

       # embedding
       A_						       =   pt_tuple_key[:-1] + ("""embedding""",)
       if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(UpperCAmelCase__					):
              return renamed_pt_tuple_key, pt_tensor

       # conv layer
       A_						       =   pt_tuple_key[:-1] + ("""kernel""",)
       if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(UpperCAmelCase__					):
              A_						       =   pt_tensor.transpose(2,							3,							1,							0					)
              return renamed_pt_tuple_key, pt_tensor

       # linear layer
       A_						       =   pt_tuple_key[:-1] + ("""kernel""",)
       if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(UpperCAmelCase__					):
              A_						       =   pt_tensor.T
              return renamed_pt_tuple_key, pt_tensor

       # old PyTorch layer norm weight
       A_						       =   pt_tuple_key[:-1] + ("""weight""",)
       if pt_tuple_key[-1] == "gamma":
              return renamed_pt_tuple_key, pt_tensor

       # old PyTorch layer norm bias
       A_						       =   pt_tuple_key[:-1] + ("""bias""",)
       if pt_tuple_key[-1] == "beta":
              return renamed_pt_tuple_key, pt_tensor

       # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
       A_						       =   None
       if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
              A_						       =   pt_tuple_key[-2] + """_g"""
       elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
              A_						       =   pt_tuple_key[-2] + """_v"""
       if name is not None:
              A_						       =   pt_tuple_key[:-3] + (name,)
              return renamed_pt_tuple_key, pt_tensor

       return pt_tuple_key, pt_tensor
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__					)      ->      Optional[int]:
       # convert pytorch tensor to numpy
       A_						       =   {k: v.numpy() for k, v in pt_state_dict.items()}

       A_						       =   flax_model.base_model_prefix

       # use params dict if the model contains batch norm layers
       if "params" in flax_model.params:
              A_						       =   flax_model.params["""params"""]
       else:
              A_						       =   flax_model.params
       A_						       =   flatten_dict(UpperCAmelCase__					)

       # add batch_stats keys,values to dict
       if "batch_stats" in flax_model.params:
              A_						       =   flatten_dict(flax_model.params["""batch_stats"""]					)
              random_flax_state_dict.update(UpperCAmelCase__					)

       A_						       =   {}

       A_						       =   (model_prefix not in flax_model_params) and (
           model_prefix in {k.split("""."""					)[0] for k in pt_state_dict.keys()}
       )
       A_						       =   (model_prefix in flax_model_params) and (
           model_prefix not in {k.split("""."""					)[0] for k in pt_state_dict.keys()}
       )

       # Need to change some parameters name to match Flax names
       for pt_key, pt_tensor in pt_state_dict.items():
              A_						       =   tuple(pt_key.split("""."""					)					)

              # remove base model prefix if necessary
              A_						       =   pt_tuple_key[0] == model_prefix
              if load_model_with_head_into_base_model and has_base_model_prefix:
                     A_						       =   pt_tuple_key[1:]

              # Correctly rename weight parameters
              A_       ,      A_						       =   rename_key_and_reshape_tensor(
                  UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)

              # add model prefix if necessary
              A_						       =   (model_prefix,) + flax_key in random_flax_state_dict
              if load_base_model_into_model_with_head and require_base_model_prefix:
                     A_						       =   (model_prefix,) + flax_key

              if flax_key in random_flax_state_dict:
                     if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                            raise ValueError(
                                F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
                                F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''					)

        # add batch stats if the model contains batchnorm layers
              if "batch_stats" in flax_model.params:
                     if "mean" in flax_key[-1] or "var" in flax_key[-1]:
                            A_						       =   jnp.asarray(UpperCAmelCase__					)
                            continue
                     # remove num_batches_tracked key
                     if "num_batches_tracked" in flax_key[-1]:
                            flax_state_dict.pop(UpperCAmelCase__,							UpperCAmelCase__					)
                            continue

                     # also add unexpected weight so that warning is thrown
                     A_						       =   jnp.asarray(UpperCAmelCase__					)

              else:
                     # also add unexpected weight so that warning is thrown
                     A_						       =   jnp.asarray(UpperCAmelCase__					)

       return unflatten_dict(UpperCAmelCase__					)
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__					)      ->      Dict:
       import torch

       # Load the index
       A_						       =   {}
       for shard_file in shard_filenames:
              # load using msgpack utils
              A_						       =   torch.load(UpperCAmelCase__					)
              A_						       =   {k: v.numpy() for k, v in pt_state_dict.items()}

              A_						       =   flax_model.base_model_prefix

              # use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
              if "batch_stats" in flax_model.params:
                     A_						       =   flax_model.params["""params"""]

                     A_						       =   flatten_dict(UpperCAmelCase__					)
                     random_flax_state_dict.update(flatten_dict(flax_model.params["""batch_stats"""]					)					)
              else:
                     A_						       =   flax_model.params
                     A_						       =   flatten_dict(UpperCAmelCase__					)

              A_						       =   (model_prefix not in flax_model_params) and (
                  model_prefix in {k.split("""."""					)[0] for k in pt_state_dict.keys()}
              )
              A_						       =   (model_prefix in flax_model_params) and (
                  model_prefix not in {k.split("""."""					)[0] for k in pt_state_dict.keys()}
              )
              # Need to change some parameters name to match Flax names
              for pt_key, pt_tensor in pt_state_dict.items():
                     A_						       =   tuple(pt_key.split("""."""					)					)

                     # remove base model prefix if necessary
                     A_						       =   pt_tuple_key[0] == model_prefix
                     if load_model_with_head_into_base_model and has_base_model_prefix:
                            A_						       =   pt_tuple_key[1:]

                     # Correctly rename weight parameters
                     A_       ,      A_						       =   rename_key_and_reshape_tensor(
                         UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)
                     # add model prefix if necessary
                     A_						       =   (model_prefix,) + flax_key in random_flax_state_dict
                     if load_base_model_into_model_with_head and require_base_model_prefix:
                            A_						       =   (model_prefix,) + flax_key

                     if flax_key in random_flax_state_dict:
                            if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                                   raise ValueError(
                                       F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
                                       F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''					)

            # add batch stats if the model contains batchnorm layers
                     if "batch_stats" in flax_model.params:
                            if "mean" in flax_key[-1]:
                                   A_						       =   jnp.asarray(UpperCAmelCase__					)
                                   continue
                            if "var" in flax_key[-1]:
                                   A_						       =   jnp.asarray(UpperCAmelCase__					)
                                   continue
                            # remove num_batches_tracked key
                            if "num_batches_tracked" in flax_key[-1]:
                                   flax_state_dict.pop(UpperCAmelCase__,							UpperCAmelCase__					)
                                   continue

                            # also add unexpected weight so that warning is thrown
                            A_						       =   jnp.asarray(UpperCAmelCase__					)

                     else:
                            # also add unexpected weight so that warning is thrown
                            A_						       =   jnp.asarray(UpperCAmelCase__					)
       return unflatten_dict(UpperCAmelCase__					)
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__					)      ->      List[str]:
       A_						       =   os.path.abspath(UpperCAmelCase__					)
       logger.info(F'''Loading Flax weights from {flax_checkpoint_path}'''					)

       # import correct flax class
       A_						       =   getattr(UpperCAmelCase__,							"""Flax""" + model.__class__.__name__					)

       # load flax weight dict
       with open(UpperCAmelCase__,							"""rb"""					) as state_f:
              try:
                     A_						       =   from_bytes(UpperCAmelCase__,							state_f.read()					)
              except UnpicklingError:
                     raise EnvironmentError(F'''Unable to convert {flax_checkpoint_path} to Flax deserializable object. '''					)

       return load_flax_weights_in_pytorch_model(UpperCAmelCase__,							UpperCAmelCase__					)







def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__					)      ->      List[Any]:

       try:
              import torch  # noqa: F401
       except ImportError:
              logger.error(
                  """Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"""
                  """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
                  """ instructions."""					)
              raise

       # check if we have bf16 weights
       A_						       =   flatten_dict(jax.tree_util.tree_map(lambda UpperCAmelCase__					: x.dtype == jnp.bfloataa,							UpperCAmelCase__					)					).values()
       if any(UpperCAmelCase__					):
              # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
              # and bf16 is not fully supported in PT yet.
              logger.warning(
                  """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
                  """before loading those in PyTorch model."""					)
              A_						       =   jax.tree_util.tree_map(
                  lambda UpperCAmelCase__					: params.astype(np.floataa					) if params.dtype == jnp.bfloataa else params,							UpperCAmelCase__					)

       A_						       =   flatten_dict(UpperCAmelCase__					)
       A_						       =   pt_model.state_dict()

       A_						       =   (pt_model.base_model_prefix in flax_state) and (
           pt_model.base_model_prefix not in {k.split("""."""					)[0] for k in pt_model_dict.keys()}
       )
       A_						       =   (pt_model.base_model_prefix not in flax_state) and (
           pt_model.base_model_prefix in {k.split("""."""					)[0] for k in pt_model_dict.keys()}
       )

       # keep track of unexpected & missing keys
       A_						       =   []
       A_						       =   set(pt_model_dict.keys()					)

       for flax_key_tuple, flax_tensor in flax_state_dict.items():
              A_						       =   flax_key_tuple[0] == pt_model.base_model_prefix
              A_						       =   """.""".join((pt_model.base_model_prefix,) + flax_key_tuple					) in pt_model_dict

              # adapt flax_key to prepare for loading from/to base model only
              if load_model_with_head_into_base_model and has_base_model_prefix:
                     A_						       =   flax_key_tuple[1:]
              elif load_base_model_into_model_with_head and require_base_model_prefix:
                     A_						       =   (pt_model.base_model_prefix,) + flax_key_tuple

              # rename flax weights to PyTorch format
              if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(UpperCAmelCase__					) not in pt_model_dict:
                     # conv layer
                     A_						       =   flax_key_tuple[:-1] + ("""weight""",)
                     A_						       =   jnp.transpose(UpperCAmelCase__,							(3, 2, 0, 1)					)
              elif flax_key_tuple[-1] == "kernel" and ".".join(UpperCAmelCase__					) not in pt_model_dict:
                     # linear layer
                     A_						       =   flax_key_tuple[:-1] + ("""weight""",)
                     A_						       =   flax_tensor.T
              elif flax_key_tuple[-1] in ["scale", "embedding"]:
                     A_						       =   flax_key_tuple[:-1] + ("""weight""",)

              # adding batch stats from flax batch norm to pt
              elif "mean" in flax_key_tuple[-1]:
                     A_						       =   flax_key_tuple[:-1] + ("""running_mean""",)
              elif "var" in flax_key_tuple[-1]:
                     A_						       =   flax_key_tuple[:-1] + ("""running_var""",)

              if "batch_stats" in flax_state:
                     A_						       =   """.""".join(flax_key_tuple[1:]					)  # Remove the params/batch_stats header
              else:
                     A_						       =   """.""".join(UpperCAmelCase__					)

              # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
              A_						       =   {}
              # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
              for key in pt_model_dict:
                     A_						       =   key.split("""."""					)
                     A_						       =   None
                     if key_components[-3::2] == ["parametrizations", "original0"]:
                            A_						       =   key_components[-2] + """_g"""
                     elif key_components[-3::2] == ["parametrizations", "original1"]:
                            A_						       =   key_components[-2] + """_v"""
                     if name is not None:
                            A_						       =   key_components[:-3] + [name]
                            A_						       =   """.""".join(UpperCAmelCase__					)
                            A_						       =   key

              if flax_key in special_pt_names:
                     A_						       =   special_pt_names[flax_key]

              if flax_key in pt_model_dict:
                     if flax_tensor.shape != pt_model_dict[flax_key].shape:
                            raise ValueError(
                                F'''Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '''
                                F'''to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'''					)
                     else:
                            # add weight to pytorch dict
                            A_						       =   np.asarray(UpperCAmelCase__					) if not isinstance(UpperCAmelCase__,							np.ndarray					) else flax_tensor
                            A_						       =   torch.from_numpy(UpperCAmelCase__					)
                            # remove from missing keys
                            missing_keys.remove(UpperCAmelCase__					)
              else:
                     # weight is not expected by PyTorch model
                     unexpected_keys.append(UpperCAmelCase__					)

       pt_model.load_state_dict(UpperCAmelCase__					)

       # re-transform missing_keys to list
       A_						       =   list(UpperCAmelCase__					)

       if len(UpperCAmelCase__					) > 0:
              logger.warning(
                  """Some weights of the Flax model were not used when initializing the PyTorch model"""
                  F''' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'''
                  F''' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'''
                  """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
                  F''' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'''
                  """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
                  """ FlaxBertForSequenceClassification model)."""					)
       else:
              logger.warning(F'''All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'''					)
       if len(UpperCAmelCase__					) > 0:
              logger.warning(
                  F'''Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'''
                  F''' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'''
                  """ use it for predictions and inference."""					)
       else:
              logger.warning(
                  F'''All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'''
                  """If your task is similar to the task the model of the checkpoint was trained on, """
                  F'''you can already use {pt_model.__class__.__name__} for predictions without further training.'''					)

       return pt_model



'''simple docstring'''







import os
import sys
import unittest


__lowerCamelCase									=     os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, '''utils'''))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
__lowerCamelCase									=     os.path.join(git_repo_path, '''src''', '''transformers''')

__lowerCamelCase									=     '''
{0} = None
'''

__lowerCamelCase									=     '''
class {0}(metaclass=DummyObject):
    _backends = {1}

    def __init__(self, *args, **kwargs):
        requires_backends(self, {1})
'''


__lowerCamelCase									=     '''
def {0}(*args, **kwargs):
    requires_backends({0}, {1})
'''


class        A__   (   unittest.TestCase      ):



    def        snake_case_      (       self							)    ->   Dict:





           '''simple docstring'''
           A_						       =   find_backend("""    _import_structure[\"models.albert\"].append(\"AlbertTokenizerFast\")"""							)
           self.assertIsNone(UpperCamelCase__							)

           A_						       =   find_backend("""    if not is_tokenizers_available():"""							)
           self.assertEqual(UpperCamelCase__		,					"""tokenizers"""							)

           A_						       =   find_backend("""    if not is_tensorflow_text_available():"""							)
           self.assertEqual(UpperCamelCase__		,					"""tensorflow_text"""							)

           A_						       =   find_backend("""    if not (is_sentencepiece_available() and is_tokenizers_available()):"""							)
           self.assertEqual(UpperCamelCase__		,					"""sentencepiece_and_tokenizers"""							)

           A_						       =   find_backend(
               """    if not (is_sentencepiece_available() and is_tensorflow_text_available()):"""							)
           self.assertEqual(UpperCamelCase__		,					"""sentencepiece_and_tensorflow_text"""							)

           A_						       =   find_backend(
               """    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):"""							)
           self.assertEqual(UpperCamelCase__		,					"""sentencepiece_and_tokenizers_and_vision"""							)



    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           A_						       =   read_init()
           # We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
           self.assertIn("""torch"""		,					UpperCamelCase__							)
           self.assertIn("""tensorflow_text"""		,					UpperCamelCase__							)
           self.assertIn("""sentencepiece_and_tokenizers"""		,					UpperCamelCase__							)

           # Likewise, we can't assert on the exact content of a key
           self.assertIn("""BertModel"""		,					objects["""torch"""]							)
           self.assertIn("""TFBertModel"""		,					objects["""tf"""]							)
           self.assertIn("""FlaxBertModel"""		,					objects["""flax"""]							)
           self.assertIn("""BertModel"""		,					objects["""torch"""]							)
           self.assertIn("""TFBertTokenizer"""		,					objects["""tensorflow_text"""]							)
           self.assertIn("""convert_slow_tokenizer"""		,					objects["""sentencepiece_and_tokenizers"""]							)



    def        snake_case_      (       self							)    ->   str:





           '''simple docstring'''
           A_						       =   create_dummy_object("""CONSTANT"""		,					"""'torch'"""							)
           self.assertEqual(UpperCamelCase__		,					"""\nCONSTANT = None\n"""							)

           A_						       =   create_dummy_object("""function"""		,					"""'torch'"""							)
           self.assertEqual(
               UpperCamelCase__		,					"""\ndef function(*args, **kwargs):\n    requires_backends(function, 'torch')\n"""							)

           A_						       =   """
class FakeClass(metaclass=DummyObject):
    _backends = 'torch'

    def __init__(self, *args, **kwargs):
        requires_backends(self, 'torch')
"""
           A_						       =   create_dummy_object("""FakeClass"""		,					"""'torch'"""							)
           self.assertEqual(UpperCamelCase__		,					UpperCamelCase__							)







    def        snake_case_      (       self							)    ->   str:





           '''simple docstring'''
           A_						       =   """# This file is autogenerated by the command `make fix-copies`, do not edit.
from ..utils import DummyObject, requires_backends


CONSTANT = None


def function(*args, **kwargs):
    requires_backends(function, [\"torch\"])


class FakeClass(metaclass=DummyObject):
    _backends = [\"torch\"]

    def __init__(self, *args, **kwargs):
        requires_backends(self, [\"torch\"])
"""
           A_						       =   create_dummy_files({"""torch""": ["""CONSTANT""", """function""", """FakeClass"""]}							)
           self.assertEqual(dummy_files["""torch"""]		,					UpperCamelCase__							)
