




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__ :    Optional[int]         =      16
UpperCamelCase__ :    Optional[Any]         =      32







def      _a				(						SCREAMING_SNAKE_CASE_							:    Union[str, Any]					):
					return int(x / 2**20					)
class       a__   :






						def __enter__(       self      ):




											"""simple docstring"""



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






						def __exit__(       self		,      *_A      ):




											"""simple docstring"""



											gc.collect()
											torch.cuda.empty_cache()
											__lowerCAmelCase           =	torch.cuda.memory_allocated()
											__lowerCAmelCase           =	torch.cuda.max_memory_allocated()
											__lowerCAmelCase           =	bamb(self.end - self.begin      )
											__lowerCAmelCase           =	bamb(self.peak - self.begin      )
											# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")







def      _a				(						SCREAMING_SNAKE_CASE_							:    Accelerator				,				SCREAMING_SNAKE_CASE_							:    int = 16				,				SCREAMING_SNAKE_CASE_							:    str = "bert-base-cased"				,				SCREAMING_SNAKE_CASE_							:    int = 3_20				,				SCREAMING_SNAKE_CASE_							:    int = 1_60				,				):
					__lowerCAmelCase           =	AutoTokenizer.from_pretrained(lowerCamelCase_					)
					__lowerCAmelCase           =	load_dataset(
					    "glue"				,				"mrpc"				,				split={"train": F"""train[:{n_train}]""", "validation": F"""validation[:{n_val}]"""}					)

					def tokenize_function(SCREAMING_SNAKE_CASE_							:    Any					):
										# max_length=None => use the model max length (it's actually the default)
										__lowerCAmelCase           =	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           =	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           =	tokenized_datasets.rename_column("label"				,				"labels"					)

					def collate_fn(SCREAMING_SNAKE_CASE_							:    List[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.
					__lowerCAmelCase           =	DataLoader(
					    tokenized_datasets["train"]				,				shuffle=lowerCamelCase_				,				collate_fn=lowerCamelCase_				,				batch_size=lowerCamelCase_					)
					__lowerCAmelCase           =	DataLoader(
					    tokenized_datasets["validation"]				,				shuffle=lowerCamelCase_				,				collate_fn=lowerCamelCase_				,				batch_size=lowerCamelCase_					)

					return train_dataloader, eval_dataloader







def      _a				(						SCREAMING_SNAKE_CASE_							:    Dict				,				SCREAMING_SNAKE_CASE_							:    Any					):
					__lowerCAmelCase           =	Accelerator()

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

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

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

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

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

					# Now we train the model
					__lowerCAmelCase           =	{}
					for epoch in range(lowerCamelCase_				,				lowerCamelCase_					):
										with TorchTracemalloc() as tracemalloc:
															model.train()
															for step, batch in enumerate(lowerCamelCase_					):
																				__lowerCAmelCase           =	model(**lowerCamelCase_					)
																				__lowerCAmelCase           =	outputs.loss
																				__lowerCAmelCase           =	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           =	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				(						):
					__lowerCAmelCase           =	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."				,				)
					__lowerCAmelCase           =	parser.parse_args()
					__lowerCAmelCase           =	{"""lr""": 2E-5, """num_epochs""": args.num_epochs, """seed""": 42, """batch_size""": 16}
					training_function(lowerCamelCase_				,				lowerCamelCase_					)


if __name__ == "__main__":
					main()










from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
					from transformers import (
					    TFAutoModel,
					    TFAutoModelForCausalLM,
					    TFAutoModelForMaskedLM,
					    TFAutoModelForPreTraining,
					    TFAutoModelForQuestionAnswering,
					    TFAutoModelForSeqaSeqLM,
					    TFAutoModelForSequenceClassification,
					    TFAutoModelForTableQuestionAnswering,
					    TFAutoModelForTokenClassification,
					    TFAutoModelWithLMHead,
					    TFBertForMaskedLM,
					    TFBertForPreTraining,
					    TFBertForQuestionAnswering,
					    TFBertForSequenceClassification,
					    TFBertModel,
					    TFFunnelBaseModel,
					    TFFunnelModel,
					    TFGPTaLMHeadModel,
					    TFRobertaForMaskedLM,
					    TFTaForConditionalGeneration,
					    TFTapasForQuestionAnswering,
					)
					from transformers.models.auto.modeling_tf_auto import (
					    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
					    TF_MODEL_FOR_MASKED_LM_MAPPING,
					    TF_MODEL_FOR_PRETRAINING_MAPPING,
					    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
					    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
					    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
					    TF_MODEL_MAPPING,
					)
					from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
					from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
					from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
					from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST
class       a__   ( snake_case__				):
						_a							:   Optional[int] =   """new-model"""


if is_tf_available():
					class       a__   ( snake_case__				):
											_a							:   Dict =   NewModelConfig
@require_tf
class       a__   ( unittest.TestCase				):






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	"bert-base-cased"
											__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
											self.assertIsNotNone(_A      )
											self.assertIsInstance(_A		,      _A      )

											__lowerCAmelCase           =	TFAutoModel.from_pretrained(_A      )
											self.assertIsNotNone(_A      )
											self.assertIsInstance(_A		,      _A      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	"bert-base-cased"
											__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
											self.assertIsNotNone(_A      )
											self.assertIsInstance(_A		,      _A      )

											__lowerCAmelCase           =	TFAutoModelForPreTraining.from_pretrained(_A      )
											self.assertIsNotNone(_A      )
											self.assertIsInstance(_A		,      _A      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )

																__lowerCAmelCase           =	TFAutoModelForCausalLM.from_pretrained(_A      )
																__lowerCAmelCase       ,							__lowerCAmelCase           =	TFAutoModelForCausalLM.from_pretrained(_A		,      output_loading_info=_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )

																__lowerCAmelCase           =	TFAutoModelWithLMHead.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )

																__lowerCAmelCase           =	TFAutoModelForMaskedLM.from_pretrained(_A      )
																__lowerCAmelCase       ,							__lowerCAmelCase           =	TFAutoModelForMaskedLM.from_pretrained(_A		,      output_loading_info=_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )

																__lowerCAmelCase           =	TFAutoModelForSeqaSeqLM.from_pretrained(_A      )
																__lowerCAmelCase       ,							__lowerCAmelCase           =	TFAutoModelForSeqaSeqLM.from_pretrained(_A		,      output_loading_info=_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_name in ["bert-base-uncased"]:
																__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )

																__lowerCAmelCase           =	TFAutoModelForSequenceClassification.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_name in ["bert-base-uncased"]:
																__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )

																__lowerCAmelCase           =	TFAutoModelForQuestionAnswering.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )






						@slow
						@require_tensorflow_probability
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
																__lowerCAmelCase           =	AutoConfig.from_pretrained(_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )

																__lowerCAmelCase           =	TFAutoModelForTableQuestionAnswering.from_pretrained(_A      )
																__lowerCAmelCase       ,							__lowerCAmelCase           =	TFAutoModelForTableQuestionAnswering.from_pretrained(
																    _A		,      output_loading_info=_A      )
																self.assertIsNotNone(_A      )
																self.assertIsInstance(_A		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	TFAutoModelWithLMHead.from_pretrained(_A      )
											self.assertIsInstance(_A		,      _A      )
											self.assertEqual(model.num_parameters()		,      1_4_4_1_0      )
											self.assertEqual(model.num_parameters(only_trainable=_A      )		,      1_4_4_1_0      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	TFAutoModelWithLMHead.from_pretrained(_A      )
											self.assertIsInstance(_A		,      _A      )
											self.assertEqual(model.num_parameters()		,      1_4_4_1_0      )
											self.assertEqual(model.num_parameters(only_trainable=_A      )		,      1_4_4_1_0      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	TFAutoModel.from_pretrained("sgugger/funnel-random-tiny"      )
											self.assertIsInstance(_A		,      _A      )

											__lowerCAmelCase           =	copy.deepcopy(model.config      )
											__lowerCAmelCase           =	["FunnelBaseModel"]
											__lowerCAmelCase           =	TFAutoModel.from_config(_A      )
											self.assertIsInstance(_A		,      _A      )

											with tempfile.TemporaryDirectory() as tmp_dir:
																model.save_pretrained(_A      )
																__lowerCAmelCase           =	TFAutoModel.from_pretrained(_A      )
																self.assertIsInstance(_A		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											try:
																AutoConfig.register("new-model"		,      _A      )

																__lowerCAmelCase           =	[
																    TFAutoModel,
																    TFAutoModelForCausalLM,
																    TFAutoModelForMaskedLM,
																    TFAutoModelForPreTraining,
																    TFAutoModelForQuestionAnswering,
																    TFAutoModelForSequenceClassification,
																    TFAutoModelForTokenClassification,
																]

																for auto_class in auto_classes:
																					with self.subTest(auto_class.__name__      ):
																										# Wrong config class will raise an error
																										with self.assertRaises(_A      ):
																															auto_class.register(_A		,      _A      )
																										auto_class.register(_A		,      _A      )
																										# Trying to register something existing in the Transformers library will raise an error
																										with self.assertRaises(_A      ):
																															auto_class.register(_A		,      _A      )

																										# Now that the config is registered, it can be used as any other config with the auto-API
																										__lowerCAmelCase           =	BertModelTester(self      ).get_config()
																										__lowerCAmelCase           =	NewModelConfig(**tiny_config.to_dict()      )
																										__lowerCAmelCase           =	auto_class.from_config(_A      )
																										self.assertIsInstance(_A		,      _A      )

																										with tempfile.TemporaryDirectory() as tmp_dir:
																															model.save_pretrained(_A      )
																															__lowerCAmelCase           =	auto_class.from_pretrained(_A      )
																															self.assertIsInstance(_A		,      _A      )

											finally:
																if "new-model" in CONFIG_MAPPING._extra_content:
																					del CONFIG_MAPPING._extra_content["new-model"]
																for mapping in (
																    TF_MODEL_MAPPING,
																    TF_MODEL_FOR_PRETRAINING_MAPPING,
																    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
																    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
																    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
																    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
																    TF_MODEL_FOR_MASKED_LM_MAPPING,
																):
																					if NewModelConfig in mapping._extra_content:
																										del mapping._extra_content[NewModelConfig]






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											with self.assertRaisesRegex(
											    _A		,      "bert-base is not a local folder and is not a valid model identifier"      ):
																__lowerCAmelCase           =	TFAutoModel.from_pretrained("bert-base"      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											with self.assertRaisesRegex(
											    _A		,      R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"      ):
																__lowerCAmelCase           =	TFAutoModel.from_pretrained(_A		,      revision="aaaaaa"      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											with self.assertRaisesRegex(
											    _A		,      "hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"		,      ):
																__lowerCAmelCase           =	TFAutoModel.from_pretrained("hf-internal-testing/config-no-model"      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											with self.assertRaisesRegex(_A		,      "Use `from_pt=True` to load this model"      ):
																__lowerCAmelCase           =	TFAutoModel.from_pretrained("hf-internal-testing/tiny-bert-pt-only"      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"      )
											with RequestCounter() as counter:
																__lowerCAmelCase           =	TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"      )
																self.assertEqual(counter.get_request_count		,      0      )
																self.assertEqual(counter.head_request_count		,      1      )
																self.assertEqual(counter.other_request_count		,      0      )

											# With a sharded checkpoint
											__lowerCAmelCase           =	TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"      )
											with RequestCounter() as counter:
																__lowerCAmelCase           =	TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"      )
																self.assertEqual(counter.get_request_count		,      0      )
																self.assertEqual(counter.head_request_count		,      1      )
																self.assertEqual(counter.other_request_count		,      0      )




