


'''simple docstring'''
import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase__    					=      logging.get_logger(__name__)





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Dict:
    UpperCAmelCase__			:		Optional[Any]													=				os.path.abspath(lowerCAmelCase__						)
    logger.info(F"""Converting TensorFlow checkpoint from {tf_path}"""						)
    # Load weights from TF model
    UpperCAmelCase__			:		int													=				tf.train.list_variables(lowerCAmelCase__						)
    UpperCAmelCase__			:		List[str]													=				[]
    UpperCAmelCase__			:		Optional[Any]													=				[]
    UpperCAmelCase__			:		List[str]													=				[]
    for full_name, shape in init_vars:
        # logger.info(f"Loading TF weight {name} with shape {shape}")
        UpperCAmelCase__			:		str													=				full_name.split('''/'''						)
        if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
            logger.info(F"""Skipping non-model layer {full_name}"""						)
            continue
        if "optimizer" in full_name:
            logger.info(F"""Skipping optimization layer {full_name}"""						)
            continue
        if name[0] == "model":
            # ignore initial 'model'
            UpperCAmelCase__			:		Tuple													=				name[1:]
        # figure out how many levels deep the name is
        UpperCAmelCase__			:		List[Any]													=				0
        for _name in name:
            if _name.startswith('''layer_with_weights'''						):
                depth += 1
            else:
                break
        layer_depth.append(lowerCAmelCase__						)
        # read data
        UpperCAmelCase__			:		List[str]													=				tf.train.load_variable(lowerCAmelCase__  ,					lowerCAmelCase__						)
        names.append('''/'''.join(lowerCAmelCase__						)						)
        arrays.append(lowerCAmelCase__						)
    logger.info(F"""Read a total of {len(lowerCAmelCase__						):,} layers"""						)

    # Sanity check
    if len(set(lowerCAmelCase__						)						) != 1:
        raise ValueError(F"""Found layer names with different depths (layer depth {list(set(lowerCAmelCase__						)						)})"""						)
    UpperCAmelCase__			:		List[str]													=				list(set(lowerCAmelCase__						)						)[0]
    if layer_depth != 1:
        raise ValueError(
            '''The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP'''
            ''' heads.'''						)

    # convert layers
    logger.info('''Converting weights...'''						)
    for full_name, array in zip(lowerCAmelCase__  ,					lowerCAmelCase__						):
        UpperCAmelCase__			:		Optional[int]													=				full_name.split('''/'''						)
        UpperCAmelCase__			:		List[str]													=				model
        UpperCAmelCase__			:		int													=				[]
        for i, m_name in enumerate(lowerCAmelCase__						):
            if m_name == ".ATTRIBUTES":
                # variable names end with .ATTRIBUTES/VARIABLE_VALUE
                break
            if m_name.startswith('''layer_with_weights'''						):
                UpperCAmelCase__			:		int													=				int(m_name.split('''-'''						)[-1]						)
                if layer_num <= 2:
                    # embedding layers
                    # layer_num 0: word_embeddings
                    # layer_num 1: position_embeddings
                    # layer_num 2: token_type_embeddings
                    continue
                elif layer_num == 3:
                    # embedding LayerNorm
                    trace.extend(['''embeddings''', '''LayerNorm''']						)
                    UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					'''embeddings'''						)
                    UpperCAmelCase__			:		Tuple													=				getattr(lowerCAmelCase__  ,					'''LayerNorm'''						)
                elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
                    # encoder layers
                    trace.extend(['''encoder''', '''layer''', str(layer_num - 4						)]						)
                    UpperCAmelCase__			:		Tuple													=				getattr(lowerCAmelCase__  ,					'''encoder'''						)
                    UpperCAmelCase__			:		Tuple													=				getattr(lowerCAmelCase__  ,					'''layer'''						)
                    UpperCAmelCase__			:		Union[str, Any]													=				pointer[layer_num - 4]
                elif layer_num == config.num_hidden_layers + 4:
                    # pooler layer
                    trace.extend(['''pooler''', '''dense''']						)
                    UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					'''pooler'''						)
                    UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					'''dense'''						)
            elif m_name == "embeddings":
                trace.append('''embeddings'''						)
                UpperCAmelCase__			:		List[str]													=				getattr(lowerCAmelCase__  ,					'''embeddings'''						)
                if layer_num == 0:
                    trace.append('''word_embeddings'''						)
                    UpperCAmelCase__			:		List[str]													=				getattr(lowerCAmelCase__  ,					'''word_embeddings'''						)
                elif layer_num == 1:
                    trace.append('''position_embeddings'''						)
                    UpperCAmelCase__			:		List[Any]													=				getattr(lowerCAmelCase__  ,					'''position_embeddings'''						)
                elif layer_num == 2:
                    trace.append('''token_type_embeddings'''						)
                    UpperCAmelCase__			:		Any													=				getattr(lowerCAmelCase__  ,					'''token_type_embeddings'''						)
                else:
                    raise ValueError(F"""Unknown embedding layer with name {full_name}"""						)
                trace.append('''weight'''						)
                UpperCAmelCase__			:		Dict													=				getattr(lowerCAmelCase__  ,					'''weight'''						)
            elif m_name == "_attention_layer":
                # self-attention layer
                trace.extend(['''attention''', '''self''']						)
                UpperCAmelCase__			:		Union[str, Any]													=				getattr(lowerCAmelCase__  ,					'''attention'''						)
                UpperCAmelCase__			:		Optional[int]													=				getattr(lowerCAmelCase__  ,					'''self'''						)
            elif m_name == "_attention_layer_norm":
                # output attention norm
                trace.extend(['''attention''', '''output''', '''LayerNorm''']						)
                UpperCAmelCase__			:		str													=				getattr(lowerCAmelCase__  ,					'''attention'''						)
                UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					'''output'''						)
                UpperCAmelCase__			:		List[Any]													=				getattr(lowerCAmelCase__  ,					'''LayerNorm'''						)
            elif m_name == "_attention_output_dense":
                # output attention dense
                trace.extend(['''attention''', '''output''', '''dense''']						)
                UpperCAmelCase__			:		Tuple													=				getattr(lowerCAmelCase__  ,					'''attention'''						)
                UpperCAmelCase__			:		Any													=				getattr(lowerCAmelCase__  ,					'''output'''						)
                UpperCAmelCase__			:		Optional[int]													=				getattr(lowerCAmelCase__  ,					'''dense'''						)
            elif m_name == "_output_dense":
                # output dense
                trace.extend(['''output''', '''dense''']						)
                UpperCAmelCase__			:		List[Any]													=				getattr(lowerCAmelCase__  ,					'''output'''						)
                UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					'''dense'''						)
            elif m_name == "_output_layer_norm":
                # output dense
                trace.extend(['''output''', '''LayerNorm''']						)
                UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					'''output'''						)
                UpperCAmelCase__			:		str													=				getattr(lowerCAmelCase__  ,					'''LayerNorm'''						)
            elif m_name == "_key_dense":
                # attention key
                trace.append('''key'''						)
                UpperCAmelCase__			:		Tuple													=				getattr(lowerCAmelCase__  ,					'''key'''						)
            elif m_name == "_query_dense":
                # attention query
                trace.append('''query'''						)
                UpperCAmelCase__			:		Union[str, Any]													=				getattr(lowerCAmelCase__  ,					'''query'''						)
            elif m_name == "_value_dense":
                # attention value
                trace.append('''value'''						)
                UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					'''value'''						)
            elif m_name == "_intermediate_dense":
                # attention intermediate dense
                trace.extend(['''intermediate''', '''dense''']						)
                UpperCAmelCase__			:		int													=				getattr(lowerCAmelCase__  ,					'''intermediate'''						)
                UpperCAmelCase__			:		Tuple													=				getattr(lowerCAmelCase__  ,					'''dense'''						)
            elif m_name == "_output_layer_norm":
                # output layer norm
                trace.append('''output'''						)
                UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					'''output'''						)
            # weights & biases
            elif m_name in ["bias", "beta"]:
                trace.append('''bias'''						)
                UpperCAmelCase__			:		str													=				getattr(lowerCAmelCase__  ,					'''bias'''						)
            elif m_name in ["kernel", "gamma"]:
                trace.append('''weight'''						)
                UpperCAmelCase__			:		Any													=				getattr(lowerCAmelCase__  ,					'''weight'''						)
            else:
                logger.warning(F"""Ignored {m_name}"""						)
        # for certain layers reshape is necessary
        UpperCAmelCase__			:		Optional[int]													=				'''.'''.join(lowerCAmelCase__						)
        if re.match(R'''(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)'''  ,					lowerCAmelCase__						) or re.match(
            R'''(\S+)\.attention\.output\.dense\.weight'''  ,					lowerCAmelCase__						):
            UpperCAmelCase__			:		Union[str, Any]													=				array.reshape(pointer.data.shape						)
        if "kernel" in full_name:
            UpperCAmelCase__			:		Dict													=				array.transpose()
        if pointer.shape == array.shape:
            UpperCAmelCase__			:		List[Any]													=				torch.from_numpy(lowerCAmelCase__						)
        else:
            raise ValueError(
                F"""Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:"""
                F""" {array.shape}"""						)
        logger.info(F"""Successfully set variable {full_name} to PyTorch layer {trace}"""						)
    return model





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Optional[int]:
    # Instantiate model
    logger.info(F"""Loading model based on config from {config_path}..."""						)
    UpperCAmelCase__			:		Union[str, Any]													=				BertConfig.from_json_file(lowerCAmelCase__						)
    UpperCAmelCase__			:		List[str]													=				BertModel(lowerCAmelCase__						)

    # Load weights from checkpoint
    logger.info(F"""Loading weights from checkpoint {tf_checkpoint_path}..."""						)
    load_tfa_weights_in_bert(lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)

    # Save pytorch-model
    logger.info(F"""Saving PyTorch model to {pytorch_dump_path}..."""						)
    torch.save(model.state_dict()  ,					lowerCAmelCase__						)


if __name__ == "__main__":
  UpperCamelCase__    					=      argparse.ArgumentParser()
  parser.add_argument(
      '''--tf_checkpoint_path''', type=str, required=True, help='''Path to the TensorFlow 2.x checkpoint path.'''
  )
  parser.add_argument(
      '''--bert_config_file''',
      type=str,
      required=True,
      help='''The config json file corresponding to the BERT model. This specifies the model architecture.''',
  )
  parser.add_argument(
      '''--pytorch_dump_path''',
      type=str,
      required=True,
      help='''Path to the output PyTorch model (must include filename).''',
  )
  UpperCamelCase__    					=      parser.parse_args()
  convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)




'''simple docstring'''
import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCamelCase__    					=      logging.get_logger(__name__)





def 					a__       (					lowerCAmelCase__						)    ->	Optional[Any]:
    UpperCAmelCase__			:		str													=				R'''\w+[.]\d+'''
    UpperCAmelCase__			:		List[Any]													=				re.findall(lowerCAmelCase__  ,					lowerCAmelCase__						)
    for pat in pats:
        UpperCAmelCase__			:		Union[str, Any]													=				key.replace(lowerCAmelCase__  ,					'''_'''.join(pat.split('''.'''						)						)						)
    return key





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Union[str, Any]:
    UpperCAmelCase__			:		Optional[Any]													=				pt_tuple_key[:-1] + ('''scale''',)
    if (
        any('''norm''' in str_ for str_ in pt_tuple_key						)
        and (pt_tuple_key[-1] == "bias")
        and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
        and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
    ):
        UpperCAmelCase__			:		Union[str, Any]													=				pt_tuple_key[:-1] + ('''scale''',)
        return renamed_pt_tuple_key, pt_tensor
    elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
        UpperCAmelCase__			:		Optional[int]													=				pt_tuple_key[:-1] + ('''scale''',)
        return renamed_pt_tuple_key, pt_tensor

    # embedding
    if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
        UpperCAmelCase__			:		str													=				pt_tuple_key[:-1] + ('''embedding''',)
        return renamed_pt_tuple_key, pt_tensor

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

    # linear layer
    UpperCAmelCase__			:		int													=				pt_tuple_key[:-1] + ('''kernel''',)
    if pt_tuple_key[-1] == "weight":
        UpperCAmelCase__			:		Optional[Any]													=				pt_tensor.T
        return renamed_pt_tuple_key, pt_tensor

    # old PyTorch layer norm weight
    UpperCAmelCase__			:		Union[str, Any]													=				pt_tuple_key[:-1] + ('''weight''',)
    if pt_tuple_key[-1] == "gamma":
        return renamed_pt_tuple_key, pt_tensor

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

    return pt_tuple_key, pt_tensor







def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__=42						)    ->	Tuple:
    # Step 1: Convert pytorch tensor to numpy
    UpperCAmelCase__			:		int													=				{k: v.numpy() for k, v in pt_state_dict.items()}

    # Step 2: Since the model is stateless, get random Flax params
    UpperCAmelCase__			:		Tuple													=				flax_model.init_weights(PRNGKey(lowerCAmelCase__						)						)

    UpperCAmelCase__			:		Optional[Any]													=				flatten_dict(lowerCAmelCase__						)
    UpperCAmelCase__			:		List[str]													=				{}

    # Need to change some parameters name to match Flax names
    for pt_key, pt_tensor in pt_state_dict.items():
        UpperCAmelCase__			:		Optional[int]													=				rename_key(lowerCAmelCase__						)
        UpperCAmelCase__			:		str													=				tuple(renamed_pt_key.split('''.'''						)						)

        # Correctly rename weight parameters
        UpperCAmelCase__						,						UpperCAmelCase__			:		Optional[int]													=				rename_key_and_reshape_tensor(lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)

        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}."""						)

        # also add unexpected weight so that warning is thrown
        UpperCAmelCase__			:		List[str]													=				jnp.asarray(lowerCAmelCase__						)

    return unflatten_dict(lowerCAmelCase__						)
