


"""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


UpperCAmelCase      		=			16
UpperCAmelCase      		=			32






def       lowerCamelCase					(a_						:List[Any])     ->							Tuple:
       return int(x / 2**20)
class 					__magic_name__		:






       def __enter__(   self  :    Optional[Any]				):

              '''simple docstring'''






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







       def __exit__(   self  :    str  ,     *snake_case__  :    int				):

              '''simple docstring'''






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







def       lowerCamelCase					(a_						:Accelerator				,				a_						:int = 16				,				a_						:str = "bert-base-cased"				,				a_						:int = 320				,				a_						:int = 160				,				)     ->							Any:
       lowercase		:Any       							=				AutoTokenizer.from_pretrained(a_)
       lowercase		:List[str]       							=				load_dataset(
           '''glue'''				,				'''mrpc'''				,				split={'''train''': F"""train[:{n_train}]""", '''validation''': F"""validation[:{n_val}]"""})

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

       # Apply the method we just defined to all the examples in all the splits of the dataset
       lowercase		:Any       							=				datasets.map(
           a_				,				batched=a_				,				remove_columns=['''idx''', '''sentence1''', '''sentence2''']				,				load_from_cache_file=a_)

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

       def collate_fn(a_						:Union[str, Any]):
              # 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(a_				,				padding='''max_length'''				,				max_length=128				,				return_tensors='''pt''')
              return tokenizer.pad(a_				,				padding='''longest'''				,				return_tensors='''pt''')

       # Instantiate dataloaders.
       lowercase		:Any       							=				DataLoader(
           tokenized_datasets['''train''']				,				shuffle=a_				,				collate_fn=a_				,				batch_size=a_)
       lowercase		:int       							=				DataLoader(
           tokenized_datasets['''validation''']				,				shuffle=a_				,				collate_fn=a_				,				batch_size=a_)

       return train_dataloader, eval_dataloader






def       lowerCamelCase					(a_						:Optional[Any]				,				a_						:int)     ->							Tuple:
       # Initialize accelerator
       lowercase		:int       							=				Accelerator()

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

       set_seed(a_)
       lowercase							,  lowercase		:List[str]       							=				get_dataloaders(a_				,				a_				,				a_				,				args.n_train				,				args.n_val)

       # Instantiate the model (we build the model here so that the seed also control new weights initialization)
       lowercase		:List[str]       							=				AutoModelForSequenceClassification.from_pretrained(a_				,				return_dict=a_)

       # Instantiate optimizer
       lowercase		:List[str]       							=				(
           AdamW
           if accelerator.state.deepspeed_plugin is None
           or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
           else DummyOptim
       )
       lowercase		:Union[str, Any]       							=				optimizer_cls(params=model.parameters()				,				lr=a_)

       if accelerator.state.deepspeed_plugin is not None:
              lowercase		:Optional[Any]       							=				accelerator.state.deepspeed_plugin.deepspeed_config[
                  '''gradient_accumulation_steps'''
              ]
       else:
              lowercase		:Optional[Any]       							=				1
       lowercase		:Optional[Any]       							=				(len(a_) * num_epochs) // gradient_accumulation_steps

       # Instantiate scheduler
       if (
           accelerator.state.deepspeed_plugin is None
           or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
       ):
              lowercase		:List[str]       							=				get_linear_schedule_with_warmup(
                  optimizer=a_				,				num_warmup_steps=0				,				num_training_steps=a_				,				)
       else:
              lowercase		:List[Any]       							=				DummyScheduler(a_				,				total_num_steps=a_				,				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.
       lowercase							,  lowercase							,  lowercase							,  lowercase							,  lowercase		:Dict       							=				accelerator.prepare(
           a_				,				a_				,				a_				,				a_				,				a_)

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

       # Now we train the model
       lowercase		:str       							=				{}
       for epoch in range(a_				,				a_):
              with TorchTracemalloc() as tracemalloc:
                     model.train()
                     for step, batch in enumerate(a_):
                            lowercase		:Any       							=				model(**a_)
                            lowercase		:List[Any]       							=				outputs.loss
                            lowercase		:List[str]       							=				loss / gradient_accumulation_steps
                            accelerator.backward(a_)
                            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)))
              lowercase		:str       							=				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(a_				,				a_)






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


if __name__ == "__main__":
 main()





"""simple docstring"""



# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# 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.
import re

from ..utils import cached_file


# docstyle-ignore
UpperCAmelCase      		=			'''
Human: <<task>>

Assistant: '''


UpperCAmelCase      		=			'''huggingface-tools/default-prompts'''
UpperCAmelCase      		=			{'''chat''': '''chat_prompt_template.txt''', '''run''': '''run_prompt_template.txt'''}






def       lowerCamelCase					(a_						:int				,				a_						:str				,				a_						:Dict="run")     ->							Optional[Any]:
       if prompt_or_repo_id is None:
              lowercase		:Tuple       							=				DEFAULT_PROMPTS_REPO

       # prompt is considered a repo ID when it does not contain any kind of space
       if re.search('''\\s'''				,				a_) is not None:
              return prompt_or_repo_id

       lowercase		:List[str]       							=				cached_file(
           a_				,				PROMPT_FILES[mode]				,				repo_type='''dataset'''				,				user_agent={'''agent''': agent_name})
       with open(a_				,				'''r'''				,				encoding='''utf-8''') as f:
              return f.read()

