



import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
							from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_  ):
			if is_torch_version("<"						,  "2.0.0"  ) or not hasattr(SCREAMING_SNAKE_CASE_						,  "_dynamo"  ):
						return False
			return isinstance(SCREAMING_SNAKE_CASE_						,  torch._dynamo.eval_frame.OptimizedModule  )



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_ = True  ):
			lowercase__					  =				(torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

			lowercase__					  =				is_compiled_module(SCREAMING_SNAKE_CASE_  )
			if is_compiled:
						lowercase__					  =				model
						lowercase__					  =				model._orig_mod

			if is_deepspeed_available():
						options += (DeepSpeedEngine,)

			while isinstance(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  ):
						lowercase__					  =				model.module

			if not keep_fpaa_wrapper:
						lowercase__					  =				getattr(SCREAMING_SNAKE_CASE_						,  "forward"  )
						lowercase__					  =				model.__dict__.pop("_original_forward"						,  SCREAMING_SNAKE_CASE_  )
						if original_forward is not None:
									while hasattr(SCREAMING_SNAKE_CASE_						,  "__wrapped__"  ):
												lowercase__					  =				forward.__wrapped__
												if forward == original_forward:
															break
									lowercase__					  =				forward
						if getattr(SCREAMING_SNAKE_CASE_						,  "_converted_to_transformer_engine"						,  SCREAMING_SNAKE_CASE_  ):
									convert_model(SCREAMING_SNAKE_CASE_						,  to_transformer_engine=SCREAMING_SNAKE_CASE_  )

			if is_compiled:
						lowercase__					  =				model
						lowercase__					  =				compiled_model

			return model



def __lowerCAmelCase (     ):
			PartialState().wait_for_everyone()



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  ):
			if PartialState().distributed_type == DistributedType.TPU:
						xm.save(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  )
			elif PartialState().local_process_index == 0:
						torch.save(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  )



@contextmanager
def __lowerCAmelCase (     **SCREAMING_SNAKE_CASE_  ):
			for key, value in kwargs.items():
						lowercase__					  =				str(SCREAMING_SNAKE_CASE_  )

			yield

			for key in kwargs:
						if key.upper() in os.environ:
									del os.environ[key.upper()]



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_  ):
			if not hasattr(SCREAMING_SNAKE_CASE_						,  "__qualname__"  ) and not hasattr(SCREAMING_SNAKE_CASE_						,  "__name__"  ):
						lowercase__					  =				getattr(SCREAMING_SNAKE_CASE_						,  "__class__"						,  SCREAMING_SNAKE_CASE_  )
			if hasattr(SCREAMING_SNAKE_CASE_						,  "__qualname__"  ):
						return obj.__qualname__
			if hasattr(SCREAMING_SNAKE_CASE_						,  "__name__"  ):
						return obj.__name__
			return str(SCREAMING_SNAKE_CASE_  )



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  ):
			for key, value in source.items():
						if isinstance(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  ):
									lowercase__					  =				destination.setdefault(SCREAMING_SNAKE_CASE_						,  {}  )
									merge_dicts(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  )
						else:
									lowercase__					  =				value

			return destination




def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_ = None  ):
			if port is None:
						lowercase__					  =				2_9500
			with socket.socket(socket.AF_INET						,  socket.SOCK_STREAM  ) as s:
						return s.connect_ex(("localhost", port)  ) == 0









import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  ):
			lowercase__					  =				("dense.weight", "attention.self.query", "attention.self.key", "attention.self.value")

			lowercase__					  =				(
			    ("layer.", "layer_"),
			    ("word_embeddings.weight", "word_embeddings"),
			    ("position_embeddings.weight", "position_embeddings"),
			    ("token_type_embeddings.weight", "token_type_embeddings"),
			    (".", "/"),
			    ("LayerNorm/weight", "LayerNorm/gamma"),
			    ("LayerNorm/bias", "LayerNorm/beta"),
			    ("weight", "kernel"),
			)

			if not os.path.isdir(SCREAMING_SNAKE_CASE_  ):
						os.makedirs(SCREAMING_SNAKE_CASE_  )

			lowercase__					  =				model.state_dict()

			def to_tf_var_name(SCREAMING_SNAKE_CASE_  ):
						for patt, repl in iter(SCREAMING_SNAKE_CASE_  ):
									lowercase__					  =				name.replace(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  )
						return f'''bert/{name}'''

			def create_tf_var(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  ):
						lowercase__					  =				tf.dtypes.as_dtype(tensor.dtype  )
						lowercase__					  =				tf.get_variable(dtype=SCREAMING_SNAKE_CASE_						,  shape=tensor.shape						,  name=SCREAMING_SNAKE_CASE_						,  initializer=tf.zeros_initializer()  )
						session.run(tf.variables_initializer([tf_var]  )  )
						session.run(SCREAMING_SNAKE_CASE_  )
						return tf_var

			tf.reset_default_graph()
			with tf.Session() as session:
						for var_name in state_dict:
									lowercase__					  =				to_tf_var_name(SCREAMING_SNAKE_CASE_  )
									lowercase__					  =				state_dict[var_name].numpy()
									if any(x in var_name for x in tensors_to_transpose  ):
												lowercase__					  =				torch_tensor.T
									lowercase__					  =				create_tf_var(tensor=SCREAMING_SNAKE_CASE_						,  name=SCREAMING_SNAKE_CASE_						,  session=SCREAMING_SNAKE_CASE_  )
									tf.keras.backend.set_value(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  )
									lowercase__					  =				session.run(SCREAMING_SNAKE_CASE_  )
									print(f'''Successfully created {tf_name}: {np.allclose(SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  )}'''  )

						lowercase__					  =				tf.train.Saver(tf.trainable_variables()  )
						saver.save(SCREAMING_SNAKE_CASE_						,  os.path.join(SCREAMING_SNAKE_CASE_						,  model_name.replace("-"						,  "_"  ) + ".ckpt"  )  )



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_=None  ):
			lowercase__					  =				argparse.ArgumentParser()
			parser.add_argument("--model_name"						,  type=SCREAMING_SNAKE_CASE_						,  required=SCREAMING_SNAKE_CASE_						,  help="model name e.g. bert-base-uncased"  )
			parser.add_argument(
			    "--cache_dir"						,  type=SCREAMING_SNAKE_CASE_						,  default=SCREAMING_SNAKE_CASE_						,  required=SCREAMING_SNAKE_CASE_						,  help="Directory containing pytorch model"  )
			parser.add_argument("--pytorch_model_path"						,  type=SCREAMING_SNAKE_CASE_						,  required=SCREAMING_SNAKE_CASE_						,  help="/path/to/<pytorch-model-name>.bin"  )
			parser.add_argument("--tf_cache_dir"						,  type=SCREAMING_SNAKE_CASE_						,  required=SCREAMING_SNAKE_CASE_						,  help="Directory in which to save tensorflow model"  )
			lowercase__					  =				parser.parse_args(SCREAMING_SNAKE_CASE_  )

			lowercase__					  =				BertModel.from_pretrained(
			    pretrained_model_name_or_path=args.model_name						,  state_dict=torch.load(args.pytorch_model_path  )						,  cache_dir=args.cache_dir						,  )

			convert_pytorch_checkpoint_to_tf(model=SCREAMING_SNAKE_CASE_						,  ckpt_dir=args.tf_cache_dir						,  model_name=args.model_name  )


if __name__ == "__main__":
							main()




