





'''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_model_parallelism.py''',
            '''model_name_or_path''': '''roberta-large''',
            '''instance_type''': '''ml.p3dn.24xlarge''',
            '''results''': {'''train_runtime''': 1600, '''eval_accuracy''': 0.3, '''eval_loss''': 1.2},
        },
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_glue.py''',
            '''model_name_or_path''': '''roberta-large''',
            '''instance_type''': '''ml.p3dn.24xlarge''',
            '''results''': {'''train_runtime''': 1600, '''eval_accuracy''': 0.3, '''eval_loss''': 1.2},
        },
    ]       )
class 			UpperCAmelCase					(      unittest.TestCase       ):




      def     lowercase__   (      self      :			Dict  )    ->				Optional[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=__snake_case      ,      )
            assert hasattr(self      ,      """env"""  )




      def     lowercase__   (      self      :			Optional[int]      ,      __snake_case      :			Tuple  )    ->				Any:
            # configuration for running training on smdistributed Model Parallel
            _lowerCAmelCase												=							{
                """enabled""": True,
                """processes_per_host""": 8,
            }
            _lowerCAmelCase												=							{
                """enabled""": True,
                """parameters""": {
                    """microbatches""": 4,
                    """placement_strategy""": """spread""",
                    """pipeline""": """interleaved""",
                    """optimize""": """speed""",
                    """partitions""": 4,
                    """ddp""": True,
                },
            }

            _lowerCAmelCase												=							{"""smdistributed""": {"""modelparallel""": smp_options}, """mpi""": mpi_options}

            _lowerCAmelCase												=							"""trainer""" if self.script == """run_glue.py""" else """smtrainer"""
            # 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=f"{self.env.base_job_name}-{instance_count}-smp-{name_extension}"      ,      instance_count=__snake_case      ,      instance_type=self.instance_type      ,      debugger_hook_config=__snake_case      ,      hyperparameters={
                    **self.env.hyperparameters,
                    """model_name_or_path""": self.model_name_or_path,
                    """max_steps""": 5_00,
                }      ,      metric_definitions=self.env.metric_definitions      ,      distribution=__snake_case      ,      py_version="""py36"""      ,      )




      def     lowercase__   (      self      :			int      ,      __snake_case      :			Optional[int]  )    ->				Union[str, Any]:
            TrainingJobAnalytics(__snake_case  ).export_csv(f"{self.env.test_path}/{job_name}_metrics.csv"  )




      @parameterized.expand([(1,)]  )
      def     lowercase__   (      self      :			str      ,      __snake_case      :			str  )    ->				List[Any]:
            # create estimator
            _lowerCAmelCase												=							self.create_estimator(__snake_case  )

            # run training
            estimator.fit()

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

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

            # 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}      ,      __snake_case  )













'''simple docstring'''






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        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      def get_masked_lm_array(lowerCAmelCase  ):
            _lowerCAmelCase												=							f"masked_lm/{name}/.ATTRIBUTES/VARIABLE_VALUE"
            _lowerCAmelCase												=							tf.train.load_variable(lowerCAmelCase							,     lowerCAmelCase  )

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

            return torch.from_numpy(lowerCAmelCase  )

      def get_encoder_array(lowerCAmelCase  ):
            _lowerCAmelCase												=							f"encoder/{name}/.ATTRIBUTES/VARIABLE_VALUE"
            _lowerCAmelCase												=							tf.train.load_variable(lowerCAmelCase							,     lowerCAmelCase  )

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

            return torch.from_numpy(lowerCAmelCase  )

      def get_encoder_layer_array(lowerCAmelCase							,     lowerCAmelCase  ):
            _lowerCAmelCase												=							f"encoder/_transformer_layers/{layer_index}/{name}/.ATTRIBUTES/VARIABLE_VALUE"
            _lowerCAmelCase												=							tf.train.load_variable(lowerCAmelCase							,     lowerCAmelCase  )

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

            return torch.from_numpy(lowerCAmelCase  )

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

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

            return torch.from_numpy(lowerCAmelCase  )

      print(f"Loading model based on config from {config_path}..."  )
      _lowerCAmelCase												=							BertConfig.from_json_file(lowerCAmelCase  )
      _lowerCAmelCase												=							BertForMaskedLM(lowerCAmelCase  )

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

            # Self-attention
            _lowerCAmelCase												=							layer.attention.self

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

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

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

            _lowerCAmelCase												=							get_encoder_layer_array(lowerCAmelCase							,     """_attention_layer_norm/gamma"""  )
            _lowerCAmelCase												=							get_encoder_layer_array(lowerCAmelCase							,     """_attention_layer_norm/beta"""  )

            # Intermediate
            _lowerCAmelCase												=							layer.intermediate

            _lowerCAmelCase												=							get_encoder_layer_array(lowerCAmelCase							,     """_intermediate_dense/kernel"""  )
            _lowerCAmelCase												=							get_encoder_layer_array(lowerCAmelCase							,     """_intermediate_dense/bias"""  )

            # Output
            _lowerCAmelCase												=							layer.output

            _lowerCAmelCase												=							get_encoder_layer_array(lowerCAmelCase							,     """_output_dense/kernel"""  )
            _lowerCAmelCase												=							get_encoder_layer_array(lowerCAmelCase							,     """_output_dense/bias"""  )

            _lowerCAmelCase												=							get_encoder_layer_array(lowerCAmelCase							,     """_output_layer_norm/gamma"""  )
            _lowerCAmelCase												=							get_encoder_layer_array(lowerCAmelCase							,     """_output_layer_norm/beta"""  )

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

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

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

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

      _lowerCAmelCase												=							get_masked_lm_array("""embedding_table"""  )

      # Pooling
      _lowerCAmelCase												=							BertPooler(config=lowerCAmelCase  )
      _lowerCAmelCase												=							get_encoder_array("""_pooler_layer/kernel"""  )
      _lowerCAmelCase												=							get_encoder_array("""_pooler_layer/bias"""  )

      # Export final model
      model.save_pretrained(lowerCAmelCase  )

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

      print("""Model conversion was done sucessfully!"""  )


if __name__ == "__main__":
      A__    :       str              =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.''',
      )
      A__    :       Dict              =parser.parse_args()
      convert_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)






