


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_				:			Any   =	logging.get_logger(__name__)







def 	SCREAMING_SNAKE_CASE_						(      __A			:      List[str]  ,   __A			:      Union[str, Any]  ,   __A			:      List[Any]    )     ->		int:

       """simple docstring"""

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

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

       """simple docstring"""

       logger.info(F"""Loading model based on config from {config_path}..."""    )
       a_ :		Optional[Any]							= BertConfig.from_json_file(__A    )
       a_ :		List[str]							= BertModel(__A    )

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

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


if __name__ == "__main__":
   UpperCAmelCase_				:			Optional[int]   =	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_				:			Dict   =	parser.parse_args()
   convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)






import subprocess
import sys

from transformers import BertConfig, BertModel, BertTokenizer, pipeline
from transformers.testing_utils import TestCasePlus, require_torch
class 			SCREAMING_SNAKE_CASE__  (    lowercase__      ):
      @require_torch
      def 	SCREAMING_SNAKE_CASE							(	self    :	Optional[Any]					)		->      List[Any]:
             # this test is a bit tricky since TRANSFORMERS_OFFLINE can only be changed before
             # `transformers` is loaded, and it's too late for inside pytest - so we are changing it
             # while running an external program

             # python one-liner segments

             # this must be loaded before socket.socket is monkey-patched
             a_ :		str							=  '\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

             a_ :		Any							=  '\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

             a_ :		Dict							=  '\nimport socket\ndef offline_socket(*args, **kwargs): raise RuntimeError("Offline mode is enabled, we shouldn\'t access internet")\nsocket.socket = offline_socket\n        '

             # Force fetching the files so that we can use the cache
             a_ :		Dict							=  'hf-internal-testing/tiny-random-bert'
             BertConfig.from_pretrained(SCREAMING_SNAKE_CASE__					)
             BertModel.from_pretrained(SCREAMING_SNAKE_CASE__					)
             BertTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__					)
             pipeline(task='fill-mask'					,							model=SCREAMING_SNAKE_CASE__					)

             # baseline - just load from_pretrained with normal network
             a_ :		Union[str, Any]							=  [sys.executable, '-c', '\n'.join([load, run, mock]					)]

             # should succeed
             a_ :		List[str]							=  self.get_env()
             # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
             a_ :		Any							=  '1'
             a_ :		Optional[Any]							=  subprocess.run(SCREAMING_SNAKE_CASE__					,							env=SCREAMING_SNAKE_CASE__					,							check=SCREAMING_SNAKE_CASE__					,							capture_output=SCREAMING_SNAKE_CASE__					)
             self.assertEqual(result.returncode					,							0					,							result.stderr					)
             self.assertIn('success'					,							result.stdout.decode()					)
      @require_torch
      def 	SCREAMING_SNAKE_CASE							(	self    :	str					)		->      List[Any]:
             # python one-liner segments
             # this must be loaded before socket.socket is monkey-patched
             a_ :		Tuple							=  '\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

             a_ :		str							=  '\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

             a_ :		Any							=  '\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Faking flaky internet")\nsocket.socket = offline_socket\n        '

             # Force fetching the files so that we can use the cache
             a_ :		Tuple							=  'hf-internal-testing/tiny-random-bert'
             BertConfig.from_pretrained(SCREAMING_SNAKE_CASE__					)
             BertModel.from_pretrained(SCREAMING_SNAKE_CASE__					)
             BertTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__					)
             pipeline(task='fill-mask'					,							model=SCREAMING_SNAKE_CASE__					)

             # baseline - just load from_pretrained with normal network
             a_ :		Optional[Any]							=  [sys.executable, '-c', '\n'.join([load, run, mock]					)]

             # should succeed
             a_ :		Dict							=  self.get_env()
             a_ :		Dict							=  subprocess.run(SCREAMING_SNAKE_CASE__					,							env=SCREAMING_SNAKE_CASE__					,							check=SCREAMING_SNAKE_CASE__					,							capture_output=SCREAMING_SNAKE_CASE__					)
             self.assertEqual(result.returncode					,							0					,							result.stderr					)
             self.assertIn('success'					,							result.stdout.decode()					)
      @require_torch
      def 	SCREAMING_SNAKE_CASE							(	self    :	int					)		->      Union[str, Any]:
             # this test is a bit tricky since TRANSFORMERS_OFFLINE can only be changed before
             # `transformers` is loaded, and it's too late for inside pytest - so we are changing it
             # while running an external program

             # python one-liner segments

             # this must be loaded before socket.socket is monkey-patched
             a_ :		Any							=  '\nfrom transformers import BertConfig, BertModel, BertTokenizer\n        '

             a_ :		Dict							=  '\nmname = "hf-internal-testing/tiny-random-bert-sharded"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nprint("success")\n        '

             a_ :		Any							=  '\nimport socket\ndef offline_socket(*args, **kwargs): raise ValueError("Offline mode is enabled")\nsocket.socket = offline_socket\n        '

             # baseline - just load from_pretrained with normal network
             a_ :		Any							=  [sys.executable, '-c', '\n'.join([load, run]					)]

             # should succeed
             a_ :		Tuple							=  self.get_env()
             a_ :		int							=  subprocess.run(SCREAMING_SNAKE_CASE__					,							env=SCREAMING_SNAKE_CASE__					,							check=SCREAMING_SNAKE_CASE__					,							capture_output=SCREAMING_SNAKE_CASE__					)
             self.assertEqual(result.returncode					,							0					,							result.stderr					)
             self.assertIn('success'					,							result.stdout.decode()					)

             # next emulate no network
             a_ :		str							=  [sys.executable, '-c', '\n'.join([load, mock, run]					)]

             # Doesn't fail anymore since the model is in the cache due to other tests, so commenting this.
             # env["TRANSFORMERS_OFFLINE"] = "0"
             # result = subprocess.run(cmd, env=env, check=False, capture_output=True)
             # self.assertEqual(result.returncode, 1, result.stderr)

             # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
             a_ :		Dict							=  '1'
             a_ :		int							=  subprocess.run(SCREAMING_SNAKE_CASE__					,							env=SCREAMING_SNAKE_CASE__					,							check=SCREAMING_SNAKE_CASE__					,							capture_output=SCREAMING_SNAKE_CASE__					)
             self.assertEqual(result.returncode					,							0					,							result.stderr					)
             self.assertIn('success'					,							result.stdout.decode()					)
      @require_torch
      def 	SCREAMING_SNAKE_CASE							(	self    :	Optional[int]					)		->      Union[str, Any]:
             a_ :		Union[str, Any]							=  '\nfrom transformers import pipeline\n        '
             a_ :		Optional[Any]							=  '\nmname = "hf-internal-testing/tiny-random-bert"\npipe = pipeline(model=mname)\n        '

             a_ :		Any							=  '\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Offline mode is enabled")\nsocket.socket = offline_socket\n        '
             a_ :		Union[str, Any]							=  self.get_env()
             a_ :		Optional[Any]							=  '1'
             a_ :		int							=  [sys.executable, '-c', '\n'.join([load, mock, run]					)]
             a_ :		Union[str, Any]							=  subprocess.run(SCREAMING_SNAKE_CASE__					,							env=SCREAMING_SNAKE_CASE__					,							check=SCREAMING_SNAKE_CASE__					,							capture_output=SCREAMING_SNAKE_CASE__					)
             self.assertEqual(result.returncode					,							1					,							result.stderr					)
             self.assertIn(
                 'You cannot infer task automatically within `pipeline` when using offline mode'					,							result.stderr.decode().replace('\n'					,							''					)					,							)

      @require_torch
      def 	SCREAMING_SNAKE_CASE							(	self    :	int					)		->      Dict:
             a_ :		Optional[int]							=  '\nfrom transformers import AutoModel\n        '
             a_ :		Dict							=  '\nmname = "hf-internal-testing/test_dynamic_model"\nAutoModel.from_pretrained(mname, trust_remote_code=True)\nprint("success")\n        '

             # baseline - just load from_pretrained with normal network
             a_ :		Tuple							=  [sys.executable, '-c', '\n'.join([load, run]					)]

             # should succeed
             a_ :		Optional[Any]							=  self.get_env()
             a_ :		Dict							=  subprocess.run(SCREAMING_SNAKE_CASE__					,							env=SCREAMING_SNAKE_CASE__					,							check=SCREAMING_SNAKE_CASE__					,							capture_output=SCREAMING_SNAKE_CASE__					)
             self.assertEqual(result.returncode					,							0					,							result.stderr					)
             self.assertIn('success'					,							result.stdout.decode()					)

             # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
             a_ :		Optional[int]							=  '1'
             a_ :		Optional[Any]							=  subprocess.run(SCREAMING_SNAKE_CASE__					,							env=SCREAMING_SNAKE_CASE__					,							check=SCREAMING_SNAKE_CASE__					,							capture_output=SCREAMING_SNAKE_CASE__					)
             self.assertEqual(result.returncode					,							0					,							result.stderr					)
             self.assertIn('success'					,							result.stdout.decode()					)


