



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
#
########################################################################

_A				:							Dict							= 32


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_ = 1_6  )							->	List[Any]:
    SCREAMING_SNAKE_CASE__       							= AutoTokenizer.from_pretrained('''bert-base-cased'''  )
    SCREAMING_SNAKE_CASE__       							= load_dataset('''glue'''   ,       '''mrpc'''  )

    def tokenize_function(lowerCAmelCase_  ):
        # max_length=None => use the model max length (it's actually the default)
        SCREAMING_SNAKE_CASE__       							= 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():
        SCREAMING_SNAKE_CASE__       							= 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
    SCREAMING_SNAKE_CASE__       							= 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":
            SCREAMING_SNAKE_CASE__       							= 1_6
        elif accelerator.mixed_precision != "no":
            SCREAMING_SNAKE_CASE__       							= 8
        else:
            SCREAMING_SNAKE_CASE__       							= None

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

    # Instantiate dataloaders.
    SCREAMING_SNAKE_CASE__       							= DataLoader(
        tokenized_datasets['''train''']   ,       shuffle=lowerCAmelCase_   ,       collate_fn=lowerCAmelCase_   ,       batch_size=lowerCAmelCase_  )
    SCREAMING_SNAKE_CASE__       							= 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

  _A				:							Union[str, Any]							= mocked_dataloaders  # noqa: F811


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

    SCREAMING_SNAKE_CASE__       							= 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
        SCREAMING_SNAKE_CASE__       							= observed_batch_size // batch_size

        # And then set it in the Accelerator directly:
        SCREAMING_SNAKE_CASE__       							= 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)
        SCREAMING_SNAKE_CASE__       							= 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).
        SCREAMING_SNAKE_CASE__       							= model.to(accelerator.device  )

        # Instantiate optimizer
        SCREAMING_SNAKE_CASE__       							= AdamW(params=model.parameters()   ,       lr=lowerCAmelCase_  )
        SCREAMING_SNAKE_CASE__       							= get_dataloaders(lowerCAmelCase_   ,       lowerCAmelCase_  )

        # Instantiate scheduler
        SCREAMING_SNAKE_CASE__       							= get_linear_schedule_with_warmup(
            optimizer=lowerCAmelCase_   ,       num_warmup_steps=1_0_0   ,       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.
        SCREAMING_SNAKE_CASE__       							= 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  )
                    SCREAMING_SNAKE_CASE__       							= model(**lowerCAmelCase_  )
                    SCREAMING_SNAKE_CASE__       							= 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():
                    SCREAMING_SNAKE_CASE__       							= model(**lowerCAmelCase_  )
                SCREAMING_SNAKE_CASE__       							= outputs.logits.argmax(dim=-1  )
                SCREAMING_SNAKE_CASE__       							= accelerator.gather_for_metrics((predictions, batch['''labels'''])  )
                metric.add_batch(
                    predictions=lowerCAmelCase_   ,       references=lowerCAmelCase_   ,       )

            SCREAMING_SNAKE_CASE__       							= 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        __snake_case       (						)							->	Tuple:
    SCREAMING_SNAKE_CASE__       							= 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.'''  )
    SCREAMING_SNAKE_CASE__       							= 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
    SCREAMING_SNAKE_CASE__       							= {"lr": 2e-5, "num_epochs": 3, "seed": 4_2, "batch_size": 2_5_6}
    training_function(lowerCAmelCase_   ,       lowerCAmelCase_  )


if __name__ == "__main__":
  main()







from __future__ import annotations

from collections.abc import Generator


def        __snake_case       (						lowerCAmelCase_  )							->	Generator[int, None, None]:
    if not isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ) or n < 1:
        raise Exception('''Sequence only defined for positive integers'''  )

    yield n
    while n != 1:
        if n % 2 == 0:
            n //= 2
        else:
            SCREAMING_SNAKE_CASE__       							= 3 * n + 1
        yield n


def        __snake_case       (						)							->	str:
    SCREAMING_SNAKE_CASE__       							= 4_3
    SCREAMING_SNAKE_CASE__       							= tuple(collatz_sequence(lowerCAmelCase_  )  )
    print(lowerCAmelCase_  )
    print(f'''Collatz sequence from {n} took {len(lowerCAmelCase_  )} steps.'''  )


if __name__ == "__main__":
  main()
