




import argparse
import os

# New Code #
import evaluate
import torch
from datasets import load_dataset
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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, 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)
#
# 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
#
########################################################################


__magic_name__  :  Tuple					=       16
__magic_name__  :  str					=       32
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase = 16				):
  lowerCAmelCase__										=      AutoTokenizer.from_pretrained('''bert-base-cased'''				)
  lowerCAmelCase__										=      load_dataset('''glue'''   , '''mrpc'''				)

  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				)

  return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("""TESTING_MOCKED_DATALOADERS""", None) == "1":
      from accelerate.test_utils.training import mocked_dataloaders

      __magic_name__  :  List[Any]					=       mocked_dataloaders  # noqa: F811
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase				):
  # For testing only
  if os.environ.get('''TESTING_MOCKED_DATALOADERS'''   , __lowerCAmelCase				) == "1":
    lowerCAmelCase__										=      2
  # 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'''				)

  # New Code #
  # We now can define an inner training loop function. It should take a batch size as the only parameter,
  # and build the dataloaders in there.
  # It also gets our decorator
  @find_executable_batch_size(starting_batch_size=__lowerCAmelCase				)
  def inner_training_loop(__lowerCAmelCase				):
    # And now just move everything below under this function
    # We need to bring in the Accelerator object from earlier
    nonlocal accelerator
    # And reset all of its attributes that could hold onto any memory:
    accelerator.free_memory()

    # Then we can declare the model, optimizer, and everything else:
    set_seed(__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				)
    lowerCAmelCase__				,      lowerCAmelCase__										=      get_dataloaders(__lowerCAmelCase   , __lowerCAmelCase				)

    # Instantiate scheduler
    lowerCAmelCase__										=      get_linear_schedule_with_warmup(
        optimizer=__lowerCAmelCase   , num_warmup_steps=1_00   , num_training_steps=(len(__lowerCAmelCase				) * num_epochs)   , )

    # 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
        accelerator.backward(__lowerCAmelCase				)
        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 #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
  inner_training_loop()
def 		a_							( ):
  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.'''				)
  lowerCAmelCase__										=      parser.parse_args()
  lowerCAmelCase__										=      {'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
  training_function(__lowerCAmelCase   , __lowerCAmelCase				)


if __name__ == "__main__":
      main()







import argparse
import os

import gluonnlp as nlp
import mxnet as mx
import numpy as np
import torch
from gluonnlp.base import get_home_dir
from gluonnlp.model.bert import BERTEncoder
from gluonnlp.model.utils import _load_vocab
from gluonnlp.vocab import Vocab
from packaging import version
from torch import nn

from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(nlp.__version__) != version.parse("""0.8.3"""):
      raise Exception("""requires gluonnlp == 0.8.3""")

if version.parse(mx.__version__) != version.parse("""1.5.0"""):
      raise Exception("""requires mxnet == 1.5.0""")

logging.set_verbosity_info()
__magic_name__  :  Tuple					=       logging.get_logger(__name__)

__magic_name__  :  Optional[Any]					=       """The Nymphenburg Palace is a beautiful palace in Munich!"""
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase				):
  lowerCAmelCase__										=      {
      '''attention_cell''': '''multi_head''',
      '''num_layers''': 4,
      '''units''': 10_24,
      '''hidden_size''': 7_68,
      '''max_length''': 5_12,
      '''num_heads''': 8,
      '''scaled''': True,
      '''dropout''': 0.1,
      '''use_residual''': True,
      '''embed_size''': 10_24,
      '''embed_dropout''': 0.1,
      '''word_embed''': None,
      '''layer_norm_eps''': 1E-5,
      '''token_type_vocab_size''': 2,
  }

  lowerCAmelCase__										=      bort_4_8_768_1024_hparams

  # Let's construct the original Bort model here
  # Taken from official BERT implementation, see:
  # https://github.com/alexa/bort/blob/master/bort/bort.py
  lowerCAmelCase__										=      BERTEncoder(
      attention_cell=predefined_args['''attention_cell''']   , num_layers=predefined_args['''num_layers''']   , units=predefined_args['''units''']   , hidden_size=predefined_args['''hidden_size''']   , max_length=predefined_args['''max_length''']   , num_heads=predefined_args['''num_heads''']   , scaled=predefined_args['''scaled''']   , dropout=predefined_args['''dropout''']   , output_attention=__lowerCAmelCase   , output_all_encodings=__lowerCAmelCase   , use_residual=predefined_args['''use_residual''']   , activation=predefined_args.get('''activation'''   , '''gelu'''				)   , layer_norm_eps=predefined_args.get('''layer_norm_eps'''   , __lowerCAmelCase				)   , )

  # Vocab information needs to be fetched first
  # It's the same as RoBERTa, so RobertaTokenizer can be used later
  lowerCAmelCase__										=      '''openwebtext_ccnews_stories_books_cased'''

  # Specify download folder to Gluonnlp's vocab
  lowerCAmelCase__										=      os.path.join(get_home_dir()   , '''models'''				)
  lowerCAmelCase__										=      _load_vocab(__lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase   , cls=__lowerCAmelCase				)

  lowerCAmelCase__										=      nlp.model.BERTModel(
      __lowerCAmelCase   , len(__lowerCAmelCase				)   , units=predefined_args['''units''']   , embed_size=predefined_args['''embed_size''']   , embed_dropout=predefined_args['''embed_dropout''']   , word_embed=predefined_args['''word_embed''']   , use_pooler=__lowerCAmelCase   , use_token_type_embed=__lowerCAmelCase   , token_type_vocab_size=predefined_args['''token_type_vocab_size''']   , use_classifier=__lowerCAmelCase   , use_decoder=__lowerCAmelCase   , )

  original_bort.load_parameters(__lowerCAmelCase   , cast_dtype=__lowerCAmelCase   , ignore_extra=__lowerCAmelCase				)
  lowerCAmelCase__										=      original_bort._collect_params_with_prefix()

  # Build our config 🤗
  lowerCAmelCase__										=      {
      '''architectures''': ['''BertForMaskedLM'''],
      '''attention_probs_dropout_prob''': predefined_args['''dropout'''],
      '''hidden_act''': '''gelu''',
      '''hidden_dropout_prob''': predefined_args['''dropout'''],
      '''hidden_size''': predefined_args['''embed_size'''],
      '''initializer_range''': 0.02,
      '''intermediate_size''': predefined_args['''hidden_size'''],
      '''layer_norm_eps''': predefined_args['''layer_norm_eps'''],
      '''max_position_embeddings''': predefined_args['''max_length'''],
      '''model_type''': '''bort''',
      '''num_attention_heads''': predefined_args['''num_heads'''],
      '''num_hidden_layers''': predefined_args['''num_layers'''],
      '''pad_token_id''': 1,  # 2 = BERT, 1 = RoBERTa
      '''type_vocab_size''': 1,  # 2 = BERT, 1 = RoBERTa
      '''vocab_size''': len(__lowerCAmelCase				),
  }

  lowerCAmelCase__										=      BertConfig.from_dict(__lowerCAmelCase				)
  lowerCAmelCase__										=      BertForMaskedLM(__lowerCAmelCase				)
  hf_bort_model.eval()

  # Parameter mapping table (Gluonnlp to Transformers)
  # * denotes layer index
  #
  # | Gluon Parameter                                                | Transformers Parameter
  # | -------------------------------------------------------------- | ----------------------
  # | `encoder.layer_norm.beta`                                      | `bert.embeddings.LayerNorm.bias`
  # | `encoder.layer_norm.gamma`                                     | `bert.embeddings.LayerNorm.weight`
  # | `encoder.position_weight`                                      | `bert.embeddings.position_embeddings.weight`
  # | `word_embed.0.weight`                                          | `bert.embeddings.word_embeddings.weight`
  # | `encoder.transformer_cells.*.attention_cell.proj_key.bias`     | `bert.encoder.layer.*.attention.self.key.bias`
  # | `encoder.transformer_cells.*.attention_cell.proj_key.weight`   | `bert.encoder.layer.*.attention.self.key.weight`
  # | `encoder.transformer_cells.*.attention_cell.proj_query.bias`   | `bert.encoder.layer.*.attention.self.query.bias`
  # | `encoder.transformer_cells.*.attention_cell.proj_query.weight` | `bert.encoder.layer.*.attention.self.query.weight`
  # | `encoder.transformer_cells.*.attention_cell.proj_value.bias`   | `bert.encoder.layer.*.attention.self.value.bias`
  # | `encoder.transformer_cells.*.attention_cell.proj_value.weight` | `bert.encoder.layer.*.attention.self.value.weight`
  # | `encoder.transformer_cells.*.ffn.ffn_2.bias`                   | `bert.encoder.layer.*.attention.output.dense.bias`
  # | `encoder.transformer_cells.*.ffn.ffn_2.weight`                 | `bert.encoder.layer.*.attention.output.dense.weight`
  # | `encoder.transformer_cells.*.layer_norm.beta`                  | `bert.encoder.layer.*.attention.output.LayerNorm.bias`
  # | `encoder.transformer_cells.*.layer_norm.gamma`                 | `bert.encoder.layer.*.attention.output.LayerNorm.weight`
  # | `encoder.transformer_cells.*.ffn.ffn_1.bias`                   | `bert.encoder.layer.*.intermediate.dense.bias`
  # | `encoder.transformer_cells.*.ffn.ffn_1.weight`                 | `bert.encoder.layer.*.intermediate.dense.weight`
  # | `encoder.transformer_cells.*.ffn.layer_norm.beta`              | `bert.encoder.layer.*.output.LayerNorm.bias`
  # | `encoder.transformer_cells.*.ffn.layer_norm.gamma`             | `bert.encoder.layer.*.output.LayerNorm.weight`
  # | `encoder.transformer_cells.*.proj.bias`                        | `bert.encoder.layer.*.output.dense.bias`
  # | `encoder.transformer_cells.*.proj.weight`                      | `bert.encoder.layer.*.output.dense.weight`

  # Helper function to convert MXNET Arrays to PyTorch
  def to_torch(__lowerCAmelCase				) -> nn.Parameter:
    return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()				)				)

  # Check param shapes and map new HF param back
  def check_and_map_params(__lowerCAmelCase   , __lowerCAmelCase				):
    lowerCAmelCase__										=      hf_param.shape

    lowerCAmelCase__										=      to_torch(params[gluon_param]				)
    lowerCAmelCase__										=      gluon_param.shape

    assert (
        shape_hf == shape_gluon
    ), F"""The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers"""

    return gluon_param

  lowerCAmelCase__										=      check_and_map_params(
      hf_bort_model.bert.embeddings.word_embeddings.weight   , '''word_embed.0.weight'''				)
  lowerCAmelCase__										=      check_and_map_params(
      hf_bort_model.bert.embeddings.position_embeddings.weight   , '''encoder.position_weight'''				)
  lowerCAmelCase__										=      check_and_map_params(
      hf_bort_model.bert.embeddings.LayerNorm.bias   , '''encoder.layer_norm.beta'''				)
  lowerCAmelCase__										=      check_and_map_params(
      hf_bort_model.bert.embeddings.LayerNorm.weight   , '''encoder.layer_norm.gamma'''				)

  # Inspired by RoBERTa conversion script, we just zero them out (Bort does not use them)
  lowerCAmelCase__										=      torch.zeros_like(
      hf_bort_model.bert.embeddings.token_type_embeddings.weight.data				)

  for i in range(hf_bort_config.num_hidden_layers				):
    lowerCAmelCase__										=      hf_bort_model.bert.encoder.layer[i]

    # self attention
    lowerCAmelCase__										=      layer.attention.self

    lowerCAmelCase__										=      check_and_map_params(
        self_attn.key.bias.data   , F"""encoder.transformer_cells.{i}.attention_cell.proj_key.bias"""				)

    lowerCAmelCase__										=      check_and_map_params(
        self_attn.key.weight.data   , F"""encoder.transformer_cells.{i}.attention_cell.proj_key.weight"""				)
    lowerCAmelCase__										=      check_and_map_params(
        self_attn.query.bias.data   , F"""encoder.transformer_cells.{i}.attention_cell.proj_query.bias"""				)
    lowerCAmelCase__										=      check_and_map_params(
        self_attn.query.weight.data   , F"""encoder.transformer_cells.{i}.attention_cell.proj_query.weight"""				)
    lowerCAmelCase__										=      check_and_map_params(
        self_attn.value.bias.data   , F"""encoder.transformer_cells.{i}.attention_cell.proj_value.bias"""				)
    lowerCAmelCase__										=      check_and_map_params(
        self_attn.value.weight.data   , F"""encoder.transformer_cells.{i}.attention_cell.proj_value.weight"""				)

    # self attention output
    lowerCAmelCase__										=      layer.attention.output

    lowerCAmelCase__										=      check_and_map_params(
        self_output.dense.bias   , F"""encoder.transformer_cells.{i}.proj.bias"""				)
    lowerCAmelCase__										=      check_and_map_params(
        self_output.dense.weight   , F"""encoder.transformer_cells.{i}.proj.weight"""				)
    lowerCAmelCase__										=      check_and_map_params(
        self_output.LayerNorm.bias   , F"""encoder.transformer_cells.{i}.layer_norm.beta"""				)
    lowerCAmelCase__										=      check_and_map_params(
        self_output.LayerNorm.weight   , F"""encoder.transformer_cells.{i}.layer_norm.gamma"""				)

    # intermediate
    lowerCAmelCase__										=      layer.intermediate

    lowerCAmelCase__										=      check_and_map_params(
        intermediate.dense.bias   , F"""encoder.transformer_cells.{i}.ffn.ffn_1.bias"""				)
    lowerCAmelCase__										=      check_and_map_params(
        intermediate.dense.weight   , F"""encoder.transformer_cells.{i}.ffn.ffn_1.weight"""				)

    # output
    lowerCAmelCase__										=      layer.output

    lowerCAmelCase__										=      check_and_map_params(
        bert_output.dense.bias   , F"""encoder.transformer_cells.{i}.ffn.ffn_2.bias"""				)
    lowerCAmelCase__										=      check_and_map_params(
        bert_output.dense.weight   , F"""encoder.transformer_cells.{i}.ffn.ffn_2.weight"""				)
    lowerCAmelCase__										=      check_and_map_params(
        bert_output.LayerNorm.bias   , F"""encoder.transformer_cells.{i}.ffn.layer_norm.beta"""				)
    lowerCAmelCase__										=      check_and_map_params(
        bert_output.LayerNorm.weight   , F"""encoder.transformer_cells.{i}.ffn.layer_norm.gamma"""				)

  # Save space and energy 🎄
  hf_bort_model.half()

  # Compare output of both models
  lowerCAmelCase__										=      RobertaTokenizer.from_pretrained('''roberta-base'''				)

  lowerCAmelCase__										=      tokenizer.encode_plus(__lowerCAmelCase				)['''input_ids''']

  # Get gluon output
  lowerCAmelCase__										=      mx.nd.array([input_ids]				)
  lowerCAmelCase__										=      original_bort(inputs=__lowerCAmelCase   , token_types=[]				)

  # Get Transformer output (save and reload model again)
  hf_bort_model.save_pretrained(__lowerCAmelCase				)
  lowerCAmelCase__										=      BertModel.from_pretrained(__lowerCAmelCase				)
  hf_bort_model.eval()

  lowerCAmelCase__										=      tokenizer.encode_plus(__lowerCAmelCase   , return_tensors='''pt'''				)
  lowerCAmelCase__										=      hf_bort_model(**__lowerCAmelCase				)[0]

  lowerCAmelCase__										=      output_gluon[0].asnumpy()
  lowerCAmelCase__										=      output_hf[0].detach().numpy()

  lowerCAmelCase__										=      np.max(np.abs(hf_layer - gluon_layer				)				).item()
  lowerCAmelCase__										=      np.allclose(__lowerCAmelCase   , __lowerCAmelCase   , atol=1E-3				)

  if success:
    print('''✔️ Both model do output the same tensors'''				)
  else:
    print('''❌ Both model do **NOT** output the same tensors'''				)
    print('''Absolute difference is:'''   , __lowerCAmelCase				)


if __name__ == "__main__":
      __magic_name__  :  Any					=       argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--bort_checkpoint_path""", default=None, type=str, required=True, help="""Path the official Bort params file."""
      )
      parser.add_argument(
          """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
      )
      __magic_name__  :  int					=       parser.parse_args()
      convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)

