
'''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()
SCREAMING_SNAKE_CASE     =       logging.get_logger(__name__)







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

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

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

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


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE     =       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).""",
      )
      SCREAMING_SNAKE_CASE     =       parser.parse_args()
      convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)





'''simple docstring'''





import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
      import jax

      from transformers.models.big_bird.modeling_flax_big_bird import (
          FlaxBigBirdForCausalLM,
          FlaxBigBirdForMaskedLM,
          FlaxBigBirdForMultipleChoice,
          FlaxBigBirdForPreTraining,
          FlaxBigBirdForQuestionAnswering,
          FlaxBigBirdForSequenceClassification,
          FlaxBigBirdForTokenClassification,
          FlaxBigBirdModel,
      )






class 							UpperCAmelCase_							(    unittest.TestCase				):

 '''simple docstring'''







 def __init__(					self		:  Tuple			,				snake_case__		:  List[str]			,				snake_case__		:  Union[str, Any]=2			,				snake_case__		:  int=56			,				snake_case__		:  Any=True			,				snake_case__		:  Union[str, Any]=True			,				snake_case__		:  Optional[int]=True			,				snake_case__		:  Any=True			,				snake_case__		:  Tuple=99			,				snake_case__		:  Any=32			,				snake_case__		:  List[Any]=2			,				snake_case__		:  Optional[Any]=2			,				snake_case__		:  Any=7			,				snake_case__		:  Union[str, Any]="gelu_new"			,				snake_case__		:  List[Any]=0.1			,				snake_case__		:  Optional[Any]=0.1			,				snake_case__		:  Optional[int]=5_12			,				snake_case__		:  Optional[int]=16			,				snake_case__		:  int=2			,				snake_case__		:  List[Any]=0.02			,				snake_case__		:  int=4			,				snake_case__		:  Optional[Any]="block_sparse"			,				snake_case__		:  Union[str, Any]=True			,				snake_case__		:  List[Any]=False			,				snake_case__		:  Optional[Any]=2			,				snake_case__		:  Optional[Any]=3			,				):



   '''simple docstring'''




   UpperCAmelCase__      :					List[Any]				       =							parent
   UpperCAmelCase__      :					Dict				       =							batch_size
   UpperCAmelCase__      :					Tuple				       =							seq_length
   UpperCAmelCase__      :					Optional[int]				       =							is_training
   UpperCAmelCase__      :					Optional[int]				       =							use_attention_mask
   UpperCAmelCase__      :					List[Any]				       =							use_token_type_ids
   UpperCAmelCase__      :					Dict				       =							use_labels
   UpperCAmelCase__      :					Tuple				       =							vocab_size
   UpperCAmelCase__      :					str				       =							hidden_size
   UpperCAmelCase__      :					List[str]				       =							num_hidden_layers
   UpperCAmelCase__      :					Union[str, Any]				       =							num_attention_heads
   UpperCAmelCase__      :					Optional[int]				       =							intermediate_size
   UpperCAmelCase__      :					Optional[int]				       =							hidden_act
   UpperCAmelCase__      :					Union[str, Any]				       =							hidden_dropout_prob
   UpperCAmelCase__      :					str				       =							attention_probs_dropout_prob
   UpperCAmelCase__      :					str				       =							max_position_embeddings
   UpperCAmelCase__      :					Any				       =							type_vocab_size
   UpperCAmelCase__      :					Dict				       =							type_sequence_label_size
   UpperCAmelCase__      :					Dict				       =							initializer_range
   UpperCAmelCase__      :					Any				       =							num_choices

   UpperCAmelCase__      :					List[str]				       =							rescale_embeddings
   UpperCAmelCase__      :					Union[str, Any]				       =							attention_type
   UpperCAmelCase__      :					Union[str, Any]				       =							use_bias
   UpperCAmelCase__      :					Union[str, Any]				       =							block_size
   UpperCAmelCase__      :					Any				       =							num_random_blocks







 def 							UpperCamelCase  (					self		:  Optional[Any] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Optional[int]				       =							ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size )

   UpperCAmelCase__      :					Optional[Any]				       =							None
   if self.use_attention_mask:
     UpperCAmelCase__      :					List[Any]				       =							random_attention_mask([self.batch_size, self.seq_length] )

   UpperCAmelCase__      :					Union[str, Any]				       =							None
   if self.use_token_type_ids:
     UpperCAmelCase__      :					Any				       =							ids_tensor([self.batch_size, self.seq_length]			,				self.type_vocab_size )

   UpperCAmelCase__      :					Union[str, Any]				       =							BigBirdConfig(
       vocab_size=self.vocab_size			,				hidden_size=self.hidden_size			,				num_hidden_layers=self.num_hidden_layers			,				num_attention_heads=self.num_attention_heads			,				intermediate_size=self.intermediate_size			,				hidden_act=self.hidden_act			,				hidden_dropout_prob=self.hidden_dropout_prob			,				attention_probs_dropout_prob=self.attention_probs_dropout_prob			,				max_position_embeddings=self.max_position_embeddings			,				type_vocab_size=self.type_vocab_size			,				is_decoder=snake_case__			,				initializer_range=self.initializer_range			,				attention_type=self.attention_type			,				block_size=self.block_size			,				num_random_blocks=self.num_random_blocks			,				use_bias=self.use_bias			,				rescale_embeddings=self.rescale_embeddings			,				)

   return config, input_ids, token_type_ids, attention_mask







 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   UpperCAmelCase__      :					int				       =							self.prepare_config_and_inputs()
   UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__      :					Tuple				       =							config_and_inputs
   UpperCAmelCase__      :					str				       =							{
       "input_ids": input_ids,
       "token_type_ids": token_type_ids,
       "attention_mask": attention_mask,
   }
   return config, inputs_dict






@require_flax
class 							UpperCAmelCase_							(    A					,				unittest.TestCase				):

 '''simple docstring'''
 lowercase_					:    Any				=							(
     (
         FlaxBigBirdForCausalLM,
         FlaxBigBirdModel,
         FlaxBigBirdForPreTraining,
         FlaxBigBirdForMaskedLM,
         FlaxBigBirdForMultipleChoice,
         FlaxBigBirdForQuestionAnswering,
         FlaxBigBirdForSequenceClassification,
         FlaxBigBirdForTokenClassification,
     )
     if is_flax_available()
     else ()
 )

 lowercase_					:    Optional[Any]				=							False
 lowercase_					:    str				=							False







 def 							UpperCamelCase  (					self		:  Optional[Any] ):



   '''simple docstring'''




   UpperCAmelCase__      :					int				       =							FlaxBigBirdModelTester(self )







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   super().test_from_pretrained_save_pretrained()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   super().test_from_pretrained_with_no_automatic_init()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   super().test_no_automatic_init()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   super().test_hidden_states_output()







 @slow
 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   for model_class_name in self.all_model_classes:
     UpperCAmelCase__      :					List[Any]				       =							model_class_name.from_pretrained("google/bigbird-roberta-base" )
     self.assertIsNotNone(snake_case__ )







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



   '''simple docstring'''




   if self.test_attn_probs:
     super().test_attention_outputs()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 							UpperCamelCase  (					self		:  int ):



   '''simple docstring'''




   UpperCAmelCase__       ,			UpperCAmelCase__      :					str				       =							self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__ ):
       UpperCAmelCase__      :					Union[str, Any]				       =							self._prepare_for_class(snake_case__			,				snake_case__ )
       UpperCAmelCase__      :					List[Any]				       =							model_class(snake_case__ )

       @jax.jit
       def model_jitted(snake_case__		:  int			,				snake_case__		:  Optional[int]=None			,				**snake_case__		:  Optional[Any] ):
         return model(input_ids=snake_case__			,				attention_mask=snake_case__			,				**snake_case__ )

       with self.subTest("JIT Enabled" ):
         UpperCAmelCase__      :					List[Any]				       =							model_jitted(**snake_case__ ).to_tuple()

       with self.subTest("JIT Disabled" ):
         with jax.disable_jit():
           UpperCAmelCase__      :					Tuple				       =							model_jitted(**snake_case__ ).to_tuple()

       self.assertEqual(len(snake_case__ )			,				len(snake_case__ ) )
       for jitted_output, output in zip(snake_case__			,				snake_case__ ):
         self.assertEqual(jitted_output.shape			,				output.shape )







 def 							UpperCamelCase  (					self		:  List[str]			,				snake_case__		:  int			,				snake_case__		:  Dict			,				snake_case__		:  Any			,				snake_case__		:  Union[str, Any]=1e-5			,				snake_case__		:  Any="outputs"			,				snake_case__		:  Optional[Any]=None ):



   '''simple docstring'''




   if name.startswith("outputs.attentions" ):
     return
   else:
     super().check_pt_flax_outputs(snake_case__			,				snake_case__			,				snake_case__			,				snake_case__			,				snake_case__			,				snake_case__ )



