


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


a_   :      int				    =      1_6
a_   :      List[str]				    =      3_2






def 		_A			(lowerCAmelCase__					:Optional[int]						)				->					Dict:




     '''simple docstring'''



     return int(x / 2**20						)




class 	a       :
 def __enter__(	self  )   ->							Any:
      gc.collect()
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
      _a      =							torch.cuda.memory_allocated()
      return self
 def __exit__(	self ,			*__magic_name__  )   ->							Tuple:
      gc.collect()
      torch.cuda.empty_cache()
      _a      =							torch.cuda.memory_allocated()
      _a      =							torch.cuda.max_memory_allocated()
      _a      =							bamb(self.end - self.begin  )
      _a      =							bamb(self.peak - self.begin  )
      # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")






def 		_A			(lowerCAmelCase__					:Accelerator			,  lowerCAmelCase__					:int = 16			,  lowerCAmelCase__					:str = "bert-base-cased"			,  lowerCAmelCase__					:int = 3_20			,  lowerCAmelCase__					:int = 1_60			,  )				->					Union[str, Any]:




     '''simple docstring'''



     _a      =							AutoTokenizer.from_pretrained(lowerCAmelCase__						)
     _a      =							load_dataset(
         'glue'			,  'mrpc'			,  split={'train': f'train[:{n_train}]', 'validation': f'validation[:{n_val}]'}						)

     def tokenize_function(lowerCAmelCase__					:Union[str, Any]						):
          # max_length=None => use the model max length (it's actually the default)
          _a      =							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
     _a      =							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
     _a      =							tokenized_datasets.rename_column('label'			,  'labels'						)

     def collate_fn(lowerCAmelCase__					: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(lowerCAmelCase__			,  padding='max_length'			,  max_length=1_28			,  return_tensors='pt'						)
          return tokenizer.pad(lowerCAmelCase__			,  padding='longest'			,  return_tensors='pt'						)

     # Instantiate dataloaders.
     _a      =							DataLoader(
         tokenized_datasets['train']			,  shuffle=lowerCAmelCase__			,  collate_fn=lowerCAmelCase__			,  batch_size=lowerCAmelCase__						)
     _a      =							DataLoader(
         tokenized_datasets['validation']			,  shuffle=lowerCAmelCase__			,  collate_fn=lowerCAmelCase__			,  batch_size=lowerCAmelCase__						)

     return train_dataloader, eval_dataloader






def 		_A			(lowerCAmelCase__					:Dict			,  lowerCAmelCase__					:int						)				->					Tuple:




     '''simple docstring'''



     _a      =							Accelerator()

     # 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      =							args.model_name_or_path

     set_seed(lowerCAmelCase__						)
     _a ,				_a      =							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)
     _a      =							AutoModelForSequenceClassification.from_pretrained(lowerCAmelCase__			,  return_dict=lowerCAmelCase__						)

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

     if accelerator.state.deepspeed_plugin is not None:
          _a      =							accelerator.state.deepspeed_plugin.deepspeed_config[
              'gradient_accumulation_steps'
          ]
     else:
          _a      =							1
     _a      =							(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
     ):
          _a      =							get_linear_schedule_with_warmup(
              optimizer=lowerCAmelCase__			,  num_warmup_steps=0			,  num_training_steps=lowerCAmelCase__			,  )
     else:
          _a      =							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.
     _a ,				_a ,				_a ,				_a ,				_a      =							accelerator.prepare(
         lowerCAmelCase__			,  lowerCAmelCase__			,  lowerCAmelCase__			,  lowerCAmelCase__			,  lowerCAmelCase__						)

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

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




     '''simple docstring'''



     _a      =							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=3_20			,  help='Number of training examples to use.'			,  )
     parser.add_argument(
         '--n_val'			,  type=lowerCAmelCase__			,  default=1_60			,  help='Number of validation examples to use.'			,  )
     parser.add_argument(
         '--num_epochs'			,  type=lowerCAmelCase__			,  default=1			,  help='Number of train epochs.'			,  )
     _a      =							parser.parse_args()
     _a      =							{'lr': 2E-5, 'num_epochs': args.num_epochs, 'seed': 42, 'batch_size': 16}
     training_function(lowerCAmelCase__			,  lowerCAmelCase__						)


if __name__ == "__main__":
 main()





'''simple docstring'''






import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoImageProcessor, ViTImageProcessor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / "utils"))

from test_module.custom_image_processing import CustomImageProcessor  # noqa E402


a_   :      Tuple				    =      get_tests_dir("fixtures")




class 	a       (       unittest.TestCase    ):
 def 							__UpperCAmelCase						(	self  )   ->							str:
      # A mock response for an HTTP head request to emulate server down
      _a      =							mock.Mock()
      _a      =							5_00
      _a      =							{}
      _a      =							HTTPError
      _a      =							{}

      # Download this model to make sure it's in the cache.
      _a      =							ViTImageProcessor.from_pretrained('hf-internal-testing/tiny-random-vit'  )
      # Under the mock environment we get a 500 error when trying to reach the model.
      with mock.patch('requests.Session.request' ,			return_value=__magic_name__  ) as mock_head:
           _a      =							ViTImageProcessor.from_pretrained('hf-internal-testing/tiny-random-vit'  )
           # This check we did call the fake head request
           mock_head.assert_called()
 def 							__UpperCAmelCase						(	self  )   ->							Optional[int]:
      # This test is for deprecated behavior and can be removed in v5
      _a      =							ViTImageProcessor.from_pretrained(
          'https://huggingface.co/hf-internal-testing/tiny-random-vit/resolve/main/preprocessor_config.json'  )




 def 							__UpperCAmelCase						(	self  )   ->							Dict:
      with self.assertRaises(__magic_name__  ):
           # config is in subfolder, the following should not work without specifying the subfolder
           _a      =							AutoImageProcessor.from_pretrained('hf-internal-testing/stable-diffusion-all-variants'  )

      _a      =							AutoImageProcessor.from_pretrained(
          'hf-internal-testing/stable-diffusion-all-variants' ,			subfolder='feature_extractor'  )

      self.assertIsNotNone(__magic_name__  )





@is_staging_test
class 	a       (       unittest.TestCase    ):
 @classmethod
 def 							__UpperCAmelCase						(	cls  )   ->							Dict:
      _a      =							TOKEN
      HfFolder.save_token(__magic_name__  )
 @classmethod
 def 							__UpperCAmelCase						(	cls  )   ->							List[Any]:
      try:
           delete_repo(token=cls._token ,			repo_id='test-image-processor'  )
      except HTTPError:
           pass

      try:
           delete_repo(token=cls._token ,			repo_id='valid_org/test-image-processor-org'  )
      except HTTPError:
           pass

      try:
           delete_repo(token=cls._token ,			repo_id='test-dynamic-image-processor'  )
      except HTTPError:
           pass
 def 							__UpperCAmelCase						(	self  )   ->							str:
      _a      =							ViTImageProcessor.from_pretrained(__magic_name__  )
      image_processor.push_to_hub('test-image-processor' ,			use_auth_token=self._token  )

      _a      =							ViTImageProcessor.from_pretrained(f'{USER}/test-image-processor'  )
      for k, v in image_processor.__dict__.items():
           self.assertEqual(__magic_name__ ,			getattr(__magic_name__ ,			__magic_name__  )  )

      # Reset repo
      delete_repo(token=self._token ,			repo_id='test-image-processor'  )

      # Push to hub via save_pretrained
      with tempfile.TemporaryDirectory() as tmp_dir:
           image_processor.save_pretrained(
               __magic_name__ ,			repo_id='test-image-processor' ,			push_to_hub=__magic_name__ ,			use_auth_token=self._token  )

      _a      =							ViTImageProcessor.from_pretrained(f'{USER}/test-image-processor'  )
      for k, v in image_processor.__dict__.items():
           self.assertEqual(__magic_name__ ,			getattr(__magic_name__ ,			__magic_name__  )  )
 def 							__UpperCAmelCase						(	self  )   ->							Optional[Any]:
      _a      =							ViTImageProcessor.from_pretrained(__magic_name__  )
      image_processor.push_to_hub('valid_org/test-image-processor' ,			use_auth_token=self._token  )

      _a      =							ViTImageProcessor.from_pretrained('valid_org/test-image-processor'  )
      for k, v in image_processor.__dict__.items():
           self.assertEqual(__magic_name__ ,			getattr(__magic_name__ ,			__magic_name__  )  )

      # Reset repo
      delete_repo(token=self._token ,			repo_id='valid_org/test-image-processor'  )

      # Push to hub via save_pretrained
      with tempfile.TemporaryDirectory() as tmp_dir:
           image_processor.save_pretrained(
               __magic_name__ ,			repo_id='valid_org/test-image-processor-org' ,			push_to_hub=__magic_name__ ,			use_auth_token=self._token  )

      _a      =							ViTImageProcessor.from_pretrained('valid_org/test-image-processor-org'  )
      for k, v in image_processor.__dict__.items():
           self.assertEqual(__magic_name__ ,			getattr(__magic_name__ ,			__magic_name__  )  )




 def 							__UpperCAmelCase						(	self  )   ->							Any:
      CustomImageProcessor.register_for_auto_class()
      _a      =							CustomImageProcessor.from_pretrained(__magic_name__  )

      image_processor.push_to_hub('test-dynamic-image-processor' ,			use_auth_token=self._token  )

      # This has added the proper auto_map field to the config
      self.assertDictEqual(
          image_processor.auto_map ,			{'AutoImageProcessor': 'custom_image_processing.CustomImageProcessor'} ,			)

      _a      =							AutoImageProcessor.from_pretrained(
          f'{USER}/test-dynamic-image-processor' ,			trust_remote_code=__magic_name__  )
      # Can't make an isinstance check because the new_image_processor is from the CustomImageProcessor class of a dynamic module
      self.assertEqual(new_image_processor.__class__.__name__ ,			'CustomImageProcessor'  )

