





'''simple docstring'''






import argparse
from typing import List

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset

# New Code #
# We'll be using StratifiedKFold for this example
from sklearn.model_selection import StratifiedKFold
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to perform Cross Validation,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


A__    :       List[Any]              =16
A__    :       Union[str, Any]              =32
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase = 16  ):





      """simple docstring"""


      _lowerCAmelCase												=							AutoTokenizer.from_pretrained("""bert-base-cased"""  )
      _lowerCAmelCase												=							DatasetDict(
          {
              """train""": dataset["""train"""].select(lowerCAmelCase  ),
              """validation""": dataset["""train"""].select(lowerCAmelCase  ),
              """test""": dataset["""validation"""],
          }  )

      def tokenize_function(lowerCAmelCase  ):
            # max_length=None => use the model max length (it's actually the default)
            _lowerCAmelCase												=							tokenizer(examples["""sentence1"""]							,     examples["""sentence2"""]							,     truncation=lowerCAmelCase							,     max_length=lowerCAmelCase  )
            return outputs

      # Apply the method we just defined to all the examples in all the splits of the dataset
      # starting with the main process first:
      with accelerator.main_process_first():
            _lowerCAmelCase												=							datasets.map(
                lowerCAmelCase							,     batched=lowerCAmelCase							,     remove_columns=["""idx""", """sentence1""", """sentence2"""]							,     )

      # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
      # transformers library
      _lowerCAmelCase												=							tokenized_datasets.rename_column("""label"""							,     """labels"""  )

      def collate_fn(lowerCAmelCase  ):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            _lowerCAmelCase												=							1_28 if accelerator.distributed_type == DistributedType.TPU else None
            # When using mixed precision we want round multiples of 8/16
            if accelerator.mixed_precision == "fp8":
                  _lowerCAmelCase												=							16
            elif accelerator.mixed_precision != "no":
                  _lowerCAmelCase												=							8
            else:
                  _lowerCAmelCase												=							None

            return tokenizer.pad(
                lowerCAmelCase							,     padding="""longest"""							,     max_length=lowerCAmelCase							,     pad_to_multiple_of=lowerCAmelCase							,     return_tensors="""pt"""							,     )

      # Instantiate dataloaders.
      _lowerCAmelCase												=							DataLoader(
          tokenized_datasets["""train"""]							,     shuffle=lowerCAmelCase							,     collate_fn=lowerCAmelCase							,     batch_size=lowerCAmelCase  )
      _lowerCAmelCase												=							DataLoader(
          tokenized_datasets["""validation"""]							,     shuffle=lowerCAmelCase							,     collate_fn=lowerCAmelCase							,     batch_size=lowerCAmelCase  )

      _lowerCAmelCase												=							DataLoader(
          tokenized_datasets["""test"""]							,     shuffle=lowerCAmelCase							,     collate_fn=lowerCAmelCase							,     batch_size=lowerCAmelCase  )

      return train_dataloader, eval_dataloader, test_dataloader
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      _lowerCAmelCase												=							[]
      # Download the dataset
      _lowerCAmelCase												=							load_dataset("""glue"""							,     """mrpc"""  )
      # Create our splits
      _lowerCAmelCase												=							StratifiedKFold(n_splits=int(args.num_folds  )  )
      # Initialize accelerator
      _lowerCAmelCase												=							Accelerator(cpu=args.cpu							,     mixed_precision=args.mixed_precision  )
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      _lowerCAmelCase												=							config["""lr"""]
      _lowerCAmelCase												=							int(config["""num_epochs"""]  )
      _lowerCAmelCase												=							int(config["""seed"""]  )
      _lowerCAmelCase												=							int(config["""batch_size"""]  )

      _lowerCAmelCase												=							evaluate.load("""glue"""							,     """mrpc"""  )

      # If the batch size is too big we use gradient accumulation
      _lowerCAmelCase												=							1
      if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
            _lowerCAmelCase												=							batch_size // MAX_GPU_BATCH_SIZE
            _lowerCAmelCase												=							MAX_GPU_BATCH_SIZE

      set_seed(lowerCAmelCase  )

      # New Code #
      # Create our folds:
      _lowerCAmelCase												=							kfold.split(np.zeros(datasets["""train"""].num_rows  )							,     datasets["""train"""]["""label"""]  )
      _lowerCAmelCase												=							[]
      # Iterate over them
      for i, (train_idxs, valid_idxs) in enumerate(lowerCAmelCase  ):
            _lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase												=							get_fold_dataloaders(
                lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     )
            # Instantiate the model (we build the model here so that the seed also control new weights initialization)
            _lowerCAmelCase												=							AutoModelForSequenceClassification.from_pretrained("""bert-base-cased"""							,     return_dict=lowerCAmelCase  )

            # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
            # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
            # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
            _lowerCAmelCase												=							model.to(accelerator.device  )

            # Instantiate optimizer
            _lowerCAmelCase												=							AdamW(params=model.parameters()							,     lr=lowerCAmelCase  )

            # Instantiate scheduler
            _lowerCAmelCase												=							get_linear_schedule_with_warmup(
                optimizer=lowerCAmelCase							,     num_warmup_steps=1_00							,     num_training_steps=(len(lowerCAmelCase  ) * num_epochs) // gradient_accumulation_steps							,     )

            # Prepare everything
            # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
            # prepare method.
            _lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase												=							accelerator.prepare(
                lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  )

            # Now we train the model
            for epoch in range(lowerCAmelCase  ):
                  model.train()
                  for step, batch in enumerate(lowerCAmelCase  ):
                        # We could avoid this line since we set the accelerator with `device_placement=True`.
                        batch.to(accelerator.device  )
                        _lowerCAmelCase												=							model(**lowerCAmelCase  )
                        _lowerCAmelCase												=							outputs.loss
                        _lowerCAmelCase												=							loss / gradient_accumulation_steps
                        accelerator.backward(lowerCAmelCase  )
                        if step % gradient_accumulation_steps == 0:
                              optimizer.step()
                              lr_scheduler.step()
                              optimizer.zero_grad()

                  model.eval()
                  for step, batch in enumerate(lowerCAmelCase  ):
                        # We could avoid this line since we set the accelerator with `device_placement=True`.
                        batch.to(accelerator.device  )
                        with torch.no_grad():
                              _lowerCAmelCase												=							model(**lowerCAmelCase  )
                        _lowerCAmelCase												=							outputs.logits.argmax(dim=-1  )
                        _lowerCAmelCase					,							_lowerCAmelCase												=							accelerator.gather_for_metrics((predictions, batch["""labels"""])  )
                        metric.add_batch(
                            predictions=lowerCAmelCase							,     references=lowerCAmelCase							,     )

                  _lowerCAmelCase												=							metric.compute()
                  # Use accelerator.print to print only on the main process.
                  accelerator.print(f"epoch {epoch}:"							,     lowerCAmelCase  )

            # New Code #
            # We also run predictions on the test set at the very end
            _lowerCAmelCase												=							[]
            for step, batch in enumerate(lowerCAmelCase  ):
                  # We could avoid this line since we set the accelerator with `device_placement=True`.
                  batch.to(accelerator.device  )
                  with torch.no_grad():
                        _lowerCAmelCase												=							model(**lowerCAmelCase  )
                  _lowerCAmelCase												=							outputs.logits
                  _lowerCAmelCase					,							_lowerCAmelCase												=							accelerator.gather_for_metrics((predictions, batch["""labels"""])  )
                  fold_predictions.append(predictions.cpu()  )
                  if i == 0:
                        # We need all of the test predictions
                        test_references.append(references.cpu()  )
        # Use accelerator.print to print only on the main process.
            test_predictions.append(torch.cat(lowerCAmelCase							,     dim=0  )  )
            # We now need to release all our memory and get rid of the current model, optimizer, etc
            accelerator.free_memory()
      # New Code #
      # Finally we check the accuracy of our folded results:
      _lowerCAmelCase												=							torch.cat(lowerCAmelCase							,     dim=0  )
      _lowerCAmelCase												=							torch.stack(lowerCAmelCase							,     dim=0  ).sum(dim=0  ).div(int(args.num_folds  )  ).argmax(dim=-1  )
      _lowerCAmelCase												=							metric.compute(predictions=lowerCAmelCase							,     references=lowerCAmelCase  )
      accelerator.print("""Average test metrics from all folds:"""							,     lowerCAmelCase  )
def        UpperCamelCase__				(		):





      """simple docstring"""


      _lowerCAmelCase												=							argparse.ArgumentParser(description="""Simple example of training script."""  )
      parser.add_argument(
          """--mixed_precision"""							,     type=lowerCAmelCase							,     default=lowerCAmelCase							,     choices=["""no""", """fp16""", """bf16""", """fp8"""]							,     help="""Whether to use mixed precision. Choose"""
          """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
          """and an Nvidia Ampere GPU."""							,     )
      parser.add_argument("""--cpu"""							,     action="""store_true"""							,     help="""If passed, will train on the CPU."""  )
      # New Code #
      parser.add_argument("""--num_folds"""							,     type=lowerCAmelCase							,     default=3							,     help="""The number of splits to perform across the dataset"""  )
      _lowerCAmelCase												=							parser.parse_args()
      _lowerCAmelCase												=							{"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
      training_function(lowerCAmelCase							,     lowerCAmelCase  )


if __name__ == "__main__":
      main()













'''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)






