

import argparse

import tensorflow as tf
import torch

from transformers import BertConfig, BertForMaskedLM
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertPooler,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


logging.set_verbosity_info()
def      _lowerCAmelCase(	__A   ,       __A   ,       __A			):
  def get_masked_lm_array(__A			):
    UpperCAmelCase					      =  F"masked_lm/{name}/.ATTRIBUTES/VARIABLE_VALUE"
    UpperCAmelCase					      =  tf.train.load_variable(__A   ,       __A			)

    if "kernel" in name:
      UpperCAmelCase					      =  array.transpose()

    return torch.from_numpy(__A			)

  def get_encoder_array(__A			):
    UpperCAmelCase					      =  F"encoder/{name}/.ATTRIBUTES/VARIABLE_VALUE"
    UpperCAmelCase					      =  tf.train.load_variable(__A   ,       __A			)

    if "kernel" in name:
      UpperCAmelCase					      =  array.transpose()

    return torch.from_numpy(__A			)

  def get_encoder_layer_array(__A   ,       __A			):
    UpperCAmelCase					      =  F"encoder/_transformer_layers/{layer_index}/{name}/.ATTRIBUTES/VARIABLE_VALUE"
    UpperCAmelCase					      =  tf.train.load_variable(__A   ,       __A			)

    if "kernel" in name:
      UpperCAmelCase					      =  array.transpose()

    return torch.from_numpy(__A			)

  def get_encoder_attention_layer_array(__A   ,       __A   ,       __A			):
    UpperCAmelCase					      =  F"encoder/_transformer_layers/{layer_index}/_attention_layer/{name}/.ATTRIBUTES/VARIABLE_VALUE"
    UpperCAmelCase					      =  tf.train.load_variable(__A   ,       __A			)
    UpperCAmelCase					      =  array.reshape(__A			)

    if "kernel" in name:
      UpperCAmelCase					      =  array.transpose()

    return torch.from_numpy(__A			)

  print(F"Loading model based on config from {config_path}..."			)
  UpperCAmelCase					      =  BertConfig.from_json_file(__A			)
  UpperCAmelCase					      =  BertForMaskedLM(__A			)

  # Layers
  for layer_index in range(0   ,       config.num_hidden_layers			):
    UpperCAmelCase					      =  model.bert.encoder.layer[layer_index]

    # Self-attention
    UpperCAmelCase					      =  layer.attention.self

    UpperCAmelCase					      =  get_encoder_attention_layer_array(
        __A   ,       "_query_dense/kernel"   ,       self_attn.query.weight.data.shape			)
    UpperCAmelCase					      =  get_encoder_attention_layer_array(
        __A   ,       "_query_dense/bias"   ,       self_attn.query.bias.data.shape			)
    UpperCAmelCase					      =  get_encoder_attention_layer_array(
        __A   ,       "_key_dense/kernel"   ,       self_attn.key.weight.data.shape			)
    UpperCAmelCase					      =  get_encoder_attention_layer_array(
        __A   ,       "_key_dense/bias"   ,       self_attn.key.bias.data.shape			)
    UpperCAmelCase					      =  get_encoder_attention_layer_array(
        __A   ,       "_value_dense/kernel"   ,       self_attn.value.weight.data.shape			)
    UpperCAmelCase					      =  get_encoder_attention_layer_array(
        __A   ,       "_value_dense/bias"   ,       self_attn.value.bias.data.shape			)

    # Self-attention Output
    UpperCAmelCase					      =  layer.attention.output

    UpperCAmelCase					      =  get_encoder_attention_layer_array(
        __A   ,       "_output_dense/kernel"   ,       self_output.dense.weight.data.shape			)
    UpperCAmelCase					      =  get_encoder_attention_layer_array(
        __A   ,       "_output_dense/bias"   ,       self_output.dense.bias.data.shape			)

    UpperCAmelCase					      =  get_encoder_layer_array(__A   ,       "_attention_layer_norm/gamma"			)
    UpperCAmelCase					      =  get_encoder_layer_array(__A   ,       "_attention_layer_norm/beta"			)

    # Intermediate
    UpperCAmelCase					      =  layer.intermediate

    UpperCAmelCase					      =  get_encoder_layer_array(__A   ,       "_intermediate_dense/kernel"			)
    UpperCAmelCase					      =  get_encoder_layer_array(__A   ,       "_intermediate_dense/bias"			)

    # Output
    UpperCAmelCase					      =  layer.output

    UpperCAmelCase					      =  get_encoder_layer_array(__A   ,       "_output_dense/kernel"			)
    UpperCAmelCase					      =  get_encoder_layer_array(__A   ,       "_output_dense/bias"			)

    UpperCAmelCase					      =  get_encoder_layer_array(__A   ,       "_output_layer_norm/gamma"			)
    UpperCAmelCase					      =  get_encoder_layer_array(__A   ,       "_output_layer_norm/beta"			)

  # Embeddings
  UpperCAmelCase					      =  get_encoder_array("_position_embedding_layer/embeddings"			)
  UpperCAmelCase					      =  get_encoder_array("_type_embedding_layer/embeddings"			)
  UpperCAmelCase					      =  get_encoder_array("_embedding_norm_layer/gamma"			)
  UpperCAmelCase					      =  get_encoder_array("_embedding_norm_layer/beta"			)

  # LM Head
  UpperCAmelCase					      =  model.cls.predictions.transform

  UpperCAmelCase					      =  get_masked_lm_array("dense/kernel"			)
  UpperCAmelCase					      =  get_masked_lm_array("dense/bias"			)

  UpperCAmelCase					      =  get_masked_lm_array("layer_norm/gamma"			)
  UpperCAmelCase					      =  get_masked_lm_array("layer_norm/beta"			)

  UpperCAmelCase					      =  get_masked_lm_array("embedding_table"			)

  # Pooling
  UpperCAmelCase					      =  BertPooler(config=__A			)
  UpperCAmelCase					      =  get_encoder_array("_pooler_layer/kernel"			)
  UpperCAmelCase					      =  get_encoder_array("_pooler_layer/bias"			)

  # Export final model
  model.save_pretrained(__A			)

  # Integration test - should load without any errors ;)
  UpperCAmelCase					      =  BertForMaskedLM.from_pretrained(__A			)
  print(new_model.eval()			)

  print("Model conversion was done sucessfully!"			)


if __name__ == "__main__":
      lowerCAmelCase__       = argparse.ArgumentParser()
      parser.add_argument(
          "--tf_checkpoint_path", type=str, required=True, help="Path to the TensorFlow Token Dropping 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.",
      )
      lowerCAmelCase__       = parser.parse_args()
      convert_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)






import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / "utils"))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


lowerCAmelCase__       = get_tests_dir("fixtures")






class 						__magic_name__    (       unittest.TestCase    ):


 def _UpperCamelCase	(						self		: Tuple    )    ->      Union[str, Any]:
   # A mock response for an HTTP head request to emulate server down
   UpperCAmelCase					      =  mock.Mock()
   UpperCAmelCase					      =  5_0_0
   UpperCAmelCase					      =  {}
   UpperCAmelCase					      =  HTTPError
   UpperCAmelCase					      =  {}

   # Download this model to make sure it's in the cache.
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"    )
   # Under the mock environment we get a 500 error when trying to reach the model.
   with mock.patch("requests.Session.request"			,						return_value=lowerCAmelCase__    ) as mock_head:
     UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"    )
     # This check we did call the fake head request
     mock_head.assert_called()


 def _UpperCamelCase	(						self		: List[Any]    )    ->      Dict:
   # This test is for deprecated behavior and can be removed in v5
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(
       "https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"    )






@is_staging_test
class 						__magic_name__    (       unittest.TestCase    ):


 @classmethod
 def _UpperCamelCase	(						cls		: List[str]    )    ->      List[Any]:
   UpperCAmelCase					      =  TOKEN
   HfFolder.save_token(lowerCAmelCase__    )


 @classmethod
 def _UpperCamelCase	(						cls		: Optional[int]    )    ->      Union[str, Any]:
   try:
     delete_repo(token=cls._token			,						repo_id="test-feature-extractor"    )
   except HTTPError:
     pass

   try:
     delete_repo(token=cls._token			,						repo_id="valid_org/test-feature-extractor-org"    )
   except HTTPError:
     pass

   try:
     delete_repo(token=cls._token			,						repo_id="test-dynamic-feature-extractor"    )
   except HTTPError:
     pass


 def _UpperCamelCase	(						self		: Any    )    ->      Any:
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(lowerCAmelCase__    )
   feature_extractor.push_to_hub("test-feature-extractor"			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )

   # Reset repo
   delete_repo(token=self._token			,						repo_id="test-feature-extractor"    )

   # Push to hub via save_pretrained
   with tempfile.TemporaryDirectory() as tmp_dir:
     feature_extractor.save_pretrained(
         lowerCAmelCase__			,						repo_id="test-feature-extractor"			,						push_to_hub=lowerCAmelCase__			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )


 def _UpperCamelCase	(						self		: List[Any]    )    ->      Tuple:
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(lowerCAmelCase__    )
   feature_extractor.push_to_hub("valid_org/test-feature-extractor"			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )

   # Reset repo
   delete_repo(token=self._token			,						repo_id="valid_org/test-feature-extractor"    )

   # Push to hub via save_pretrained
   with tempfile.TemporaryDirectory() as tmp_dir:
     feature_extractor.save_pretrained(
         lowerCAmelCase__			,						repo_id="valid_org/test-feature-extractor-org"			,						push_to_hub=lowerCAmelCase__			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor-org"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )


 def _UpperCamelCase	(						self		: Dict    )    ->      List[str]:
   CustomFeatureExtractor.register_for_auto_class()
   UpperCAmelCase					      =  CustomFeatureExtractor.from_pretrained(lowerCAmelCase__    )

   feature_extractor.push_to_hub("test-dynamic-feature-extractor"			,						use_auth_token=self._token    )

   # This has added the proper auto_map field to the config
   self.assertDictEqual(
       feature_extractor.auto_map			,						{"AutoFeatureExtractor": "custom_feature_extraction.CustomFeatureExtractor"}			,						)

   UpperCAmelCase					      =  AutoFeatureExtractor.from_pretrained(
       f"{USER}/test-dynamic-feature-extractor"			,						trust_remote_code=lowerCAmelCase__    )
   # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
   self.assertEqual(new_feature_extractor.__class__.__name__			,						"CustomFeatureExtractor"    )



