


"""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()
a__       :							int        =  logging.get_logger(__name__)




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




     '''simple docstring'''


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

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




     '''simple docstring'''


     logger.info(f"""Loading model based on config from {config_path}..."""     )
     __SCREAMING_SNAKE_CASE		     =       BertConfig.from_json_file(lowerCAmelCase_     )
     __SCREAMING_SNAKE_CASE		     =       BertModel(lowerCAmelCase_     )

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

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


if __name__ == "__main__":
  a__       :							Union[str, Any]        =  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).''',
  )
  a__       :							Tuple        =  parser.parse_args()
  convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)




"""simple docstring"""







import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
		from sagemaker import Session, TrainingJobAnalytics
		from sagemaker.huggingface import HuggingFace




@pytest.mark.skipif(
    literal_eval(os.getenv("TEST_SAGEMAKER"    ,						"False")) is not True    ,						reason="Skipping test because should only be run when releasing minor transformers version"    ,						)
@pytest.mark.usefixtures("sm_env")
@parameterized_class(
    [
        {
            "framework": "pytorch",
            "script": "run_glue.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 650, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "pytorch",
            "script": "run_ddp.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "tensorflow",
            "script": "run_tf_dist.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.6, "eval_loss": 0.7},
        },
    ])
class 				UpperCamelCase_			(				unittest.TestCase):

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Any:
						if self.framework == "pytorch":
											subprocess.run(
											    F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()    ,    encoding="utf-8"    ,    check=UpperCAmelCase__    ,    )
						assert hasattr(self    ,    "env"     )




	def 						UpperCAmelCase_   (				self    :       Tuple    ,    UpperCAmelCase__    :       int     ) ->			Any:
						__SCREAMING_SNAKE_CASE		     =       F"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
						# distributed data settings
						__SCREAMING_SNAKE_CASE		     =       {"smdistributed": {"dataparallel": {"enabled": True}}} if self.script != "run_ddp.py" else None

						# creates estimator
						return HuggingFace(
						    entry_point=self.script    ,    source_dir=self.env.test_path    ,    role=self.env.role    ,    image_uri=self.env.image_uri    ,    base_job_name=UpperCAmelCase__    ,    instance_count=UpperCAmelCase__    ,    instance_type=self.instance_type    ,    debugger_hook_config=UpperCAmelCase__    ,    hyperparameters={**self.env.distributed_hyperparameters, "model_name_or_path": self.model_name_or_path}    ,    metric_definitions=self.env.metric_definitions    ,    distribution=UpperCAmelCase__    ,    py_version="py36"    ,    )




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Optional[int]     ) ->			Optional[Any]:
						TrainingJobAnalytics(UpperCAmelCase__     ).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""     )




	@parameterized.expand([(2,)]     )
	def 						UpperCAmelCase_   (				self    :       Dict    ,    UpperCAmelCase__    :       Optional[Any]     ) ->			List[str]:
						# create estimator
						__SCREAMING_SNAKE_CASE		     =       self.create_estimator(UpperCAmelCase__     )

						# run training
						estimator.fit()

						# result dataframe
						__SCREAMING_SNAKE_CASE		     =       TrainingJobAnalytics(estimator.latest_training_job.name     ).dataframe()

						# extract kpis
						__SCREAMING_SNAKE_CASE		     =       list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]     )
						__SCREAMING_SNAKE_CASE		     =       list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]     )
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						__SCREAMING_SNAKE_CASE		     =       (
						    Session().describe_training_job(estimator.latest_training_job.name     ).get("TrainingTimeInSeconds"    ,    9_9_9_9_9_9     )
						)

						# assert kpis
						assert train_runtime <= self.results["train_runtime"]
						assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy     )
						assert all(t <= self.results["eval_loss"] for t in eval_loss     )

						# dump tests result into json file to share in PR
						with open(F"""{estimator.latest_training_job.name}.json"""    ,    "w"     ) as outfile:
											json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}    ,    UpperCAmelCase__     )
