


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  _snake_case       (					A	)						->   str:
							return int(x / 2**20	)




class 		a__    :





							'''simple docstring'''






							def __enter__(		self       )	->							Tuple:
														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 ,   *lowerCamelCase_       )	->							Tuple:
														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  _snake_case       (					A			,   A = 16			,   A = "bert-base-cased"			,   A = 320			,   A = 160			,   )						->   List[str]:
							lowerCAmelCase__	       =     AutoTokenizer.from_pretrained(A	)
							lowerCAmelCase__	       =     load_dataset(
							    '''glue'''			,   '''mrpc'''			,   split={'''train''': F"""train[:{n_train}]""", '''validation''': F"""validation[:{n_val}]"""}	)

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

							def collate_fn(A	):
														# 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.
							lowerCAmelCase__	       =     DataLoader(
							    tokenized_datasets['''train''']			,   shuffle=A			,   collate_fn=A			,   batch_size=A	)
							lowerCAmelCase__	       =     DataLoader(
							    tokenized_datasets['''validation''']			,   shuffle=A			,   collate_fn=A			,   batch_size=A	)

							return train_dataloader, eval_dataloader


def  _snake_case       (					A			,   A	)						->   str:
							# Initialize accelerator
							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(A	)
							lowerCAmelCase__       , lowerCAmelCase__	       =     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)
							lowerCAmelCase__	       =     AutoModelForSequenceClassification.from_pretrained(A			,   return_dict=A	)

							# 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=A	)

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

							# 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(A			,   A	):
														with TorchTracemalloc() as tracemalloc:
																					model.train()
																					for step, batch in enumerate(A	):
																												lowerCAmelCase__	       =     model(**A	)
																												lowerCAmelCase__	       =     outputs.loss
																												lowerCAmelCase__	       =     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	)	)	)
														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(A			,   A	)


def  _snake_case       (					)						->   int:
							lowerCAmelCase__	       =     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.'''			,   )
							lowerCAmelCase__	       =     parser.parse_args()
							lowerCAmelCase__	       =     {'''lr''': 2E-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}
							training_function(A			,   A	)


if __name__ == "__main__":
		main()



'''simple docstring'''






import unittest

from transformers import (
    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
    TextClassificationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, nested_simplify, require_tf, require_torch, slow

from .test_pipelines_common import ANY


# These 2 model types require different inputs than those of the usual text models.
__UpperCAmelCase       =							{'''LayoutLMv2Config''', '''LayoutLMv3Config'''}




@is_pipeline_test
class 		a__    (							unittest.TestCase  ):





							'''simple docstring'''






							lowercase__      :       Optional[Any]     						=				MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
							lowercase__      :       List[Any]     						=				TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING

							if model_mapping is not None:
														lowercase__      :       Optional[Any]     						=				{config: model for config, model in model_mapping.items() if config.__name__ not in _TO_SKIP}
							if tf_model_mapping is not None:
														lowercase__      :       Tuple     						=				{
														    config: model for config, model in tf_model_mapping.items() if config.__name__ not in _TO_SKIP
														}







							@require_torch
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[int]:
														lowerCAmelCase__	       =     pipeline(
														    task='''text-classification''' ,   model='''hf-internal-testing/tiny-random-distilbert''' ,   framework='''pt'''       )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )

														lowerCAmelCase__	       =     text_classifier('''This is great !''' ,   top_k=2       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}]       )

														lowerCAmelCase__	       =     text_classifier(['''This is great !''', '''This is bad'''] ,   top_k=2       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [
														        [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}],
														        [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}],
														    ] ,   )

														lowerCAmelCase__	       =     text_classifier('''This is great !''' ,   top_k=1       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )

														# Legacy behavior
														lowerCAmelCase__	       =     text_classifier('''This is great !''' ,   return_all_scores=lowerCamelCase_       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )

														lowerCAmelCase__	       =     text_classifier('''This is great !''' ,   return_all_scores=lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [[{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}]]       )

														lowerCAmelCase__	       =     text_classifier(['''This is great !''', '''Something else'''] ,   return_all_scores=lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [
														        [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}],
														        [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}],
														    ] ,   )

														lowerCAmelCase__	       =     text_classifier(['''This is great !''', '''Something else'''] ,   return_all_scores=lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [
														        {'''label''': '''LABEL_0''', '''score''': 0.504},
														        {'''label''': '''LABEL_0''', '''score''': 0.504},
														    ] ,   )







							@require_torch
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							int:
														import torch

														lowerCAmelCase__	       =     pipeline(
														    task='''text-classification''' ,   model='''hf-internal-testing/tiny-random-distilbert''' ,   framework='''pt''' ,   device=torch.device('''cpu'''       ) ,   )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )







							@require_tf
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							str:
														lowerCAmelCase__	       =     pipeline(
														    task='''text-classification''' ,   model='''hf-internal-testing/tiny-random-distilbert''' ,   framework='''tf'''       )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )







							@slow
							@require_torch
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Dict:
														lowerCAmelCase__	       =     pipeline('''text-classification'''       )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''POSITIVE''', '''score''': 1.0}]       )
														lowerCAmelCase__	       =     text_classifier('''This is bad !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''NEGATIVE''', '''score''': 1.0}]       )
														lowerCAmelCase__	       =     text_classifier('''Birds are a type of animal'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''POSITIVE''', '''score''': 0.988}]       )







							@slow
							@require_tf
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														lowerCAmelCase__	       =     pipeline('''text-classification''' ,   framework='''tf'''       )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''POSITIVE''', '''score''': 1.0}]       )
														lowerCAmelCase__	       =     text_classifier('''This is bad !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''NEGATIVE''', '''score''': 1.0}]       )
														lowerCAmelCase__	       =     text_classifier('''Birds are a type of animal'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''POSITIVE''', '''score''': 0.988}]       )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							Tuple:
														lowerCAmelCase__	       =     TextClassificationPipeline(model=lowerCamelCase_ ,   tokenizer=lowerCamelCase_       )
														return text_classifier, ["HuggingFace is in", "This is another test"]







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							Union[str, Any]:
														lowerCAmelCase__	       =     text_classifier.model
														# Small inputs because BartTokenizer tiny has maximum position embeddings = 22
														lowerCAmelCase__	       =     '''HuggingFace is in'''
														lowerCAmelCase__	       =     text_classifier(lowerCamelCase_       )

														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}]       )
														self.assertTrue(outputs[0]['''label'''] in model.config.idalabel.values()       )

														lowerCAmelCase__	       =     ['''HuggingFace is in ''', '''Paris is in France''']
														lowerCAmelCase__	       =     text_classifier(lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}, {'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}] ,   )
														self.assertTrue(outputs[0]['''label'''] in model.config.idalabel.values()       )
														self.assertTrue(outputs[1]['''label'''] in model.config.idalabel.values()       )

														# Forcing to get all results with `top_k=None`
														# This is NOT the legacy format
														lowerCAmelCase__	       =     text_classifier(lowerCamelCase_ ,   top_k=lowerCamelCase_       )
														lowerCAmelCase__	       =     len(model.config.idalabel.values()       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [[{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}] * N, [{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}] * N] ,   )

														lowerCAmelCase__	       =     {'''text''': '''HuggingFace is in ''', '''text_pair''': '''Paris is in France'''}
														lowerCAmelCase__	       =     text_classifier(lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   {'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )} ,   )
														self.assertTrue(outputs['''label'''] in model.config.idalabel.values()       )

														# This might be used a text pair, but tokenizer + pipe interaction
														# makes it hard to understand that it's not using the pair properly
														# https://github.com/huggingface/transformers/issues/17305
														# We disabled this usage instead as it was outputting wrong outputs.
														lowerCAmelCase__	       =     [['''HuggingFace is in ''', '''Paris is in France''']]
														with self.assertRaises(lowerCamelCase_       ):
																					text_classifier(lowerCamelCase_       )

														# This used to be valid for doing text pairs
														# We're keeping it working because of backward compatibility
														lowerCAmelCase__	       =     text_classifier([[['''HuggingFace is in ''', '''Paris is in France''']]]       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}] ,   )
														self.assertTrue(outputs[0]['''label'''] in model.config.idalabel.values()       )