





'''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()
lowerCAmelCase					:     Any												=logging.get_logger(__name__)
def 							UpperCAmelCase_  (						__lowerCamelCase					:					Optional[Any]	,__lowerCamelCase					:					List[str]	,__lowerCamelCase					:					Tuple ):
   lowercase_			:Dict    			= os.path.abspath(snake_case_ )
   logger.info(F'Converting TensorFlow checkpoint from {tf_path}' )
   # Load weights from TF model
   lowercase_			:Optional[Any]    			= tf.train.list_variables(snake_case_ )
   lowercase_			:int    			= []
   lowercase_			:Union[str, Any]    			= []
   lowercase_			:str    			= []
   for full_name, shape in init_vars:
      # logger.info(f"Loading TF weight {name} with shape {shape}")
      lowercase_			:Optional[int]    			= 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'
         lowercase_			:Union[str, Any]    			= name[1:]
      # figure out how many levels deep the name is
      lowercase_			:Optional[int]    			= 0
      for _name in name:
         if _name.startswith("layer_with_weights" ):
            depth += 1
         else:
            break
      layer_depth.append(snake_case_ )
      # read data
      lowercase_			:Optional[Any]    			= tf.train.load_variable(snake_case_	,snake_case_ )
      names.append("/".join(snake_case_ ) )
      arrays.append(snake_case_ )
   logger.info(F'Read a total of {len(snake_case_ ):,} layers' )

   # Sanity check
   if len(set(snake_case_ ) ) != 1:
      raise ValueError(F'Found layer names with different depths (layer depth {list(set(snake_case_ ) )})' )
   lowercase_			:int    			= list(set(snake_case_ ) )[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(snake_case_	,snake_case_ ):
      lowercase_			:List[str]    			= full_name.split("/" )
      lowercase_			:List[Any]    			= model
      lowercase_			:Tuple    			= []
      for i, m_name in enumerate(snake_case_ ):
         if m_name == ".ATTRIBUTES":
            # variable names end with .ATTRIBUTES/VARIABLE_VALUE
            break
         if m_name.startswith("layer_with_weights" ):
            lowercase_			: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"] )
               lowercase_			:Optional[int]    			= getattr(snake_case_	,"embeddings" )
               lowercase_			:Any    			= getattr(snake_case_	,"LayerNorm" )
            elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
               # encoder layers
               trace.extend(["encoder", "layer", str(layer_num - 4 )] )
               lowercase_			:int    			= getattr(snake_case_	,"encoder" )
               lowercase_			:Optional[Any]    			= getattr(snake_case_	,"layer" )
               lowercase_			:Optional[Any]    			= pointer[layer_num - 4]
            elif layer_num == config.num_hidden_layers + 4:
               # pooler layer
               trace.extend(["pooler", "dense"] )
               lowercase_			:int    			= getattr(snake_case_	,"pooler" )
               lowercase_			:Any    			= getattr(snake_case_	,"dense" )
         elif m_name == "embeddings":
            trace.append("embeddings" )
            lowercase_			:Any    			= getattr(snake_case_	,"embeddings" )
            if layer_num == 0:
               trace.append("word_embeddings" )
               lowercase_			:Dict    			= getattr(snake_case_	,"word_embeddings" )
            elif layer_num == 1:
               trace.append("position_embeddings" )
               lowercase_			:List[Any]    			= getattr(snake_case_	,"position_embeddings" )
            elif layer_num == 2:
               trace.append("token_type_embeddings" )
               lowercase_			:Union[str, Any]    			= getattr(snake_case_	,"token_type_embeddings" )
            else:
               raise ValueError(F'Unknown embedding layer with name {full_name}' )
            trace.append("weight" )
            lowercase_			:Any    			= getattr(snake_case_	,"weight" )
         elif m_name == "_attention_layer":
            # self-attention layer
            trace.extend(["attention", "self"] )
            lowercase_			:List[Any]    			= getattr(snake_case_	,"attention" )
            lowercase_			:Dict    			= getattr(snake_case_	,"self" )
         elif m_name == "_attention_layer_norm":
            # output attention norm
            trace.extend(["attention", "output", "LayerNorm"] )
            lowercase_			:List[str]    			= getattr(snake_case_	,"attention" )
            lowercase_			:List[str]    			= getattr(snake_case_	,"output" )
            lowercase_			:Tuple    			= getattr(snake_case_	,"LayerNorm" )
         elif m_name == "_attention_output_dense":
            # output attention dense
            trace.extend(["attention", "output", "dense"] )
            lowercase_			:int    			= getattr(snake_case_	,"attention" )
            lowercase_			:Union[str, Any]    			= getattr(snake_case_	,"output" )
            lowercase_			:List[str]    			= getattr(snake_case_	,"dense" )
         elif m_name == "_output_dense":
            # output dense
            trace.extend(["output", "dense"] )
            lowercase_			:int    			= getattr(snake_case_	,"output" )
            lowercase_			:Dict    			= getattr(snake_case_	,"dense" )
         elif m_name == "_output_layer_norm":
            # output dense
            trace.extend(["output", "LayerNorm"] )
            lowercase_			:Union[str, Any]    			= getattr(snake_case_	,"output" )
            lowercase_			:Any    			= getattr(snake_case_	,"LayerNorm" )
         elif m_name == "_key_dense":
            # attention key
            trace.append("key" )
            lowercase_			:str    			= getattr(snake_case_	,"key" )
         elif m_name == "_query_dense":
            # attention query
            trace.append("query" )
            lowercase_			:Any    			= getattr(snake_case_	,"query" )
         elif m_name == "_value_dense":
            # attention value
            trace.append("value" )
            lowercase_			:int    			= getattr(snake_case_	,"value" )
         elif m_name == "_intermediate_dense":
            # attention intermediate dense
            trace.extend(["intermediate", "dense"] )
            lowercase_			:Union[str, Any]    			= getattr(snake_case_	,"intermediate" )
            lowercase_			:List[Any]    			= getattr(snake_case_	,"dense" )
         elif m_name == "_output_layer_norm":
            # output layer norm
            trace.append("output" )
            lowercase_			:Optional[int]    			= getattr(snake_case_	,"output" )
         # weights & biases
         elif m_name in ["bias", "beta"]:
            trace.append("bias" )
            lowercase_			:Tuple    			= getattr(snake_case_	,"bias" )
         elif m_name in ["kernel", "gamma"]:
            trace.append("weight" )
            lowercase_			:Dict    			= getattr(snake_case_	,"weight" )
         else:
            logger.warning(F'Ignored {m_name}' )
        # for certain layers reshape is necessary
      lowercase_			:Dict    			= '''.'''.join(snake_case_ )
      if re.match(r"(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)"	,snake_case_ ) or re.match(
          r"(\S+)\.attention\.output\.dense\.weight"	,snake_case_ ):
         lowercase_			:Any    			= array.reshape(pointer.data.shape )
      if "kernel" in full_name:
         lowercase_			:List[Any]    			= array.transpose()
      if pointer.shape == array.shape:
         lowercase_			:str    			= torch.from_numpy(snake_case_ )
      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 							UpperCAmelCase_  (						__lowerCamelCase					:					List[Any]	,__lowerCamelCase					:					int	,__lowerCamelCase					:					Union[str, Any] ):
   logger.info(F'Loading model based on config from {config_path}...' )
   lowercase_			:Optional[Any]    			= BertConfig.from_json_file(snake_case_ )
   lowercase_			:List[str]    			= BertModel(snake_case_ )

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

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


if __name__ == "__main__":
 lowerCAmelCase					:     List[Any]												=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).''',
 )
 lowerCAmelCase					:     str												=parser.parse_args()
 convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)









"""simple docstring"""



import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
       import os

       import jax.numpy as jnp
       from jax import jit

       from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
       from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

       a_ =			"""0.12"""  # assumed parallelism: 8


if is_torch_available():
       import torch







def 		__lowercase     (   snake_case_  :  Any				,snake_case_  :  int				,snake_case_  :  List[str]=None						)   ->Union[str, Any]:



    '''simple docstring'''







    if rng is None:
        __A    :						Union[str, Any]						 =     random.Random()

    __A    :						List[str]						 =     1
    for dim in shape:
        total_dims *= dim

    __A    :						List[str]						 =     []
    for _ in range(snake_case_						):
        values.append(rng.randint(0				,vocab_size - 1						)						)

    __A    :						Optional[Any]						 =     np.array(snake_case_				,dtype=jnp.intaa						).reshape(snake_case_						)

    return output







def 		__lowercase     (   snake_case_  :  int				,snake_case_  :  Union[str, Any]=None						)   ->int:



    '''simple docstring'''







    __A    :						int						 =     ids_tensor(snake_case_				,vocab_size=2				,rng=snake_case_						)
    # make sure that at least one token is attended to for each batch
    __A    :						str						 =     1
    return attn_mask






@require_flax
class 		__snake_case					:





 """simple docstring"""


 _lowerCamelCase        =     None
 _lowerCamelCase        =     ()

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A    :						Optional[int]						 =     self.model_tester.prepare_config_and_inputs_for_common()

     # cut to half length & take max batch_size 3
     __A    :						Optional[int]						 =     2
     __A    :						Tuple						 =     inputs['''input_ids'''].shape[-1] // 2
     __A    :						List[Any]						 =     inputs['''input_ids'''][:max_batch_size, :sequence_length]

     __A    :						Optional[Any]						 =     jnp.ones_like(__lowerCamelCase  )
     __A    :						List[Any]						 =     attention_mask[:max_batch_size, :sequence_length]

     # generate max 5 tokens
     __A    :						Optional[int]						 =     input_ids.shape[-1] + 5
     if config.eos_token_id is not None and config.pad_token_id is None:
         # hack to allow generate for models such as GPT2 as is done in `generate()`
         __A    :						Optional[int]						 =     config.eos_token_id
     return config, input_ids, attention_mask, max_length

 @is_pt_flax_cross_test
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						List[Any]						 =     self._get_input_ids_and_config()
     __A    :						List[Any]						 =     False
     __A    :						str						 =     max_length
     __A    :						Tuple						 =     0

     for model_class in self.all_generative_model_classes:
         __A    :						Union[str, Any]						 =     model_class(__lowerCamelCase  )

         __A    :						Any						 =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
         __A    :						Dict						 =     getattr(__lowerCamelCase					,						__lowerCamelCase  )
         __A    :						Tuple						 =     pt_model_class(__lowerCamelCase  ).eval()
         __A    :						Optional[Any]						 =     load_flax_weights_in_pytorch_model(__lowerCamelCase					,						flax_model.params  )

         __A    :						str						 =     flax_model.generate(__lowerCamelCase  ).sequences
         __A    :						List[Any]						 =     pt_model.generate(torch.tensor(__lowerCamelCase					,						dtype=torch.long  )  )

         if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
             __A    :						Any						 =     flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

         self.assertListEqual(pt_generation_outputs.numpy().tolist()					,						flax_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Union[str, Any]						 =     self._get_input_ids_and_config()
     __A    :						List[str]						 =     False
     __A    :						Optional[Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Optional[int]						 =     model_class(__lowerCamelCase  )

         __A    :						Dict						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()
     __A    :						List[Any]						 =     True
     __A    :						Union[str, Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Dict						 =     model_class(__lowerCamelCase  )

         __A    :						List[Any]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						List[Any]						 =     jit(model.generate  )
         __A    :						List[str]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						List[Any]						 =     self._get_input_ids_and_config()
     __A    :						List[str]						 =     False
     __A    :						Any						 =     max_length
     __A    :						List[Any]						 =     2

     for model_class in self.all_generative_model_classes:
         __A    :						Tuple						 =     model_class(__lowerCamelCase  )

         __A    :						Optional[Any]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     jit(model.generate  )
         __A    :						Optional[int]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Tuple						 =     self._get_input_ids_and_config()
     __A    :						str						 =     False
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[Any]						 =     2
     __A    :						Any						 =     2

     for model_class in self.all_generative_model_classes:
         __A    :						List[str]						 =     model_class(__lowerCamelCase  )

         __A    :						Optional[int]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[0]					,						input_ids.shape[0] * config.num_return_sequences  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Any						 =     self._get_input_ids_and_config()
     __A    :						Optional[Any]						 =     True
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[str]						 =     0.8
     __A    :						List[str]						 =     10
     __A    :						Union[str, Any]						 =     0.3
     __A    :						Union[str, Any]						 =     1
     __A    :						Optional[Any]						 =     8
     __A    :						Dict						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						List[Any]						 =     model_class(__lowerCamelCase  )

         __A    :						List[str]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Any						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Union[str, Any]						 =     self._get_input_ids_and_config()
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[str]						 =     1
     __A    :						str						 =     8
     __A    :						Any						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						Union[str, Any]						 =     model_class(__lowerCamelCase  )

         __A    :						Tuple						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Optional[Any]						 =     jit(model.generate  )
         __A    :						Optional[Any]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()
     __A    :						Optional[int]						 =     max_length
     __A    :						List[str]						 =     2
     __A    :						List[Any]						 =     1
     __A    :						Optional[Any]						 =     8
     __A    :						str						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						Optional[int]						 =     model_class(__lowerCamelCase  )

         __A    :						int						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						str						 =     jit(model.generate  )
         __A    :						Any						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						Dict						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						Optional[Any]						 =     False
     __A    :						Union[str, Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						int						 =     model_class(__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Optional[int]						 =     jit(model.generate  )
         __A    :						Any						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Optional[int]						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						str						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						List[Any]						 =     True
     __A    :						Any						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						str						 =     model_class(__lowerCamelCase  )

         __A    :						List[Any]						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						List[str]						 =     jit(model.generate  )
         __A    :						Optional[int]						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Optional[Any]						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						List[str]						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						Optional[int]						 =     2
     __A    :						Dict						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Any						 =     model_class(__lowerCamelCase  )

         __A    :						int						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						str						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )






@require_flax
class 		__snake_case					(		unittest.TestCase     ):





 """simple docstring"""


 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Optional[int]						 =     AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''  )
     __A    :						Optional[Any]						 =     FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''  )

     __A    :						Any						 =     '''Hello world'''
     __A    :						Dict						 =     tokenizer(__lowerCamelCase					,						return_tensors='''np'''  ).input_ids

     # typos are quickly detected (the correct argument is `do_sample`)
     with self.assertRaisesRegex(__lowerCamelCase					,						'''do_samples'''  ):
         model.generate(__lowerCamelCase					,						do_samples=__lowerCamelCase  )

     # arbitrary arguments that will not be used anywhere are also not accepted
     with self.assertRaisesRegex(__lowerCamelCase					,						'''foo'''  ):
         __A    :						Any						 =     {'''foo''': '''bar'''}
         model.generate(__lowerCamelCase					,						**__lowerCamelCase  )




