





import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase				)    ->  Optional[Any]:
  lowerCamelCase__					:   Union[str, Any]      				=	('dense.weight', 'attention.self.query', 'attention.self.key', 'attention.self.value')

  lowerCamelCase__					:   Optional[int]      				=	(
      ('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(_UpperCAmelCase				):
    os.makedirs(_UpperCAmelCase				)

  lowerCamelCase__					:   Optional[Any]      				=	model.state_dict()

  def to_tf_var_name(_UpperCAmelCase				):
    for patt, repl in iter(_UpperCAmelCase				):
      lowerCamelCase__					:   str      				=	name.replace(_UpperCAmelCase		,			_UpperCAmelCase				)
    return F"""bert/{name}"""

  def create_tf_var(_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase				):
    lowerCamelCase__					:   Optional[int]      				=	tf.dtypes.as_dtype(tensor.dtype				)
    lowerCamelCase__					:   List[Any]      				=	tf.get_variable(dtype=_UpperCAmelCase		,			shape=tensor.shape		,			name=_UpperCAmelCase		,			initializer=tf.zeros_initializer()				)
    session.run(tf.variables_initializer([tf_var]				)				)
    session.run(_UpperCAmelCase				)
    return tf_var

  tf.reset_default_graph()
  with tf.Session() as session:
    for var_name in state_dict:
      lowerCamelCase__					:   Tuple      				=	to_tf_var_name(_UpperCAmelCase				)
      lowerCamelCase__					:   Any      				=	state_dict[var_name].numpy()
      if any(x in var_name for x in tensors_to_transpose				):
        lowerCamelCase__					:   str      				=	torch_tensor.T
      lowerCamelCase__					:   Optional[Any]      				=	create_tf_var(tensor=_UpperCAmelCase		,			name=_UpperCAmelCase		,			session=_UpperCAmelCase				)
      tf.keras.backend.set_value(_UpperCAmelCase		,			_UpperCAmelCase				)
      lowerCamelCase__					:   Dict      				=	session.run(_UpperCAmelCase				)
      print(F"""Successfully created {tf_name}: {np.allclose(_UpperCAmelCase		,			_UpperCAmelCase				)}"""				)

    lowerCamelCase__					:   Dict      				=	tf.train.Saver(tf.trainable_variables()				)
    saver.save(_UpperCAmelCase		,			os.path.join(_UpperCAmelCase		,			model_name.replace('-'		,			'_'				) + '.ckpt'				)				)


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase=None				)    ->  Optional[int]:
  lowerCamelCase__					:   Optional[int]      				=	argparse.ArgumentParser()
  parser.add_argument('--model_name'		,			type=_UpperCAmelCase		,			required=_UpperCAmelCase		,			help='model name e.g. bert-base-uncased'				)
  parser.add_argument(
      '--cache_dir'		,			type=_UpperCAmelCase		,			default=_UpperCAmelCase		,			required=_UpperCAmelCase		,			help='Directory containing pytorch model'				)
  parser.add_argument('--pytorch_model_path'		,			type=_UpperCAmelCase		,			required=_UpperCAmelCase		,			help='/path/to/<pytorch-model-name>.bin'				)
  parser.add_argument('--tf_cache_dir'		,			type=_UpperCAmelCase		,			required=_UpperCAmelCase		,			help='Directory in which to save tensorflow model'				)
  lowerCamelCase__					:   List[str]      				=	parser.parse_args(_UpperCAmelCase				)

  lowerCamelCase__					:   Optional[Any]      				=	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=_UpperCAmelCase		,			ckpt_dir=args.tf_cache_dir		,			model_name=args.model_name				)


if __name__ == "__main__":
    main()












import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


_UpperCAmelCase       : Union[str, Any]			     =       logging.get_logger(__name__)

class      lowerCAmelCase							:
 UpperCAmelCase__									=	None


@experimental
def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase				)    ->  Any:
    if ParallelBackendConfig.backend_name is None:
        return _map_with_multiprocessing_pool(
            _UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase				)

    return _map_with_joblib(_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase				)


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase				)    ->  Any:
    lowerCamelCase__					:   int      				=	num_proc if num_proc <= len(_UpperCAmelCase				) else len(_UpperCAmelCase				)
    lowerCamelCase__					:   Optional[int]      				=	[]  # We organize the splits ourselve (contiguous splits)
    for index in range(_UpperCAmelCase				):
        lowerCamelCase__					:   str      				=	len(_UpperCAmelCase				) // num_proc
        lowerCamelCase__					:   List[str]      				=	len(_UpperCAmelCase				) % num_proc
        lowerCamelCase__					:   Tuple      				=	div * index + min(_UpperCAmelCase		,			_UpperCAmelCase				)
        lowerCamelCase__					:   int      				=	start + div + (1 if index < mod else 0)
        split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)				)

    if len(_UpperCAmelCase				) != sum(len(i[1]				) for i in split_kwds				):
        raise ValueError(
            F"""Error dividing inputs iterable among processes. """
            F"""Total number of objects {len(_UpperCAmelCase				)}, """
            F"""length: {sum(len(i[1]				) for i in split_kwds				)}"""				)

    logger.info(
        F"""Spawning {num_proc} processes for {len(_UpperCAmelCase				)} objects in slices of {[len(i[1]				) for i in split_kwds]}"""				)
    lowerCamelCase__				,							lowerCamelCase__					:   List[str]      				=	None, None
    if not disable_tqdm:
        lowerCamelCase__				,							lowerCamelCase__					:   Union[str, Any]      				=	(RLock(),), tqdm.set_lock
    with Pool(_UpperCAmelCase		,			initargs=_UpperCAmelCase		,			initializer=_UpperCAmelCase				) as pool:
        lowerCamelCase__					:   Any      				=	pool.map(_UpperCAmelCase		,			_UpperCAmelCase				)
    logger.info(F"""Finished {num_proc} processes"""				)
    lowerCamelCase__					:   Union[str, Any]      				=	[obj for proc_res in mapped for obj in proc_res]
    logger.info(F"""Unpacked {len(_UpperCAmelCase				)} objects"""				)

    return mapped


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase				)    ->  int:
    # progress bar is not yet supported for _map_with_joblib, because tqdm couldn't accurately be applied to joblib,
    # and it requires monkey-patching joblib internal classes which is subject to change
    import joblib

    with joblib.parallel_backend(ParallelBackendConfig.backend_name		,			n_jobs=_UpperCAmelCase				):
        return joblib.Parallel()(
            joblib.delayed(_UpperCAmelCase				)((function, obj, types, None, True, None)				) for obj in iterable				)


@experimental
@contextlib.contextmanager
def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase				)    ->  Tuple:
    lowerCamelCase__					:   List[str]      				=	backend_name

    if backend_name == "spark":
        from joblibspark import register_spark

        register_spark()

        # TODO: call create_cache_and_write_probe if "download" in steps
        # TODO: raise NotImplementedError when Dataset.map etc is called

    try:
        yield
    finally:
        lowerCamelCase__					:   str      				=	None





