




import argparse
import os

# New Code #
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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to combine both the gradient accumulation
# and automatic batch size finder utilities of Accelerate to perfrom
# automatic gradient accumulation
#
# 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)
#
# 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
#
########################################################################

__snake_case			      =    32





def   lowerCAmelCase_ (				__lowerCAmelCase		,					__lowerCAmelCase = 16						)->  Dict:






      '''simple docstring'''





      UpperCAmelCase    : int      =AutoTokenizer.from_pretrained('''bert-base-cased'''						)
      UpperCAmelCase    : Optional[Any]      =load_dataset('''glue'''		,					'''mrpc'''						)

      def tokenize_function(__lowerCAmelCase						):
            # max_length=None => use the model max length (it's actually the default)
            UpperCAmelCase    : List[Any]      =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
      # starting with the main process first:
      with accelerator.main_process_first():
            UpperCAmelCase    : Optional[int]      =datasets.map(
                __lowerCAmelCase		,					batched=__lowerCAmelCase		,					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
      UpperCAmelCase    : List[Any]      =tokenized_datasets.rename_column('''label'''		,					'''labels'''						)

      def collate_fn(__lowerCAmelCase						):
            # When using mixed precision we want round multiples of 8/16
            if accelerator.mixed_precision == "fp8":
                  UpperCAmelCase    : Dict      =16
            elif accelerator.mixed_precision != "no":
                  UpperCAmelCase    : str      =8
            else:
                  UpperCAmelCase    : int      =None

            return tokenizer.pad(
                __lowerCAmelCase		,					padding='''longest'''		,					pad_to_multiple_of=__lowerCAmelCase		,					return_tensors='''pt'''		,					)

      # Instantiate dataloaders.
      UpperCAmelCase    : Any      =DataLoader(
          tokenized_datasets['''train''']		,					shuffle=__lowerCAmelCase		,					collate_fn=__lowerCAmelCase		,					batch_size=__lowerCAmelCase						)
      UpperCAmelCase    : Tuple      =DataLoader(
          tokenized_datasets['''validation''']		,					shuffle=__lowerCAmelCase		,					collate_fn=__lowerCAmelCase		,					batch_size=__lowerCAmelCase						)

      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

     __snake_case			      =    mocked_dataloaders  # noqa: F811





def   lowerCAmelCase_ (				__lowerCAmelCase		,					__lowerCAmelCase						)->  List[str]:






      '''simple docstring'''





      if os.environ.get('''TESTING_MOCKED_DATALOADERS'''		,					__lowerCAmelCase						) == "1":
            UpperCAmelCase    : List[Any]      =2
      # Initialize accelerator
      UpperCAmelCase    : Tuple      =Accelerator(cpu=args.cpu		,					mixed_precision=args.mixed_precision						)
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      UpperCAmelCase    : Dict      =config["lr"]
      UpperCAmelCase    : List[str]      =int(config['''num_epochs''']						)
      UpperCAmelCase    : str      =int(config['''seed''']						)
      UpperCAmelCase    : Dict      =int(config['''batch_size''']						)

      UpperCAmelCase    : Optional[int]      =evaluate.load('''glue'''		,					'''mrpc'''						)

      # New Code #
      # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
      # to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
      # half each time. From this, we can calculate the number of gradient accumulation steps needed
      # and modify the Accelerator object as a result
      @find_executable_batch_size(starting_batch_size=int(__lowerCAmelCase						)						)
      def inner_training_loop(__lowerCAmelCase						):
            # Since we need to modify the outside accelerator object, we need to bring it
            # to the local scope
            nonlocal accelerator

            # We can calculate the number of gradient accumulation steps based on the current
            # batch size vs the starting batch size
            UpperCAmelCase    : Optional[Any]      =observed_batch_size // batch_size

            # And then set it in the Accelerator directly:
            UpperCAmelCase    : List[str]      =num_gradient_accumulation_steps

            # Next we need to free all of the stored model references in the Accelerator each time
            accelerator.free_memory()

            # And set the seed so our results are reproducable each reset
            set_seed(__lowerCAmelCase						)

            # Instantiate the model (we build the model here so that the seed also control new weights initialization)
            UpperCAmelCase    : str      =AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''		,					return_dict=__lowerCAmelCase						)

            # 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).
            UpperCAmelCase    : List[Any]      =model.to(accelerator.device						)

            # Instantiate optimizer
            UpperCAmelCase    : str      =AdamW(params=model.parameters()		,					lr=__lowerCAmelCase						)
            UpperCAmelCase    : List[str]      =get_dataloaders(__lowerCAmelCase		,					__lowerCAmelCase						)

            # Instantiate scheduler
            UpperCAmelCase    : List[Any]      =get_linear_schedule_with_warmup(
                optimizer=__lowerCAmelCase		,					num_warmup_steps=1_00		,					num_training_steps=(len(__lowerCAmelCase						) * num_epochs)		,					)

            # 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.
            UpperCAmelCase    : Any      =accelerator.prepare(
                __lowerCAmelCase		,					__lowerCAmelCase		,					__lowerCAmelCase		,					__lowerCAmelCase		,					__lowerCAmelCase						)

            # Now we train the model
            for epoch in range(__lowerCAmelCase						):
                  model.train()
                  for step, batch in enumerate(__lowerCAmelCase						):
                        # And perform gradient accumulation
                        with accelerator.accumulate(__lowerCAmelCase						):
                              # We could avoid this line since we set the accelerator with `device_placement=True`.
                              batch.to(accelerator.device						)
                              UpperCAmelCase    : Any      =model(**__lowerCAmelCase						)
                              UpperCAmelCase    : List[Any]      =outputs.loss
                              accelerator.backward(__lowerCAmelCase						)
                              optimizer.step()
                              lr_scheduler.step()
                              optimizer.zero_grad()

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

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

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
      inner_training_loop()





def   lowerCAmelCase_ (				)->  Optional[int]:






      '''simple docstring'''





      UpperCAmelCase    : Tuple      =argparse.ArgumentParser(description='''Simple example of training script.'''						)
      parser.add_argument(
          '''--mixed_precision'''		,					type=__lowerCAmelCase		,					default=__lowerCAmelCase		,					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.'''						)
      UpperCAmelCase    : int      =parser.parse_args()
      # New Code #
      # We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
      UpperCAmelCase    : Dict      ={"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 2_56}
      training_function(__lowerCAmelCase		,					__lowerCAmelCase						)


if __name__ == "__main__":
     main()

from collections import defaultdict





def   lowerCAmelCase_ (				__lowerCAmelCase = 5						)->  int:






      '''simple docstring'''





      UpperCAmelCase    : List[str]      =defaultdict(__lowerCAmelCase						)
      UpperCAmelCase    : Tuple      =0

      while True:
            UpperCAmelCase    : Any      =get_digits(__lowerCAmelCase						)
            freqs[digits].append(__lowerCAmelCase						)

            if len(freqs[digits]						) == max_base:
                  UpperCAmelCase    : Optional[int]      =freqs[digits][0] ** 3
                  return base

            num += 1





def   lowerCAmelCase_ (				__lowerCAmelCase						)->  str:






      '''simple docstring'''





      return "".join(sorted(str(num**3						)						)						)


if __name__ == "__main__":
     print(f'{solution() = }')
