






"""simple docstring"""




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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to combine both the gradient accumulation
# and automatic batch size finder utilities of Accelerate to perfrom
# automatic gradient accumulation
#
# 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
#
########################################################################

lowerCamelCase__			=   32






def 						_SCREAMING_SNAKE_CASE							( UpperCamelCase     :	Accelerator		,				UpperCamelCase     :	int = 16    ):
  A__											=							AutoTokenizer.from_pretrained("""bert-base-cased"""    )
  A__											=							load_dataset("""glue"""		,				"""mrpc"""    )

  def tokenize_function(UpperCamelCase     :	List[Any]    ):
    # max_length=None => use the model max length (it's actually the default)
    A__											=							tokenizer(examples["""sentence1"""]		,				examples["""sentence2"""]		,				truncation=UpperCamelCase		,				max_length=UpperCamelCase    )
    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():
    A__											=							datasets.map(
        UpperCamelCase		,				batched=UpperCamelCase		,				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
  A__											=							tokenized_datasets.rename_column("""label"""		,				"""labels"""    )

  def collate_fn(UpperCamelCase     :	Optional[Any]    ):
    # When using mixed precision we want round multiples of 8/16
    if accelerator.mixed_precision == "fp8":
      A__											=							16
    elif accelerator.mixed_precision != "no":
      A__											=							8
    else:
      A__											=							None

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

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

  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

      lowerCamelCase__			=   mocked_dataloaders  # noqa: F811






def 						_SCREAMING_SNAKE_CASE							( UpperCamelCase     :	Any		,				UpperCamelCase     :	Optional[int]    ):
  # For testing only
  if os.environ.get("""TESTING_MOCKED_DATALOADERS"""		,				UpperCamelCase    ) == "1":
    A__											=							2
  # Initialize accelerator
  A__											=							Accelerator(cpu=args.cpu		,				mixed_precision=args.mixed_precision    )
  # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
  A__											=							config["""lr"""]
  A__											=							int(config["""num_epochs"""]    )
  A__											=							int(config["""seed"""]    )
  A__											=							int(config["""batch_size"""]    )

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

  # New Code #
  # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
  # to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
  # half each time. From this, we can calculate the number of gradient accumulation steps needed
  # and modify the Accelerator object as a result
  @find_executable_batch_size(starting_batch_size=int(UpperCamelCase    )    )
  def inner_training_loop(UpperCamelCase     :	Any    ):
    # Since we need to modify the outside accelerator object, we need to bring it
    # to the local scope
    nonlocal accelerator

    # We can calculate the number of gradient accumulation steps based on the current
    # batch size vs the starting batch size
    A__											=							observed_batch_size // batch_size

    # And then set it in the Accelerator directly:
    A__											=							num_gradient_accumulation_steps

    # Next we need to free all of the stored model references in the Accelerator each time
    accelerator.free_memory()

    # And set the seed so our results are reproducable each reset
    set_seed(UpperCamelCase    )

    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    A__											=							AutoModelForSequenceClassification.from_pretrained("""bert-base-cased"""		,				return_dict=UpperCamelCase    )

    # 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).
    A__											=							model.to(accelerator.device    )

    # Instantiate optimizer
    A__											=							AdamW(params=model.parameters()		,				lr=UpperCamelCase    )
    A__						,A__											=							get_dataloaders(UpperCamelCase		,				UpperCamelCase    )

    # Instantiate scheduler
    A__											=							get_linear_schedule_with_warmup(
        optimizer=UpperCamelCase		,				num_warmup_steps=100		,				num_training_steps=(len(UpperCamelCase    ) * 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.
    A__						,A__						,A__						,A__						,A__											=							accelerator.prepare(
        UpperCamelCase		,				UpperCamelCase		,				UpperCamelCase		,				UpperCamelCase		,				UpperCamelCase    )

    # Now we train the model
    for epoch in range(UpperCamelCase    ):
      model.train()
      for step, batch in enumerate(UpperCamelCase    ):
        # And perform gradient accumulation
        with accelerator.accumulate(UpperCamelCase    ):
          # We could avoid this line since we set the accelerator with `device_placement=True`.
          batch.to(accelerator.device    )
          A__											=							model(**UpperCamelCase    )
          A__											=							outputs.loss
          accelerator.backward(UpperCamelCase    )
          optimizer.step()
          lr_scheduler.step()
          optimizer.zero_grad()

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

      A__											=							metric.compute()
      # Use accelerator.print to print only on the main process.
      accelerator.print(F"""epoch {epoch}:"""		,				UpperCamelCase    )

    # 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 						_SCREAMING_SNAKE_CASE							( ):
  A__											=							argparse.ArgumentParser(description="""Simple example of training script."""    )
  parser.add_argument(
      """--mixed_precision"""		,				type=UpperCamelCase		,				default=UpperCamelCase		,				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."""    )
  A__											=							parser.parse_args()
  # New Code #
  # We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
  A__											=							{"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 256}
  training_function(UpperCamelCase		,				UpperCamelCase    )


if __name__ == "__main__":
      main()












"""simple docstring"""




import argparse

from . import (
    BertConfig,
    BertGenerationConfig,
    BertGenerationDecoder,
    BertGenerationEncoder,
    load_tf_weights_in_bert_generation,
    logging,
)


logging.set_verbosity_info()






def 						_SCREAMING_SNAKE_CASE							( UpperCamelCase     :	List[str]		,				UpperCamelCase     :	Tuple		,				UpperCamelCase     :	Optional[int]		,				UpperCamelCase     :	Union[str, Any]		,				UpperCamelCase     :	Optional[int]    ):
  # Initialise PyTorch model
  A__											=							BertConfig.from_pretrained(
      """bert-large-cased"""		,				vocab_size=UpperCamelCase		,				max_position_embeddings=512		,				is_decoder=UpperCamelCase		,				add_cross_attention=UpperCamelCase		,				)
  A__											=							bert_config.to_dict()
  del bert_config_dict["type_vocab_size"]
  A__											=							BertGenerationConfig(**UpperCamelCase    )
  if is_encoder:
    A__											=							BertGenerationEncoder(UpperCamelCase    )
  else:
    A__											=							BertGenerationDecoder(UpperCamelCase    )
  print(F"""Building PyTorch model from configuration: {config}"""    )

  # Load weights from tf checkpoint
  load_tf_weights_in_bert_generation(
      UpperCamelCase		,				UpperCamelCase		,				model_class="""bert"""		,				is_encoder_named_decoder=UpperCamelCase		,				is_encoder=UpperCamelCase		,				)

  # Save pytorch-model
  print(F"""Save PyTorch model and config to {pytorch_dump_path}"""    )
  model.save_pretrained(UpperCamelCase    )


if __name__ == "__main__":
      lowerCamelCase__			=   argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--tf_hub_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
      )
      parser.add_argument(
          "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
      )
      parser.add_argument(
          "--is_encoder_named_decoder",
          action="store_true",
          help="If decoder has to be renamed to encoder in PyTorch model.",
      )
      parser.add_argument("--is_encoder", action="store_true", help="If model is an encoder.")
      parser.add_argument("--vocab_size", default=50_358, type=int, help="Vocab size of model")
      lowerCamelCase__			=   parser.parse_args()
      convert_tf_checkpoint_to_pytorch(
          args.tf_hub_path,
          args.pytorch_dump_path,
          args.is_encoder_named_decoder,
          args.vocab_size,
          is_encoder=args.is_encoder,
      )




