


"""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__				:							Union[str, Any]						=     1_6
a__				:							Any						=     3_2




def    UpperCAmelCase__       (lowerCAmelCase_   ):




					'''simple docstring'''


					return int(x / 2**20   )




class  UpperCamelCase_	:

	"""simple docstring"""





	def __enter__( self			:				Optional[Any]   )       ->	List[Any]:
						gc.collect()
						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
						__SCREAMING_SNAKE_CASE						   =				torch.cuda.memory_allocated()
						return self




	def __exit__( self			:				int			,			*UpperCAmelCase__			:				int   )       ->	List[Any]:
						gc.collect()
						torch.cuda.empty_cache()
						__SCREAMING_SNAKE_CASE						   =				torch.cuda.memory_allocated()
						__SCREAMING_SNAKE_CASE						   =				torch.cuda.max_memory_allocated()
						__SCREAMING_SNAKE_CASE						   =				bamb(self.end - self.begin   )
						__SCREAMING_SNAKE_CASE						   =				bamb(self.peak - self.begin   )
						# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")




def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_ = 16  ,						lowerCAmelCase_ = "bert-base-cased"  ,						lowerCAmelCase_ = 320  ,						lowerCAmelCase_ = 160  ,						):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				AutoTokenizer.from_pretrained(_lowerCamelCase   )
					__SCREAMING_SNAKE_CASE						   =				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)
										__SCREAMING_SNAKE_CASE						   =				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
					__SCREAMING_SNAKE_CASE						   =				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
					__SCREAMING_SNAKE_CASE						   =				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.
					__SCREAMING_SNAKE_CASE						   =				DataLoader(
					    tokenized_datasets["train"]  ,						shuffle=_lowerCamelCase  ,						collate_fn=_lowerCamelCase  ,						batch_size=_lowerCamelCase   )
					__SCREAMING_SNAKE_CASE						   =				DataLoader(
					    tokenized_datasets["validation"]  ,						shuffle=_lowerCamelCase  ,						collate_fn=_lowerCamelCase  ,						batch_size=_lowerCamelCase   )

					return train_dataloader, eval_dataloader




def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_   ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				Accelerator()

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

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

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

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

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

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




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				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."  ,						)
					__SCREAMING_SNAKE_CASE						   =				parser.parse_args()
					__SCREAMING_SNAKE_CASE						   =				{"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")
