



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, DistributedType
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, 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)
#
# 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
#
########################################################################


UpperCamelCase_	   =			1_6
UpperCamelCase_	   =			3_2




def 		SCREAMING_SNAKE_CASE     (     snake_case__				,   snake_case__ = 16      )   ->	Optional[int]:
    __UpperCAmelCase			 =AutoTokenizer.from_pretrained('''bert-base-cased'''      )
    __UpperCAmelCase			 =load_dataset('''glue'''				,   '''mrpc'''      )

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

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

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

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

       UpperCamelCase_	   =			mocked_dataloaders  # noqa: F811




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

    __UpperCAmelCase			 =evaluate.load('''glue'''				,   '''mrpc'''      )

    # New Code #
    # We now can define an inner training loop function. It should take a batch size as the only parameter,
    # and build the dataloaders in there.
    # It also gets our decorator
    @find_executable_batch_size(starting_batch_size=__lowerCAmelCase      )
    def inner_training_loop(snake_case__      ):
        # And now just move everything below under this function
        # We need to bring in the Accelerator object from earlier
        nonlocal accelerator
        # And reset all of its attributes that could hold onto any memory:
        accelerator.free_memory()

        # Then we can declare the model, optimizer, and everything else:
        set_seed(__lowerCAmelCase      )

        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
        __UpperCAmelCase			 =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			 =model.to(accelerator.device      )

        # Instantiate optimizer
        __UpperCAmelCase			 =AdamW(params=model.parameters()				,   lr=__lowerCAmelCase      )
        __UpperCAmelCase			 =get_dataloaders(__lowerCAmelCase				,   __lowerCAmelCase      )

        # Instantiate scheduler
        __UpperCAmelCase			 =get_linear_schedule_with_warmup(
            optimizer=__lowerCAmelCase				,   num_warmup_steps=100				,   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			 =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      ):
                # We could avoid this line since we set the accelerator with `device_placement=True`.
                batch.to(accelerator.device      )
                __UpperCAmelCase			 =model(**__lowerCAmelCase      )
                __UpperCAmelCase			 =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			 =model(**__lowerCAmelCase      )
                __UpperCAmelCase			 =outputs.logits.argmax(dim=-1      )
                __UpperCAmelCase			 =accelerator.gather_for_metrics((predictions, batch['''labels'''])      )
                metric.add_batch(
                    predictions=__lowerCAmelCase				,   references=__lowerCAmelCase				,   )

            __UpperCAmelCase			 =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 		SCREAMING_SNAKE_CASE     (     )   ->	List[str]:
    __UpperCAmelCase			 =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			 =parser.parse_args()
    __UpperCAmelCase			 ={"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
    training_function(__lowerCAmelCase				,   __lowerCAmelCase      )


if __name__ == "__main__":
       main()











import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCamelCase    							=      logging.get_logger(__name__)


def 				__lowerCamelCase     (							__lowerCAmelCase							:					Tuple							)     ->  List[str]:
   __UpperCamelCase	:			List[str]					=					R"""\w+[.]\d+"""
   __UpperCamelCase	:			Optional[int]					=					re.findall(__lowerCAmelCase					,							__lowerCAmelCase							)
   for pat in pats:
      __UpperCamelCase	:			str					=					key.replace(__lowerCAmelCase					,							"""_""".join(pat.split("""."""							)							)							)
   return key


def 				__lowerCamelCase     (							__lowerCAmelCase							:					int					,							__lowerCAmelCase							:					Optional[int]					,							__lowerCAmelCase							:					Union[str, Any]							)     ->  str:
   __UpperCamelCase	:			Dict					=					pt_tuple_key[:-1] + ("""scale""",)
   if (
       any("""norm""" in str_ for str_ in pt_tuple_key							)
       and (pt_tuple_key[-1] == "bias")
       and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
       and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
   ):
      __UpperCamelCase	:			List[Any]					=					pt_tuple_key[:-1] + ("""scale""",)
      return renamed_pt_tuple_key, pt_tensor
   elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
      __UpperCamelCase	:			Dict					=					pt_tuple_key[:-1] + ("""scale""",)
      return renamed_pt_tuple_key, pt_tensor

   # embedding
   if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
      __UpperCamelCase	:			Any					=					pt_tuple_key[:-1] + ("""embedding""",)
      return renamed_pt_tuple_key, pt_tensor

   # conv layer
   __UpperCamelCase	:			List[Any]					=					pt_tuple_key[:-1] + ("""kernel""",)
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
      __UpperCamelCase	:			Dict					=					pt_tensor.transpose(2					,							3					,							1					,							0							)
      return renamed_pt_tuple_key, pt_tensor

   # linear layer
   __UpperCamelCase	:			Union[str, Any]					=					pt_tuple_key[:-1] + ("""kernel""",)
   if pt_tuple_key[-1] == "weight":
      __UpperCamelCase	:			int					=					pt_tensor.T
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm weight
   __UpperCamelCase	:			List[Any]					=					pt_tuple_key[:-1] + ("""weight""",)
   if pt_tuple_key[-1] == "gamma":
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm bias
   __UpperCamelCase	:			Tuple					=					pt_tuple_key[:-1] + ("""bias""",)
   if pt_tuple_key[-1] == "beta":
      return renamed_pt_tuple_key, pt_tensor

   return pt_tuple_key, pt_tensor



def 				__lowerCamelCase     (							__lowerCAmelCase							:					Dict					,							__lowerCAmelCase							:					List[Any]					,							__lowerCAmelCase							:					Optional[int]=42							)     ->  int:
   # Step 1: Convert pytorch tensor to numpy
   __UpperCamelCase	:			Dict					=					{k: v.numpy() for k, v in pt_state_dict.items()}

   # Step 2: Since the model is stateless, get random Flax params
   __UpperCamelCase	:			str					=					flax_model.init_weights(PRNGKey(__lowerCAmelCase							)							)

   __UpperCamelCase	:			Union[str, Any]					=					flatten_dict(__lowerCAmelCase							)
   __UpperCamelCase	:			List[str]					=					{}

   # Need to change some parameters name to match Flax names
   for pt_key, pt_tensor in pt_state_dict.items():
      __UpperCamelCase	:			str					=					rename_key(__lowerCAmelCase							)
      __UpperCamelCase	:			Any					=					tuple(renamed_pt_key.split("""."""							)							)

      # Correctly rename weight parameters
      __UpperCamelCase	,     __UpperCamelCase	:			Tuple					=					rename_key_and_reshape_tensor(__lowerCAmelCase					,							__lowerCAmelCase					,							__lowerCAmelCase							)

      if flax_key in random_flax_state_dict:
         if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
            raise ValueError(
                f'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                f'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'							)

        # also add unexpected weight so that warning is thrown
      __UpperCamelCase	:			Any					=					jnp.asarray(__lowerCAmelCase							)

   return unflatten_dict(__lowerCAmelCase							)



