




import argparse
import os

import evaluate
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


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing the checkpointing capability,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################

lowercase__      :			Union[str, Any]           =					16
lowercase__      :			Union[str, Any]           =					32

def 					SCREAMING_SNAKE_CASE							(   __UpperCamelCase		,					__UpperCamelCase = 16)			->     int:
     a      =      AutoTokenizer.from_pretrained("bert-base-cased")
     a      =      load_dataset("glue"		,					"mrpc")

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

     # Apply the method we just defined to all the examples in all the splits of the dataset
     # starting with the main process first:
     with accelerator.main_process_first():
          a      =      datasets.map(
              __UpperCamelCase		,					batched=__UpperCamelCase		,					remove_columns=["idx", "sentence1", "sentence2"]		,					)

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     a      =      tokenized_datasets.rename_column("label"		,					"labels")

     def collate_fn(__UpperCamelCase):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          a      =      1_28 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               a      =      16
          elif accelerator.mixed_precision != "no":
               a      =      8
          else:
               a      =      None

          return tokenizer.pad(
              __UpperCamelCase		,					padding="longest"		,					max_length=__UpperCamelCase		,					pad_to_multiple_of=__UpperCamelCase		,					return_tensors="pt"		,					)

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

     return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
   from accelerate.test_utils.training import mocked_dataloaders

   lowercase__      :			Any           =					mocked_dataloaders  # noqa: F811

def 					SCREAMING_SNAKE_CASE							(   __UpperCamelCase		,					__UpperCamelCase)			->     List[str]:
     # For testing only
     if os.environ.get("TESTING_MOCKED_DATALOADERS"		,					__UpperCamelCase) == "1":
          a      =      2
     # Initialize accelerator
     a      =      Accelerator(cpu=args.cpu		,					mixed_precision=args.mixed_precision)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     a      =      config["lr"]
     a      =      int(config["num_epochs"])
     a      =      int(config["seed"])
     a      =      int(config["batch_size"])

     # New Code #
     # Parse out whether we are saving every epoch or after a certain number of batches
     if hasattr(args.checkpointing_steps		,					"isdigit"):
          if args.checkpointing_steps == "epoch":
               a      =      args.checkpointing_steps
          elif args.checkpointing_steps.isdigit():
               a      =      int(args.checkpointing_steps)
          else:
               raise ValueError(
                   f'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.''')
     else:
          a      =      None

     set_seed(__UpperCamelCase)

     a  ,      a      =      get_dataloaders(__UpperCamelCase		,					__UpperCamelCase)
     a      =      evaluate.load("glue"		,					"mrpc")

     # If the batch size is too big we use gradient accumulation
     a      =      1
     if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
          a      =      batch_size // MAX_GPU_BATCH_SIZE
          a      =      MAX_GPU_BATCH_SIZE

     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     a      =      AutoModelForSequenceClassification.from_pretrained("bert-base-cased"		,					return_dict=__UpperCamelCase)

     # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
     # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
     # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
     a      =      model.to(accelerator.device)

     # Instantiate optimizer
     a      =      AdamW(params=model.parameters()		,					lr=__UpperCamelCase)

     # Instantiate scheduler
     a      =      get_linear_schedule_with_warmup(
         optimizer=__UpperCamelCase		,					num_warmup_steps=1_00		,					num_training_steps=(len(__UpperCamelCase) * num_epochs) // gradient_accumulation_steps		,					)

     # 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.
     a  ,      a  ,      a  ,      a  ,      a      =      accelerator.prepare(
         __UpperCamelCase		,					__UpperCamelCase		,					__UpperCamelCase		,					__UpperCamelCase		,					__UpperCamelCase)

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

     # We need to load the checkpoint back in before training here with `load_state`
     # The total number of epochs is adjusted based on where the state is being loaded from,
     # as we assume continuation of the same training script
     if args.resume_from_checkpoint:
          if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
               accelerator.print(f'''Resumed from checkpoint: {args.resume_from_checkpoint}''')
               accelerator.load_state(args.resume_from_checkpoint)
               a      =      os.path.basename(args.resume_from_checkpoint)
          else:
               # Get the most recent checkpoint
               a      =      [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
               dirs.sort(key=os.path.getctime)
               a      =      dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
          # Extract `epoch_{i}` or `step_{i}`
          a      =      os.path.splitext(__UpperCamelCase)[0]

          if "epoch" in training_difference:
               a      =      int(training_difference.replace("epoch_"		,					"")) + 1
               a      =      None
          else:
               a      =      int(training_difference.replace("step_"		,					""))
               a      =      resume_step // len(__UpperCamelCase)
               resume_step -= starting_epoch * len(__UpperCamelCase)

    # Now we train the model
     for epoch in range(__UpperCamelCase		,					__UpperCamelCase):
          model.train()
          # New Code #
          if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
               # We need to skip steps until we reach the resumed step
               a      =      accelerator.skip_first_batches(__UpperCamelCase		,					__UpperCamelCase)
               overall_step += resume_step
          else:
               # After the first iteration though, we need to go back to the original dataloader
               a      =      train_dataloader
          for step, batch in enumerate(__UpperCamelCase):
               # We could avoid this line since we set the accelerator with `device_placement=True`.
               batch.to(accelerator.device)
               a      =      model(**__UpperCamelCase)
               a      =      outputs.loss
               a      =      loss / gradient_accumulation_steps
               accelerator.backward(__UpperCamelCase)
               if step % gradient_accumulation_steps == 0:
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
               # New Code #
               overall_step += 1

               # New Code #
               # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
               # These are saved to folders named `step_{overall_step}`
               # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
               # If mixed precision was used, will also save a "scalar.bin" file
               if isinstance(__UpperCamelCase		,					__UpperCamelCase):
                    a      =      f'''step_{overall_step}'''
                    if overall_step % checkpointing_steps == 0:
                         if args.output_dir is not None:
                              a      =      os.path.join(args.output_dir		,					__UpperCamelCase)
                         accelerator.save_state(__UpperCamelCase)

          model.eval()
          for step, batch in enumerate(__UpperCamelCase):
               # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
               batch.to(accelerator.device)
               with torch.no_grad():
                    a      =      model(**__UpperCamelCase)
               a      =      outputs.logits.argmax(dim=-1)
               a  ,      a      =      accelerator.gather_for_metrics((predictions, batch["labels"]))
               metric.add_batch(
                   predictions=__UpperCamelCase		,					references=__UpperCamelCase		,					)

          a      =      metric.compute()
          # Use accelerator.print to print only on the main process.
          accelerator.print(f'''epoch {epoch}:'''		,					__UpperCamelCase)

          # New Code #
          # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
          # These are saved to folders named `epoch_{epoch}`
          # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
          # If mixed precision was used, will also save a "scalar.bin" file
          if checkpointing_steps == "epoch":
               a      =      f'''epoch_{epoch}'''
               if args.output_dir is not None:
                    a      =      os.path.join(args.output_dir		,					__UpperCamelCase)
               accelerator.save_state(__UpperCamelCase)

def 					SCREAMING_SNAKE_CASE							(   )			->     str:
     a      =      argparse.ArgumentParser(description="Simple example of training script.")
     parser.add_argument(
         "--mixed_precision"		,					type=__UpperCamelCase		,					default=__UpperCamelCase		,					choices=["no", "fp16", "bf16", "fp8"]		,					help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU."		,					)
     parser.add_argument("--cpu"		,					action="store_true"		,					help="If passed, will train on the CPU.")
     parser.add_argument(
         "--checkpointing_steps"		,					type=__UpperCamelCase		,					default=__UpperCamelCase		,					help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch."		,					)
     parser.add_argument(
         "--output_dir"		,					type=__UpperCamelCase		,					default="."		,					help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory."		,					)
     parser.add_argument(
         "--resume_from_checkpoint"		,					type=__UpperCamelCase		,					default=__UpperCamelCase		,					help="If the training should continue from a checkpoint folder."		,					)
     a      =      parser.parse_args()
     a      =      {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(__UpperCamelCase		,					__UpperCamelCase)


if __name__ == "__main__":
   main()







from __future__ import annotations

import unittest

from transformers import CTRLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers.models.ctrl.modeling_tf_ctrl import (
       TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFCTRLForSequenceClassification,
       TFCTRLLMHeadModel,
       TFCTRLModel,
   )




class 			a__       (					UpperCamelCase__      ):







     def __init__(	self  ,		A  ,		)  ->    List[Any]:



          '''simple docstring'''



          a      =      parent
          a      =      13
          a      =      7
          a      =      True
          a      =      True
          a      =      True
          a      =      True
          a      =      True
          a      =      99
          a      =      32
          a      =      2
          a      =      4
          a      =      37
          a      =      "gelu"
          a      =      0.1
          a      =      0.1
          a      =      512
          a      =      16
          a      =      2
          a      =      0.0_2
          a      =      3
          a      =      4
          a      =      None
          a      =      self.vocab_size - 1







     def 			lowerCAmelCase_				(	self )  ->    Union[str, Any]:



          '''simple docstring'''



          a      =      ids_tensor([self.batch_size, self.seq_length]  ,		self.vocab_size )

          a      =      None
          if self.use_input_mask:
               a      =      random_attention_mask([self.batch_size, self.seq_length] )

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

          a      =      None
          if self.use_mc_token_ids:
               a      =      ids_tensor([self.batch_size, self.num_choices]  ,		self.seq_length )

          a      =      None
          a      =      None
          a      =      None
          if self.use_labels:
               a      =      ids_tensor([self.batch_size]  ,		self.type_sequence_label_size )
               a      =      ids_tensor([self.batch_size, self.seq_length]  ,		self.num_labels )
               a      =      ids_tensor([self.batch_size]  ,		self.num_choices )

          a      =      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  ,		)

          a      =      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 			lowerCAmelCase_				(	self  ,		A  ,		A  ,		A  ,		A  ,		A  ,		*A )  ->    Optional[int]:



          '''simple docstring'''



          a      =      TFCTRLModel(config=A )
          a      =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          a      =      model(A )

          a      =      [input_ids, None, input_mask]  # None is the input for 'past'
          a      =      model(A )

          a      =      model(A )

          self.parent.assertEqual(result.last_hidden_state.shape  ,		(self.batch_size, self.seq_length, self.hidden_size) )







     def 			lowerCAmelCase_				(	self  ,		A  ,		A  ,		A  ,		A  ,		A  ,		*A )  ->    Tuple:



          '''simple docstring'''



          a      =      TFCTRLLMHeadModel(config=A )
          a      =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          a      =      model(A )
          self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.seq_length, self.vocab_size) )







     def 			lowerCAmelCase_				(	self  ,		A  ,		A  ,		A  ,		A  ,		A  ,		*A )  ->    List[Any]:



          '''simple docstring'''



          a      =      self.num_labels
          a      =      ids_tensor([self.batch_size]  ,		self.type_sequence_label_size )
          a      =      {
              "input_ids": input_ids,
              "token_type_ids": token_type_ids,
              "labels": sequence_labels,
          }
          a      =      TFCTRLForSequenceClassification(A )
          a      =      model(A )
          self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.num_labels) )







     def 			lowerCAmelCase_				(	self )  ->    Union[str, Any]:



          '''simple docstring'''



          a      =      self.prepare_config_and_inputs()

          (
              (
              a
          )  ,      (
              a
          )  ,      (
              a
          )  ,      (
              a
          )  ,      (
              a
          )  ,      (
              a
          )  ,      (
              a
          )  ,      (
              a
          )  ,      (
              a
          )  ,      
          )      =      config_and_inputs

          a      =      {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
          return config, inputs_dict







@require_tf
class 			a__       (					UpperCamelCase__   ,					UpperCamelCase__   ,					unittest.TestCase      ):
     a		:			Any							=	(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
     a		:			Optional[Any]							=	(TFCTRLLMHeadModel,) if is_tf_available() else ()
     a		:			List[Any]							=	(
         {
             """feature-extraction""": TFCTRLModel,
             """text-classification""": TFCTRLForSequenceClassification,
             """text-generation""": TFCTRLLMHeadModel,
             """zero-shot""": TFCTRLForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     a		:			Union[str, Any]							=	False
     a		:			str							=	False







     def 			lowerCAmelCase_				(	self  ,		A  ,		A  ,		A  ,		A  ,		A )  ->    List[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 			lowerCAmelCase_				(	self )  ->    Optional[int]:



          '''simple docstring'''



          a      =      TFCTRLModelTester(self )
          a      =      ConfigTester(self  ,		config_class=A  ,		n_embd=37 )







     def 			lowerCAmelCase_				(	self )  ->    Tuple:



          '''simple docstring'''



          self.config_tester.run_common_tests()







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



          '''simple docstring'''



          a      =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_ctrl_model(*A )







     def 			lowerCAmelCase_				(	self )  ->    List[str]:



          '''simple docstring'''



          a      =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_ctrl_lm_head(*A )







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



          '''simple docstring'''



          a      =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_ctrl_for_sequence_classification(*A )







     def 			lowerCAmelCase_				(	self )  ->    Tuple:



          '''simple docstring'''



          a  ,      a      =      self.model_tester.prepare_config_and_inputs_for_common()
          a      =      [TFCTRLLMHeadModel]
          a      =      [TFCTRLForSequenceClassification]

          for model_class in self.all_model_classes:
               a      =      model_class(A )
               model.build()  # may be needed for the get_bias() call below
               assert isinstance(model.get_input_embeddings()  ,		tf.keras.layers.Layer )

               if model_class in list_lm_models:
                    a      =      model.get_output_embeddings()
                    assert isinstance(A  ,		tf.keras.layers.Layer )
                    a      =      model.get_bias()
                    assert isinstance(A  ,		A )
                    for k, v in name.items():
                         assert isinstance(A  ,		tf.Variable )
               elif model_class in list_other_models_with_output_ebd:
                    a      =      model.get_output_embeddings()
                    assert isinstance(A  ,		tf.keras.layers.Layer )
                    a      =      model.get_bias()
                    assert name is None
               else:
                    a      =      model.get_output_embeddings()
                    assert x is None
                    a      =      model.get_bias()
                    assert name is None







     @slow
     def 			lowerCAmelCase_				(	self )  ->    Tuple:



          '''simple docstring'''



          for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               a      =      TFCTRLModel.from_pretrained(A )
               self.assertIsNotNone(A )







@require_tf
class 			a__       (					unittest.TestCase      ):







     @slow
     def 			lowerCAmelCase_				(	self )  ->    Any:



          '''simple docstring'''



          a      =      TFCTRLLMHeadModel.from_pretrained("ctrl" )
          a      =      tf.convert_to_tensor([[11859, 0, 1611, 8]]  ,		dtype=tf.intaa )  # Legal the president is
          a      =      [
              11859,
              0,
              1611,
              8,
              5,
              150,
              26449,
              2,
              19,
              348,
              469,
              3,
              2595,
              48,
              20740,
              246533,
              246533,
              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      =      model.generate(A  ,		do_sample=A )
          self.assertListEqual(output_ids[0].numpy().tolist()  ,		A )

