






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


lowerCamelCase	:							int				=16
lowerCamelCase	:							Tuple				=32







def 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase			) ->	List[str]:
     return int(x / 2**20			)





class 			__a  :







 def __enter__(		self					:   Optional[int]						):


      '''simple docstring'''






      gc.collect()
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
      UpperCamelCase__			:   Tuple       				=  torch.cuda.memory_allocated()
      return self







 def __exit__(		self					:   str							,   *SCREAMING_SNAKE_CASE					:   Dict						):


      '''simple docstring'''






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







def 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase					,    __lowerCAmelCase = 16					,    __lowerCAmelCase = "bert-base-cased"					,    __lowerCAmelCase = 320					,    __lowerCAmelCase = 160					,    ) ->	Optional[Any]:
     UpperCamelCase__			:   Optional[Any]       				=  AutoTokenizer.from_pretrained(__lowerCAmelCase			)
     UpperCamelCase__			:   List[Any]       				=  load_dataset(
         "glue"					,    "mrpc"					,    split={"train": f'train[:{n_train}]', "validation": f'validation[:{n_val}]'}			)

     def tokenize_function(__lowerCAmelCase			):
          # max_length=None => use the model max length (it's actually the default)
          UpperCamelCase__			:   Tuple       				=  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
     UpperCamelCase__			:   int       				=  datasets.map(
         __lowerCAmelCase					,    batched=__lowerCAmelCase					,    remove_columns=["idx", "sentence1", "sentence2"]					,    load_from_cache_file=__lowerCAmelCase			)

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     UpperCamelCase__			:   Optional[int]       				=  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.
          if accelerator.distributed_type == DistributedType.TPU:
               return tokenizer.pad(__lowerCAmelCase					,    padding="max_length"					,    max_length=128					,    return_tensors="pt"			)
          return tokenizer.pad(__lowerCAmelCase					,    padding="longest"					,    return_tensors="pt"			)

     # Instantiate dataloaders.
     UpperCamelCase__			:   str       				=  DataLoader(
         tokenized_datasets["train"]					,    shuffle=__lowerCAmelCase					,    collate_fn=__lowerCAmelCase					,    batch_size=__lowerCAmelCase			)
     UpperCamelCase__			:   List[str]       				=  DataLoader(
         tokenized_datasets["validation"]					,    shuffle=__lowerCAmelCase					,    collate_fn=__lowerCAmelCase					,    batch_size=__lowerCAmelCase			)

     return train_dataloader, eval_dataloader







def 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase					,    __lowerCAmelCase			) ->	Optional[Any]:
     # Initialize accelerator
     UpperCamelCase__			:   Union[str, Any]       				=  Accelerator()

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

     set_seed(__lowerCAmelCase			)
     UpperCamelCase__				,     UpperCamelCase__			:   Union[str, Any]       				=  get_dataloaders(__lowerCAmelCase					,    __lowerCAmelCase					,    __lowerCAmelCase					,    args.n_train					,    args.n_val			)

     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     UpperCamelCase__			:   Any       				=  AutoModelForSequenceClassification.from_pretrained(__lowerCAmelCase					,    return_dict=__lowerCAmelCase			)

     # Instantiate optimizer
     UpperCamelCase__			:   Optional[int]       				=  (
         AdamW
         if accelerator.state.deepspeed_plugin is None
         or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
         else DummyOptim
     )
     UpperCamelCase__			:   List[Any]       				=  optimizer_cls(params=model.parameters()					,    lr=__lowerCAmelCase			)

     if accelerator.state.deepspeed_plugin is not None:
          UpperCamelCase__			:   Dict       				=  accelerator.state.deepspeed_plugin.deepspeed_config[
              "gradient_accumulation_steps"
          ]
     else:
          UpperCamelCase__			:   Any       				=  1
     UpperCamelCase__			:   Union[str, Any]       				=  (len(__lowerCAmelCase			) * 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__			:   Dict       				=  get_linear_schedule_with_warmup(
              optimizer=__lowerCAmelCase					,    num_warmup_steps=0					,    num_training_steps=__lowerCAmelCase					,    )
     else:
          UpperCamelCase__			:   int       				=  DummyScheduler(__lowerCAmelCase					,    total_num_steps=__lowerCAmelCase					,    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__			:   str       				=  accelerator.prepare(
         __lowerCAmelCase					,    __lowerCAmelCase					,    __lowerCAmelCase					,    __lowerCAmelCase					,    __lowerCAmelCase			)

     # We need to keep track of how many total steps we have iterated over
     UpperCamelCase__			:   str       				=  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__			:   Tuple       				=  {}
     for epoch in range(__lowerCAmelCase					,    __lowerCAmelCase			):
          with TorchTracemalloc() as tracemalloc:
               model.train()
               for step, batch in enumerate(__lowerCAmelCase			):
                    UpperCamelCase__			:   str       				=  model(**__lowerCAmelCase			)
                    UpperCamelCase__			:   str       				=  outputs.loss
                    UpperCamelCase__			:   Optional[Any]       				=  loss / gradient_accumulation_steps
                    accelerator.backward(__lowerCAmelCase			)
                    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__			:   int       				=  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(__lowerCAmelCase					,    __lowerCAmelCase			)







def 		SCREAMING_SNAKE_CASE							(					) ->	int:
     UpperCamelCase__			:   Tuple       				=  argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage."			)
     parser.add_argument(
         "--model_name_or_path"					,    type=__lowerCAmelCase					,    default="bert-base-cased"					,    help="Path to pretrained model or model identifier from huggingface.co/models."					,    required=__lowerCAmelCase					,    )
     parser.add_argument(
         "--output_dir"					,    type=__lowerCAmelCase					,    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=__lowerCAmelCase					,    default=__lowerCAmelCase					,    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=__lowerCAmelCase					,    default=320					,    help="Number of training examples to use."					,    )
     parser.add_argument(
         "--n_val"					,    type=__lowerCAmelCase					,    default=160					,    help="Number of validation examples to use."					,    )
     parser.add_argument(
         "--num_epochs"					,    type=__lowerCAmelCase					,    default=1					,    help="Number of train epochs."					,    )
     UpperCamelCase__			:   Optional[Any]       				=  parser.parse_args()
     UpperCamelCase__			:   List[str]       				=  {"lr": 2E-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
     training_function(__lowerCAmelCase					,    __lowerCAmelCase			)


if __name__ == "__main__":
 main()







import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam





class 			__a  (    datasets.BeamBasedBuilder       ):







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      return datasets.DatasetInfo(
          features=datasets.Features({"content": datasets.Value("string"						)}						)							,   supervised_keys=SCREAMING_SNAKE_CASE							,   )







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,   gen_kwargs={"examples": get_test_dummy_examples()}						)]







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   int						):


      '''simple docstring'''






      import apache_beam as beam

      return pipeline | "Load Examples" >> beam.Create(SCREAMING_SNAKE_CASE						)







class 			__a  (    datasets.BeamBasedBuilder       ):







 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      return datasets.DatasetInfo(
          features=datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"						)}						)}						)							,   supervised_keys=SCREAMING_SNAKE_CASE							,   )







 def 					__lowercase					(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   List[Any]						):


      '''simple docstring'''






      return [
          datasets.SplitGenerator(name=datasets.Split.TRAIN							,   gen_kwargs={"examples": get_test_nested_examples()}						)
      ]







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      import apache_beam as beam

      return pipeline | "Load Examples" >> beam.Create(SCREAMING_SNAKE_CASE						)







def 		SCREAMING_SNAKE_CASE							(					) ->	List[str]:
     return [(i, {"content": content}) for i, content in enumerate(["foo", "bar", "foobar"]			)]







def 		SCREAMING_SNAKE_CASE							(					) ->	Optional[Any]:
     return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["foo", "bar", "foobar"]			)]





class 			__a  (    A__       ):







 @require_beam
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  len(get_test_dummy_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   List[Any]       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train.arrow'						)						)						)
           self.assertDictEqual(builder.info.features							,   datasets.Features({"content": datasets.Value("string"						)}						)						)
           UpperCamelCase__			:   Optional[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           self.assertDictEqual(dset["train"][0]							,   get_test_dummy_examples()[0][1]						)
           self.assertDictEqual(
               dset["train"][expected_num_examples - 1]							,   get_test_dummy_examples()[expected_num_examples - 1][1]						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset







 @require_beam
 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      import apache_beam as beam

      UpperCamelCase__			:   Any       				=  beam.io.parquetio.WriteToParquet

      UpperCamelCase__			:   Tuple       				=  len(get_test_dummy_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   str       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           with patch("apache_beam.io.parquetio.WriteToParquet"						) as write_parquet_mock:
                UpperCamelCase__			:   Optional[Any]       				=  partial(SCREAMING_SNAKE_CASE							,   num_shards=2						)
                builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train-00000-of-00002.arrow'						)						)						)
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train-00000-of-00002.arrow'						)						)						)
           self.assertDictEqual(builder.info.features							,   datasets.Features({"content": datasets.Value("string"						)}						)						)
           UpperCamelCase__			:   List[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           # Order is not preserved when sharding, so we just check that all the elements are there
           self.assertListEqual(sorted(dset["train"]["content"]						)							,   sorted(["foo", "bar", "foobar"]						)						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset







 @require_beam
 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   int       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE						)
           self.assertRaises(datasets.builder.MissingBeamOptions							,   builder.download_and_prepare						)







 @require_beam
 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  len(get_test_nested_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   Any       				=  NestedBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train.arrow'						)						)						)
           self.assertDictEqual(
               builder.info.features							,   datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"						)}						)}						)						)
           UpperCamelCase__			:   List[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           self.assertDictEqual(dset["train"][0]							,   get_test_nested_examples()[0][1]						)
           self.assertDictEqual(
               dset["train"][expected_num_examples - 1]							,   get_test_nested_examples()[expected_num_examples - 1][1]						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset