





"""simple docstring"""





import argparse
import gc
import json
import os

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.deepspeed import DummyOptim, DummyScheduler


lowercase__				:str				   =       1_6
lowercase__				:Union[str, Any]				   =       3_2


def       lowerCamelCase_       (   UpperCAmelCase_  )						->str:





  """simple docstring"""


  return int(x / 2**20  )




class 							snake_case							:

      '''simple docstring'''



      def __enter__(							self       :      Any ):


        '''simple docstring'''




        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
        __UpperCAmelCase    :						Optional[int]       =						torch.cuda.memory_allocated()
        return self







      def __exit__(							self       :      Optional[Any]						,				*__lowercase       :      List[Any] ):


        '''simple docstring'''




        gc.collect()
        torch.cuda.empty_cache()
        __UpperCAmelCase    :						Tuple       =						torch.cuda.memory_allocated()
        __UpperCAmelCase    :						int       =						torch.cuda.max_memory_allocated()
        __UpperCAmelCase    :						List[Any]       =						bamb(self.end - self.begin )
        __UpperCAmelCase    :						str       =						bamb(self.peak - self.begin )
        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_ = 16  ,			UpperCAmelCase_ = "bert-base-cased"  ,			UpperCAmelCase_ = 3_20  ,			UpperCAmelCase_ = 1_60  ,			)						->Tuple:





  """simple docstring"""


  __UpperCAmelCase    :						str       =						AutoTokenizer.from_pretrained(UpperCAmelCase_  )
  __UpperCAmelCase    :						Tuple       =						load_dataset(
      '''glue'''  ,			'''mrpc'''  ,			split={'''train''': f'''train[:{n_train}]''', '''validation''': f'''validation[:{n_val}]'''}  )

  def tokenize_function(UpperCAmelCase_  ):
    # max_length=None => use the model max length (it's actually the default)
    __UpperCAmelCase    :						Any       =						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
  __UpperCAmelCase    :						List[str]       =						datasets.map(
      UpperCAmelCase_  ,			batched=UpperCAmelCase_  ,			remove_columns=['''idx''', '''sentence1''', '''sentence2''']  ,			load_from_cache_file=UpperCAmelCase_  )

  # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
  # transformers library
  __UpperCAmelCase    :						Union[str, Any]       =						tokenized_datasets.rename_column('''label'''  ,			'''labels'''  )

  def collate_fn(UpperCAmelCase_  ):
    # On TPU it's best to pad everything to the same length or training will be very slow.
    if accelerator.distributed_type == DistributedType.TPU:
      return tokenizer.pad(UpperCAmelCase_  ,			padding='''max_length'''  ,			max_length=1_28  ,			return_tensors='''pt'''  )
    return tokenizer.pad(UpperCAmelCase_  ,			padding='''longest'''  ,			return_tensors='''pt'''  )

  # Instantiate dataloaders.
  __UpperCAmelCase    :						str       =						DataLoader(
      tokenized_datasets['''train''']  ,			shuffle=UpperCAmelCase_  ,			collate_fn=UpperCAmelCase_  ,			batch_size=UpperCAmelCase_  )
  __UpperCAmelCase    :						Union[str, Any]       =						DataLoader(
      tokenized_datasets['''validation''']  ,			shuffle=UpperCAmelCase_  ,			collate_fn=UpperCAmelCase_  ,			batch_size=UpperCAmelCase_  )

  return train_dataloader, eval_dataloader


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  )						->int:





  """simple docstring"""


  __UpperCAmelCase    :						Dict       =						Accelerator()

  # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
  __UpperCAmelCase    :						Any       =						config['''lr''']
  __UpperCAmelCase    :						str       =						int(config['''num_epochs''']  )
  __UpperCAmelCase    :						Tuple       =						int(config['''seed''']  )
  __UpperCAmelCase    :						Optional[Any]       =						int(config['''batch_size''']  )
  __UpperCAmelCase    :						List[Any]       =						args.model_name_or_path

  set_seed(UpperCAmelCase_  )
  __UpperCAmelCase       ,     __UpperCAmelCase    :						Tuple       =						get_dataloaders(UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  ,			args.n_train  ,			args.n_val  )

  # Instantiate the model (we build the model here so that the seed also control new weights initialization)
  __UpperCAmelCase    :						Union[str, Any]       =						AutoModelForSequenceClassification.from_pretrained(UpperCAmelCase_  ,			return_dict=UpperCAmelCase_  )

  # Instantiate optimizer
  __UpperCAmelCase    :						Dict       =						(
      AdamW
      if accelerator.state.deepspeed_plugin is None
      or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
      else DummyOptim
  )
  __UpperCAmelCase    :						Dict       =						optimizer_cls(params=model.parameters()  ,			lr=UpperCAmelCase_  )

  if accelerator.state.deepspeed_plugin is not None:
    __UpperCAmelCase    :						Dict       =						accelerator.state.deepspeed_plugin.deepspeed_config[
        '''gradient_accumulation_steps'''
    ]
  else:
    __UpperCAmelCase    :						List[str]       =						1
  __UpperCAmelCase    :						Optional[Any]       =						(len(UpperCAmelCase_  ) * num_epochs) // gradient_accumulation_steps

  # Instantiate scheduler
  if (
      accelerator.state.deepspeed_plugin is None
      or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
  ):
    __UpperCAmelCase    :						int       =						get_linear_schedule_with_warmup(
        optimizer=UpperCAmelCase_  ,			num_warmup_steps=0  ,			num_training_steps=UpperCAmelCase_  ,			)
  else:
    __UpperCAmelCase    :						List[str]       =						DummyScheduler(UpperCAmelCase_  ,			total_num_steps=UpperCAmelCase_  ,			warmup_num_steps=0  )

  # 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.
  __UpperCAmelCase       ,     __UpperCAmelCase       ,     __UpperCAmelCase       ,     __UpperCAmelCase       ,     __UpperCAmelCase    :						Dict       =						accelerator.prepare(
      UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  )

  # We need to keep track of how many total steps we have iterated over
  __UpperCAmelCase    :						Tuple       =						0
  # We also need to keep track of the stating epoch so files are named properly
  __UpperCAmelCase    :						Tuple       =						0

  # Now we train the model
  __UpperCAmelCase    :						Any       =						{}
  for epoch in range(UpperCAmelCase_  ,			UpperCAmelCase_  ):
    with TorchTracemalloc() as tracemalloc:
      model.train()
      for step, batch in enumerate(UpperCAmelCase_  ):
        __UpperCAmelCase    :						Optional[int]       =						model(**UpperCAmelCase_  )
        __UpperCAmelCase    :						Tuple       =						outputs.loss
        __UpperCAmelCase    :						Tuple       =						loss / gradient_accumulation_steps
        accelerator.backward(UpperCAmelCase_  )
        if step % gradient_accumulation_steps == 0:
          optimizer.step()
          lr_scheduler.step()
          optimizer.zero_grad()

        overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
    accelerator.print('''Memory before entering the train : {}'''.format(bamb(tracemalloc.begin  )  )  )
    accelerator.print('''Memory consumed at the end of the train (end-begin): {}'''.format(tracemalloc.used  )  )
    accelerator.print('''Peak Memory consumed during the train (max-begin): {}'''.format(tracemalloc.peaked  )  )
    accelerator.print(
        '''Total Peak Memory consumed during the train (max): {}'''.format(
            tracemalloc.peaked + bamb(tracemalloc.begin  )  )  )
    __UpperCAmelCase    :						Dict       =						tracemalloc.peaked + bamb(tracemalloc.begin  )
    if args.peak_memory_upper_bound is not None:
      assert (
          train_total_peak_memory[f'''epoch-{epoch}'''] <= args.peak_memory_upper_bound
      ), "Peak memory usage exceeded the upper bound"

  accelerator.wait_for_everyone()
  if accelerator.is_main_process:
    with open(os.path.join(args.output_dir  ,			'''peak_memory_utilization.json'''  )  ,			'''w'''  ) as f:
      json.dump(UpperCAmelCase_  ,			UpperCAmelCase_  )


def       lowerCamelCase_       (   )						->int:





  """simple docstring"""


  __UpperCAmelCase    :						List[Any]       =						argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.'''  )
  parser.add_argument(
      '''--model_name_or_path'''  ,			type=UpperCAmelCase_  ,			default='''bert-base-cased'''  ,			help='''Path to pretrained model or model identifier from huggingface.co/models.'''  ,			required=UpperCAmelCase_  ,			)
  parser.add_argument(
      '''--output_dir'''  ,			type=UpperCAmelCase_  ,			default='''.'''  ,			help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'''  ,			)
  parser.add_argument(
      '''--peak_memory_upper_bound'''  ,			type=UpperCAmelCase_  ,			default=UpperCAmelCase_  ,			help='''The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.'''  ,			)
  parser.add_argument(
      '''--n_train'''  ,			type=UpperCAmelCase_  ,			default=3_20  ,			help='''Number of training examples to use.'''  ,			)
  parser.add_argument(
      '''--n_val'''  ,			type=UpperCAmelCase_  ,			default=1_60  ,			help='''Number of validation examples to use.'''  ,			)
  parser.add_argument(
      '''--num_epochs'''  ,			type=UpperCAmelCase_  ,			default=1  ,			help='''Number of train epochs.'''  ,			)
  __UpperCAmelCase    :						Optional[Any]       =						parser.parse_args()
  __UpperCAmelCase    :						Optional[Any]       =						{'''lr''': 2e-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}
  training_function(UpperCAmelCase_  ,			UpperCAmelCase_  )


if __name__ == "__main__":
  main()






"""simple docstring"""





# coding=utf-8
# Copyright 2020 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# this script dumps information about the environment

import os
import sys

import transformers


lowercase__				:int				   =       '3'

print('Python version:', sys.version)
print('transformers version:', transformers.__version__)

try:
  import torch

  print('Torch version:', torch.__version__)
  print('Cuda available:', torch.cuda.is_available())
  print('Cuda version:', torch.version.cuda)
  print('CuDNN version:', torch.backends.cudnn.version())
  print('Number of GPUs available:', torch.cuda.device_count())
  print('NCCL version:', torch.cuda.nccl.version())
except ImportError:
  print('Torch version:', None)

try:
  import deepspeed

  print('DeepSpeed version:', deepspeed.__version__)
except ImportError:
  print('DeepSpeed version:', None)

try:
  import tensorflow as tf

  print('TensorFlow version:', tf.__version__)
  print('TF GPUs available:', bool(tf.config.list_physical_devices('GPU')))
  print('Number of TF GPUs available:', len(tf.config.list_physical_devices('GPU')))
except ImportError:
  print('TensorFlow version:', None)