



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


_snake_case		=  16
_snake_case		=  32






def 				A				(					_lowerCamelCase  ):





							'''simple docstring'''




							return int(x / 2**20  )



class   UpperCAmelCase_    :







							def __enter__(   self):

														'''simple docstring'''

														gc.collect()
														torch.cuda.empty_cache()
														torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
														_lowerCAmelCase     :							List[str]    =    torch.cuda.memory_allocated()
														return self







							def __exit__(   self,			*__a):

														'''simple docstring'''

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






def 				A				(					_lowerCamelCase				, _lowerCamelCase = 16				, _lowerCamelCase = "bert-base-cased"				, _lowerCamelCase = 320				, _lowerCamelCase = 160				, ):





							'''simple docstring'''




							_lowerCAmelCase     :							Optional[Any]    =    AutoTokenizer.from_pretrained(_lowerCamelCase  )
							_lowerCAmelCase     :							Optional[int]    =    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)
														_lowerCAmelCase     :							Optional[Any]    =    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
							_lowerCAmelCase     :							List[str]    =    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
							_lowerCAmelCase     :							Any    =    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.
							_lowerCAmelCase     :							Dict    =    DataLoader(
							    tokenized_datasets["train"]				, shuffle=_lowerCamelCase				, collate_fn=_lowerCamelCase				, batch_size=_lowerCamelCase  )
							_lowerCAmelCase     :							Dict    =    DataLoader(
							    tokenized_datasets["validation"]				, shuffle=_lowerCamelCase				, collate_fn=_lowerCamelCase				, batch_size=_lowerCamelCase  )

							return train_dataloader, eval_dataloader






def 				A				(					_lowerCamelCase				, _lowerCamelCase  ):





							'''simple docstring'''




							_lowerCAmelCase     :							Dict    =    Accelerator()

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

							set_seed(_lowerCamelCase  )
							_lowerCAmelCase   ,    _lowerCAmelCase     :							List[str]    =    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)
							_lowerCAmelCase     :							Any    =    AutoModelForSequenceClassification.from_pretrained(_lowerCamelCase				, return_dict=_lowerCamelCase  )

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

							if accelerator.state.deepspeed_plugin is not None:
														_lowerCAmelCase     :							Optional[Any]    =    accelerator.state.deepspeed_plugin.deepspeed_config[
														    "gradient_accumulation_steps"
														]
							else:
														_lowerCAmelCase     :							int    =    1
							_lowerCAmelCase     :							List[str]    =    (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
							):
														_lowerCAmelCase     :							int    =    get_linear_schedule_with_warmup(
														    optimizer=_lowerCamelCase				, num_warmup_steps=0				, num_training_steps=_lowerCamelCase				, )
							else:
														_lowerCAmelCase     :							Optional[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.
							_lowerCAmelCase   ,    _lowerCAmelCase   ,    _lowerCAmelCase   ,    _lowerCAmelCase   ,    _lowerCAmelCase     :							int    =    accelerator.prepare(
							    _lowerCamelCase				, _lowerCamelCase				, _lowerCamelCase				, _lowerCamelCase				, _lowerCamelCase  )

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

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





							'''simple docstring'''




							_lowerCAmelCase     :							Dict    =    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."				, )
							_lowerCAmelCase     :							int    =    parser.parse_args()
							_lowerCAmelCase     :							int    =    {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
							training_function(_lowerCamelCase				, _lowerCamelCase  )


if __name__ == "__main__":
						main()





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


_snake_case		=  get_tests_dir("fixtures")



class   UpperCAmelCase_    (       unittest.TestCase):







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[Any]    =    mock.Mock()
														_lowerCAmelCase     :							int    =    500
														_lowerCAmelCase     :							Tuple    =    {}
														_lowerCAmelCase     :							str    =    HTTPError
														_lowerCAmelCase     :							Union[str, Any]    =    {}

														# Download this model to make sure it's in the cache.
														_lowerCAmelCase     :							Tuple    =    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=__a) as mock_head:
																					_lowerCAmelCase     :							Optional[int]    =    ViTImageProcessor.from_pretrained("hf-internal-testing/tiny-random-vit")
																					# This check we did call the fake head request
																					mock_head.assert_called()







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Tuple    =    ViTImageProcessor.from_pretrained(
														    "https://huggingface.co/hf-internal-testing/tiny-random-vit/resolve/main/preprocessor_config.json")







							def    snake_case__      (   self):

														'''simple docstring'''

														with self.assertRaises(__a):
																					# config is in subfolder, the following should not work without specifying the subfolder
																					_lowerCAmelCase     :							int    =    AutoImageProcessor.from_pretrained("hf-internal-testing/stable-diffusion-all-variants")

														_lowerCAmelCase     :							Optional[Any]    =    AutoImageProcessor.from_pretrained(
														    "hf-internal-testing/stable-diffusion-all-variants",			subfolder="feature_extractor")

														self.assertIsNotNone(__a)



@is_staging_test
class   UpperCAmelCase_    (       unittest.TestCase):







							@classmethod
							def    snake_case__      (   cls):

														'''simple docstring'''

														_lowerCAmelCase     :							Union[str, Any]    =    TOKEN
														HfFolder.save_token(__a)







							@classmethod
							def    snake_case__      (   cls):

														'''simple docstring'''

														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    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[int]    =    ViTImageProcessor.from_pretrained(__a)
														image_processor.push_to_hub("test-image-processor",			use_auth_token=self._token)

														_lowerCAmelCase     :							str    =    ViTImageProcessor.from_pretrained(f"{USER}/test-image-processor")
														for k, v in image_processor.__dict__.items():
																					self.assertEqual(__a,			getattr(__a,			__a))

														# 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(
																					    __a,			repo_id="test-image-processor",			push_to_hub=__a,			use_auth_token=self._token)

														_lowerCAmelCase     :							Optional[int]    =    ViTImageProcessor.from_pretrained(f"{USER}/test-image-processor")
														for k, v in image_processor.__dict__.items():
																					self.assertEqual(__a,			getattr(__a,			__a))







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Any    =    ViTImageProcessor.from_pretrained(__a)
														image_processor.push_to_hub("valid_org/test-image-processor",			use_auth_token=self._token)

														_lowerCAmelCase     :							Tuple    =    ViTImageProcessor.from_pretrained("valid_org/test-image-processor")
														for k, v in image_processor.__dict__.items():
																					self.assertEqual(__a,			getattr(__a,			__a))

														# 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(
																					    __a,			repo_id="valid_org/test-image-processor-org",			push_to_hub=__a,			use_auth_token=self._token)

														_lowerCAmelCase     :							Optional[int]    =    ViTImageProcessor.from_pretrained("valid_org/test-image-processor-org")
														for k, v in image_processor.__dict__.items():
																					self.assertEqual(__a,			getattr(__a,			__a))







							def    snake_case__      (   self):

														'''simple docstring'''

														CustomImageProcessor.register_for_auto_class()
														_lowerCAmelCase     :							List[str]    =    CustomImageProcessor.from_pretrained(__a)

														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"},			)

														_lowerCAmelCase     :							Tuple    =    AutoImageProcessor.from_pretrained(
														    f"{USER}/test-dynamic-image-processor",			trust_remote_code=__a)
														# 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")
