


import gc
import unittest

from transformers import CTRLConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
         CTRLForSequenceClassification,
         CTRLLMHeadModel,
         CTRLModel,
     )






class   _lowerCamelCase :




       """simple docstring"""



       def __init__(      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=14				,       _SCREAMING_SNAKE_CASE=7				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=99				,       _SCREAMING_SNAKE_CASE=32				,       _SCREAMING_SNAKE_CASE=5				,       _SCREAMING_SNAKE_CASE=4				,       _SCREAMING_SNAKE_CASE=37				,       _SCREAMING_SNAKE_CASE="gelu"				,       _SCREAMING_SNAKE_CASE=0.1				,       _SCREAMING_SNAKE_CASE=0.1				,       _SCREAMING_SNAKE_CASE=512				,       _SCREAMING_SNAKE_CASE=16				,       _SCREAMING_SNAKE_CASE=2				,       _SCREAMING_SNAKE_CASE=0.0_2				,       _SCREAMING_SNAKE_CASE=3				,       _SCREAMING_SNAKE_CASE=4				,       _SCREAMING_SNAKE_CASE=None				,       )->Optional[Any]:




         '''simple docstring'''





         A_      :				List[Any]		     =					parent
         A_      :				Union[str, Any]		     =					batch_size
         A_      :				Tuple		     =					seq_length
         A_      :				Dict		     =					is_training
         A_      :				Optional[int]		     =					use_token_type_ids
         A_      :				Union[str, Any]		     =					use_input_mask
         A_      :				List[str]		     =					use_labels
         A_      :				List[str]		     =					use_mc_token_ids
         A_      :				str		     =					vocab_size
         A_      :				Dict		     =					hidden_size
         A_      :				int		     =					num_hidden_layers
         A_      :				Any		     =					num_attention_heads
         A_      :				List[Any]		     =					intermediate_size
         A_      :				str		     =					hidden_act
         A_      :				Any		     =					hidden_dropout_prob
         A_      :				Tuple		     =					attention_probs_dropout_prob
         A_      :				Any		     =					max_position_embeddings
         A_      :				Optional[Any]		     =					type_vocab_size
         A_      :				List[Any]		     =					type_sequence_label_size
         A_      :				List[str]		     =					initializer_range
         A_      :				Union[str, Any]		     =					num_labels
         A_      :				int		     =					num_choices
         A_      :				Tuple		     =					scope
         A_      :				Dict		     =					self.vocab_size - 1



       def        _snake_case       (      self		)->Dict:




         '''simple docstring'''





         A_      :				Tuple		     =					ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size		)

         A_      :				str		     =					None
         if self.use_input_mask:
           A_      :				Optional[Any]		     =					random_attention_mask([self.batch_size, self.seq_length]		)

         A_      :				str		     =					None
         if self.use_token_type_ids:
           A_      :				str		     =					ids_tensor([self.batch_size, self.seq_length]				,       self.type_vocab_size		)

         A_      :				Union[str, Any]		     =					None
         if self.use_mc_token_ids:
           A_      :				int		     =					ids_tensor([self.batch_size, self.num_choices]				,       self.seq_length		)

         A_      :				List[Any]		     =					None
         A_      :				Tuple		     =					None
         A_      :				Union[str, Any]		     =					None
         if self.use_labels:
           A_      :				Any		     =					ids_tensor([self.batch_size]				,       self.type_sequence_label_size		)
           A_      :				Union[str, Any]		     =					ids_tensor([self.batch_size, self.seq_length]				,       self.num_labels		)
           A_      :				Optional[Any]		     =					ids_tensor([self.batch_size]				,       self.num_choices		)

         A_      :				Tuple		     =					self.get_config()

         A_      :				List[str]		     =					ids_tensor([self.num_hidden_layers, self.num_attention_heads]				,       2		)

         return (
             config,
             input_ids,
             input_mask,
             head_mask,
             token_type_ids,
             mc_token_ids,
             sequence_labels,
             token_labels,
             choice_labels,
         )



       def        _snake_case       (      self		)->Tuple:




         '''simple docstring'''





         return CTRLConfig(
             vocab_size=self.vocab_size				,       n_embd=self.hidden_size				,       n_layer=self.num_hidden_layers				,       n_head=self.num_attention_heads				,       n_positions=self.max_position_embeddings				,       pad_token_id=self.pad_token_id				,       )



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       *_SCREAMING_SNAKE_CASE		)->List[str]:




         '''simple docstring'''





         A_      :				List[str]		     =					CTRLModel(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()

         model(SCREAMING_SNAKE_CASE__				,       token_type_ids=SCREAMING_SNAKE_CASE__				,       head_mask=SCREAMING_SNAKE_CASE__		)
         model(SCREAMING_SNAKE_CASE__				,       token_type_ids=SCREAMING_SNAKE_CASE__		)
         A_      :				Optional[Any]		     =					model(SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)		)
         self.parent.assertEqual(len(result.past_key_values		)				,       config.n_layer		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       *_SCREAMING_SNAKE_CASE		)->Any:




         '''simple docstring'''





         A_      :				List[str]		     =					CTRLLMHeadModel(SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()

         A_      :				Tuple		     =					model(SCREAMING_SNAKE_CASE__				,       token_type_ids=SCREAMING_SNAKE_CASE__				,       labels=SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.loss.shape				,       ()		)
         self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)		)



       def        _snake_case       (      self		)->List[Any]:




         '''simple docstring'''





         A_      :				str		     =					self.prepare_config_and_inputs()

         (
             (
             A_       
         )     ,       (
             A_       
         )     ,       (
             A_       
         )     ,       (
             A_       
         )     ,       (
             A_       
         )     ,       (
             A_       
         )     ,       (
             A_       
         )     ,       (
             A_       
         )     ,       (
             A_       
         )     ,              
         )      :				Union[str, Any]		     =					config_and_inputs

         A_      :				str		     =					{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''head_mask''': head_mask}

         return config, inputs_dict







       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       *_SCREAMING_SNAKE_CASE		)->Tuple:




         '''simple docstring'''





         A_      :				List[str]		     =					self.num_labels
         A_      :				Tuple		     =					CTRLForSequenceClassification(SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         A_      :				int		     =					ids_tensor([self.batch_size]				,       self.type_sequence_label_size		)
         A_      :				str		     =					model(SCREAMING_SNAKE_CASE__				,       token_type_ids=SCREAMING_SNAKE_CASE__				,       labels=SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)		)






@require_torch
class   _lowerCamelCase (					UpperCAmelCase__					,							UpperCAmelCase__					,							UpperCAmelCase__					,							unittest.TestCase			):




       """simple docstring"""
       snake_case      =    (CTRLModel, CTRLLMHeadModel, CTRLForSequenceClassification) if is_torch_available() else ()
       snake_case      =    (CTRLLMHeadModel,) if is_torch_available() else ()
       snake_case      =    (
           {
               """feature-extraction""": CTRLModel,
               """text-classification""": CTRLForSequenceClassification,
               """text-generation""": CTRLLMHeadModel,
               """zero-shot""": CTRLForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       snake_case      =    True
       snake_case      =    False
       snake_case      =    False



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)->str:




         '''simple docstring'''





         if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
           # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
           # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
           # config could not be created.
           return True

         return False



       def        _snake_case       (      self		)->Tuple:




         '''simple docstring'''





         A_      :				Optional[Any]		     =					CTRLModelTester(self		)
         A_      :				int		     =					ConfigTester(self				,       config_class=SCREAMING_SNAKE_CASE__				,       n_embd=37		)



       def        _snake_case       (      self		)->str:




         '''simple docstring'''





         super().tearDown()
         # clean-up as much as possible GPU memory occupied by PyTorch
         gc.collect()
         torch.cuda.empty_cache()



       def        _snake_case       (      self		)->Tuple:




         '''simple docstring'''





         self.config_tester.run_common_tests()



       def        _snake_case       (      self		)->Optional[Any]:




         '''simple docstring'''





         A_      :				str		     =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_ctrl_model(*SCREAMING_SNAKE_CASE__		)



       def        _snake_case       (      self		)->Dict:




         '''simple docstring'''





         A_      :				List[Any]		     =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_lm_head_model(*SCREAMING_SNAKE_CASE__		)



       @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''		)
       def        _snake_case       (      self		)->int:




         '''simple docstring'''





         pass



       @slow
       def        _snake_case       (      self		)->List[str]:




         '''simple docstring'''





         for model_name in CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           A_      :				Optional[Any]		     =					CTRLModel.from_pretrained(SCREAMING_SNAKE_CASE__		)
           self.assertIsNotNone(SCREAMING_SNAKE_CASE__		)







       @unittest.skip('''The model doesn\'t support left padding'''		)  # and it's not used enough to be worth fixing :)
       def        _snake_case       (      self		)->Optional[int]:




         '''simple docstring'''





         pass






@require_torch
class   _lowerCamelCase (					unittest.TestCase			):




       """simple docstring"""



       def        _snake_case       (      self		)->Optional[int]:




         '''simple docstring'''





         super().tearDown()
         # clean-up as much as possible GPU memory occupied by PyTorch
         gc.collect()
         torch.cuda.empty_cache()







       @slow
       def        _snake_case       (      self		)->List[str]:




         '''simple docstring'''





         A_      :				Tuple		     =					CTRLLMHeadModel.from_pretrained('''ctrl'''		)
         model.to(SCREAMING_SNAKE_CASE__		)
         A_      :				Optional[int]		     =					torch.tensor(
             [[1_1859, 0, 1611, 8]]				,       dtype=torch.long				,       device=SCREAMING_SNAKE_CASE__		)  # Legal the president is
         A_      :				Optional[Any]		     =					[
             1_1859,
             0,
             1611,
             8,
             5,
             150,
             2_6449,
             2,
             19,
             348,
             469,
             3,
             2595,
             48,
             2_0740,
             24_6533,
             24_6533,
             19,
             30,
             5,
         ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

         A_      :				Dict		     =					model.generate(SCREAMING_SNAKE_CASE__				,       do_sample=SCREAMING_SNAKE_CASE__		)
         self.assertListEqual(output_ids[0].tolist()				,       SCREAMING_SNAKE_CASE__		)





'''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			:      Optional[Any]					      =				16
_A			:      Union[str, Any]					      =				32







def      UpperCamelCase_						(					snake_case_       :       List[str]      )					->     str:






 '''simple docstring'''

 return int(x / 2**20      )







class 				_lowercase     :


 '''simple docstring'''






 def __enter__(      self   :	List[Any]							)  ->    int:
  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   :	Tuple      ,       *SCREAMING_SNAKE_CASE__   :	List[Any]							)  ->    Optional[Any]:
  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      UpperCamelCase_						(					snake_case_       :       Accelerator						,		snake_case_       :       int = 16						,		snake_case_       :       str = "bert-base-cased"						,		snake_case_       :       int = 3_20						,		snake_case_       :       int = 1_60						,		)					->     Optional[int]:






 '''simple docstring'''

 __lowerCAmelCase											=						AutoTokenizer.from_pretrained(snake_case_      )
 __lowerCAmelCase											=						load_dataset(
     """glue"""						,		"""mrpc"""						,		split={"""train""": f"""train[:{n_train}]""", """validation""": f"""validation[:{n_val}]"""}      )

 def tokenize_function(snake_case_       :       List[Any]      ):
  # max_length=None => use the model max length (it's actually the default)
  __lowerCAmelCase											=						tokenizer(examples["""sentence1"""]						,		examples["""sentence2"""]						,		truncation=snake_case_						,		max_length=snake_case_      )
  return outputs

 # Apply the method we just defined to all the examples in all the splits of the dataset
 __lowerCAmelCase											=						datasets.map(
     snake_case_						,		batched=snake_case_						,		remove_columns=["""idx""", """sentence1""", """sentence2"""]						,		load_from_cache_file=snake_case_      )

 # 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(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(snake_case_						,		padding="""max_length"""						,		max_length=1_28						,		return_tensors="""pt"""      )
  return tokenizer.pad(snake_case_						,		padding="""longest"""						,		return_tensors="""pt"""      )

 # Instantiate dataloaders.
 __lowerCAmelCase											=						DataLoader(
     tokenized_datasets["""train"""]						,		shuffle=snake_case_						,		collate_fn=snake_case_						,		batch_size=snake_case_      )
 __lowerCAmelCase											=						DataLoader(
     tokenized_datasets["""validation"""]						,		shuffle=snake_case_						,		collate_fn=snake_case_						,		batch_size=snake_case_      )

 return train_dataloader, eval_dataloader







def      UpperCamelCase_						(					snake_case_       :       List[Any]						,		snake_case_       :       Tuple      )					->     Optional[int]:






 '''simple docstring'''

 __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(snake_case_      )
 __lowerCAmelCase							,       __lowerCAmelCase											=						get_dataloaders(snake_case_						,		snake_case_						,		snake_case_						,		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(snake_case_						,		return_dict=snake_case_      )

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

 if accelerator.state.deepspeed_plugin is not None:
  __lowerCAmelCase											=						accelerator.state.deepspeed_plugin.deepspeed_config[
      """gradient_accumulation_steps"""
  ]
 else:
  __lowerCAmelCase											=						1
 __lowerCAmelCase											=						(len(snake_case_      ) * 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=snake_case_						,		num_warmup_steps=0						,		num_training_steps=snake_case_						,		)
 else:
  __lowerCAmelCase											=						DummyScheduler(snake_case_						,		total_num_steps=snake_case_						,		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(
     snake_case_						,		snake_case_						,		snake_case_						,		snake_case_						,		snake_case_      )

 # 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(snake_case_						,		snake_case_      ):
  with TorchTracemalloc() as tracemalloc:
   model.train()
   for step, batch in enumerate(snake_case_      ):
    __lowerCAmelCase											=						model(**snake_case_      )
    __lowerCAmelCase											=						outputs.loss
    __lowerCAmelCase											=						loss / gradient_accumulation_steps
    accelerator.backward(snake_case_      )
    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(snake_case_						,		snake_case_      )







def      UpperCamelCase_						(					)					->     Any:






 '''simple docstring'''

 __lowerCAmelCase											=						argparse.ArgumentParser(description="""Simple example of training script tracking peak GPU memory usage."""      )
 parser.add_argument(
     """--model_name_or_path"""						,		type=snake_case_						,		default="""bert-base-cased"""						,		help="""Path to pretrained model or model identifier from huggingface.co/models."""						,		required=snake_case_						,		)
 parser.add_argument(
     """--output_dir"""						,		type=snake_case_						,		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=snake_case_						,		default=snake_case_						,		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=snake_case_						,		default=3_20						,		help="""Number of training examples to use."""						,		)
 parser.add_argument(
     """--n_val"""						,		type=snake_case_						,		default=1_60						,		help="""Number of validation examples to use."""						,		)
 parser.add_argument(
     """--num_epochs"""						,		type=snake_case_						,		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(snake_case_						,		snake_case_      )


if __name__ == "__main__":
 main()






