


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

__UpperCamelCase					:	Any          =				32
def   snake_case_    (			__lowercase      ,							__lowercase = 1_6							):
   UpperCAmelCase_				:   List[str]							= AutoTokenizer.from_pretrained('''bert-base-cased'''							)
   UpperCAmelCase_				:   Tuple							= load_dataset('''glue'''      ,							'''mrpc'''							)

   def tokenize_function(__lowercase							):
      # max_length=None => use the model max length (it's actually the default)
      UpperCAmelCase_				:   Optional[int]							= tokenizer(examples['''sentence1''']      ,							examples['''sentence2''']      ,							truncation=__lowercase      ,							max_length=__lowercase							)
      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[Any]							= datasets.map(
          __lowercase      ,							batched=__lowercase      ,							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_				:   Tuple							= tokenized_datasets.rename_column('''label'''      ,							'''labels'''							)

   def collate_fn(__lowercase							):
      # When using mixed precision we want round multiples of 8/16
      if accelerator.mixed_precision == "fp8":
         UpperCAmelCase_				:   Tuple							= 1_6
      elif accelerator.mixed_precision != "no":
         UpperCAmelCase_				:   Optional[Any]							= 8
      else:
         UpperCAmelCase_				:   Tuple							= None

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

   # Instantiate dataloaders.
   UpperCAmelCase_				:   Optional[int]							= DataLoader(
       tokenized_datasets['''train''']      ,							shuffle=__lowercase      ,							collate_fn=__lowercase      ,							batch_size=__lowercase							)
   UpperCAmelCase_				:   str							= DataLoader(
       tokenized_datasets['''validation''']      ,							shuffle=__lowercase      ,							collate_fn=__lowercase      ,							batch_size=__lowercase							)

   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					:	str          =				mocked_dataloaders  # noqa: F811
def   snake_case_    (			__lowercase      ,							__lowercase							):
   # For testing only
   if os.environ.get('''TESTING_MOCKED_DATALOADERS'''      ,							__lowercase							) == "1":
      UpperCAmelCase_				:   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_				:   int							= config['''lr''']
   UpperCAmelCase_				:   str							= int(config['''num_epochs''']							)
   UpperCAmelCase_				:   Dict							= int(config['''seed''']							)
   UpperCAmelCase_				:   Optional[int]							= int(config['''batch_size''']							)

   UpperCAmelCase_				:   Optional[Any]							= 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(__lowercase							)							)
   def inner_training_loop(__lowercase							):
      # 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_				:   Tuple							= observed_batch_size // batch_size

      # And then set it in the Accelerator directly:
      UpperCAmelCase_				:   Any							= 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(__lowercase							)

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

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

      # Instantiate optimizer
      UpperCAmelCase_				:   Any							= AdamW(params=model.parameters()      ,							lr=__lowercase							)
      UpperCAmelCase_		,   UpperCAmelCase_				:   Optional[Any]							= get_dataloaders(__lowercase      ,							__lowercase							)

      # Instantiate scheduler
      UpperCAmelCase_				:   Dict							= get_linear_schedule_with_warmup(
          optimizer=__lowercase      ,							num_warmup_steps=1_0_0      ,							num_training_steps=(len(__lowercase							) * 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_		,   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_				:   str							= accelerator.prepare(
          __lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase							)

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

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

         UpperCAmelCase_				:   Dict							= metric.compute()
         # Use accelerator.print to print only on the main process.
         accelerator.print(F'''epoch {epoch}:'''      ,							__lowercase							)

    # 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_    (			):
   UpperCAmelCase_				:   Any							= argparse.ArgumentParser(description='''Simple example of training script.'''							)
   parser.add_argument(
       '''--mixed_precision'''      ,							type=__lowercase      ,							default=__lowercase      ,							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_				:   Any							= 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_				:   int							= {'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 4_2, '''batch_size''': 2_5_6}
   training_function(__lowercase      ,							__lowercase							)


if __name__ == "__main__":
       main()



import json
import os
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


if TYPE_CHECKING:
       from transformers.pipelines.conversational import Conversation

__UpperCamelCase					:	Dict          =				logging.get_logger(__name__)

__UpperCamelCase					:	List[str]          =				{
    'vocab_file': 'vocab.json',
    'merges_file': 'merges.txt',
}

__UpperCamelCase					:	Optional[int]          =				{
    'vocab_file': {
        'gpt2': 'https://huggingface.co/gpt2/resolve/main/vocab.json',
        'gpt2-medium': 'https://huggingface.co/gpt2-medium/resolve/main/vocab.json',
        'gpt2-large': 'https://huggingface.co/gpt2-large/resolve/main/vocab.json',
        'gpt2-xl': 'https://huggingface.co/gpt2-xl/resolve/main/vocab.json',
        'distilgpt2': 'https://huggingface.co/distilgpt2/resolve/main/vocab.json',
    },
    'merges_file': {
        'gpt2': 'https://huggingface.co/gpt2/resolve/main/merges.txt',
        'gpt2-medium': 'https://huggingface.co/gpt2-medium/resolve/main/merges.txt',
        'gpt2-large': 'https://huggingface.co/gpt2-large/resolve/main/merges.txt',
        'gpt2-xl': 'https://huggingface.co/gpt2-xl/resolve/main/merges.txt',
        'distilgpt2': 'https://huggingface.co/distilgpt2/resolve/main/merges.txt',
    },
}

__UpperCamelCase					:	Any          =				{
    'gpt2': 1024,
    'gpt2-medium': 1024,
    'gpt2-large': 1024,
    'gpt2-xl': 1024,
    'distilgpt2': 1024,
}
@lru_cache()
def   snake_case_    (			):
   UpperCAmelCase_				:   Union[str, Any]							= (
       list(range(ord('''!'''							)      ,							ord('''~'''							) + 1							)							) + list(range(ord('''¡'''							)      ,							ord('''¬'''							) + 1							)							) + list(range(ord('''®'''							)      ,							ord('''ÿ'''							) + 1							)							)
   )
   UpperCAmelCase_				:   Optional[int]							= bs[:]
   UpperCAmelCase_				:   Union[str, Any]							= 0
   for b in range(2**8							):
      if b not in bs:
         bs.append(__lowercase							)
         cs.append(2**8 + n							)
         n += 1
   UpperCAmelCase_				:   Dict							= [chr(__lowercase							) for n in cs]
   return dict(zip(__lowercase      ,							__lowercase							)							)
def   snake_case_    (			__lowercase							):
   UpperCAmelCase_				:   List[str]							= set()
   UpperCAmelCase_				:   Optional[Any]							= word[0]
   for char in word[1:]:
      pairs.add((prev_char, char)							)
      UpperCAmelCase_				:   Union[str, Any]							= char
   return pairs





class 			lowerCAmelCase__(     snake_case__    ):






    '''simple docstring'''





    A_	:  Optional[int]								=			VOCAB_FILES_NAMES
    A_	:  Any								=			PRETRAINED_VOCAB_FILES_MAP
    A_	:  Dict								=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    A_	:  List[str]								=			['input_ids', 'attention_mask']





    def __init__(						self					: Optional[int]							, __snake_case					: Dict							, __snake_case					: List[Any]							, __snake_case					: List[Any]="replace"							, __snake_case					: Any="<|endoftext|>"							, __snake_case					: Union[str, Any]="<|endoftext|>"							, __snake_case					: List[str]="<|endoftext|>"							, __snake_case					: Any=None							, __snake_case					: List[str]=False							, __snake_case					: Optional[int]=False							, **__snake_case					: str							, ):





       '''simple docstring'''
       UpperCAmelCase_				:   Optional[int]							= AddedToken(__snake_case							, lstrip=__snake_case							, rstrip=__snake_case						) if isinstance(__snake_case							, __snake_case						) else bos_token
       UpperCAmelCase_				:   List[Any]							= AddedToken(__snake_case							, lstrip=__snake_case							, rstrip=__snake_case						) if isinstance(__snake_case							, __snake_case						) else eos_token
       UpperCAmelCase_				:   Optional[int]							= AddedToken(__snake_case							, lstrip=__snake_case							, rstrip=__snake_case						) if isinstance(__snake_case							, __snake_case						) else unk_token
       UpperCAmelCase_				:   Tuple							= AddedToken(__snake_case							, lstrip=__snake_case							, rstrip=__snake_case						) if isinstance(__snake_case							, __snake_case						) else pad_token
       super().__init__(
           errors=__snake_case							, unk_token=__snake_case							, bos_token=__snake_case							, eos_token=__snake_case							, pad_token=__snake_case							, add_prefix_space=__snake_case							, add_bos_token=__snake_case							, **__snake_case							, )
       UpperCAmelCase_				:   List[str]							= add_bos_token

       with open(__snake_case							, encoding='''utf-8'''						) as vocab_handle:
          UpperCAmelCase_				:   Optional[int]							= json.load(__snake_case						)
       UpperCAmelCase_				:   int							= {v: k for k, v in self.encoder.items()}
       UpperCAmelCase_				:   Union[str, Any]							= errors  # how to handle errors in decoding
       UpperCAmelCase_				:   Optional[int]							= bytes_to_unicode()
       UpperCAmelCase_				:   str							= {v: k for k, v in self.byte_encoder.items()}
       with open(__snake_case							, encoding='''utf-8'''						) as merges_handle:
          UpperCAmelCase_				:   Any							= merges_handle.read().split('''\n'''						)[1:-1]
       UpperCAmelCase_				:   List[str]							= [tuple(merge.split()						) for merge in bpe_merges]
       UpperCAmelCase_				:   Optional[int]							= dict(zip(__snake_case							, range(len(__snake_case						)						)						)						)
       UpperCAmelCase_				:   Any							= {}
       UpperCAmelCase_				:   int							= add_prefix_space

       # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
       UpperCAmelCase_				:   int							= re.compile(R'''\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'''						)





    @property
    def        _lowerCamelCase			(						self					: int						):





       '''simple docstring'''
       return len(self.encoder						)





    def        _lowerCamelCase			(						self					: str						):





       '''simple docstring'''
       return dict(self.encoder							, **self.added_tokens_encoder						)





    def        _lowerCamelCase			(						self					: Optional[Any]							, __snake_case					: Optional[Any]						):





       '''simple docstring'''
       if token in self.cache:
          return self.cache[token]
       UpperCAmelCase_				:   Any							= tuple(__snake_case						)
       UpperCAmelCase_				:   Any							= get_pairs(__snake_case						)

       if not pairs:
          return token

       while True:
          UpperCAmelCase_				:   Any							= min(__snake_case							, key=lambda __snake_case						: self.bpe_ranks.get(__snake_case							, float('''inf'''						)						)						)
          if bigram not in self.bpe_ranks:
             break
          UpperCAmelCase_		,   UpperCAmelCase_				:   Dict							= bigram
          UpperCAmelCase_				:   List[Any]							= []
          UpperCAmelCase_				:   Union[str, Any]							= 0
          while i < len(__snake_case						):
             try:
                UpperCAmelCase_				:   Optional[int]							= word.index(__snake_case							, __snake_case						)
             except ValueError:
                new_word.extend(word[i:]						)
                break
             else:
                new_word.extend(word[i:j]						)
                UpperCAmelCase_				:   str							= j

             if word[i] == first and i < len(__snake_case						) - 1 and word[i + 1] == second:
                new_word.append(first + second						)
                i += 2
             else:
                new_word.append(word[i]						)
                i += 1
          UpperCAmelCase_				:   Any							= tuple(__snake_case						)
          UpperCAmelCase_				:   Tuple							= new_word
          if len(__snake_case						) == 1:
             break
          else:
             UpperCAmelCase_				:   str							= get_pairs(__snake_case						)
       UpperCAmelCase_				:   Union[str, Any]							= ''' '''.join(__snake_case						)
       UpperCAmelCase_				:   Optional[Any]							= word
       return word





    def        _lowerCamelCase			(						self					: List[Any]							, __snake_case					: Any							, __snake_case					: List[str]=None						):





       '''simple docstring'''
       if self.add_bos_token:
          UpperCAmelCase_				:   List[Any]							= [self.bos_token_id]
       else:
          UpperCAmelCase_				:   Dict							= []

       UpperCAmelCase_				:   Dict							= bos_token_ids + token_ids_a

       if token_ids_a is None:
          return output

       return output + bos_token_ids + token_ids_a





    def        _lowerCamelCase			(						self					: Optional[int]							, __snake_case					: List[int]							, __snake_case					: Optional[List[int]] = None							, __snake_case					: bool = False						):





       '''simple docstring'''
       if already_has_special_tokens:
          return super().get_special_tokens_mask(
              token_ids_a=__snake_case							, token_ids_a=__snake_case							, already_has_special_tokens=__snake_case						)

       if not self.add_bos_token:
          return super().get_special_tokens_mask(
              token_ids_a=__snake_case							, token_ids_a=__snake_case							, already_has_special_tokens=__snake_case						)

       if token_ids_a is None:
          return [1] + ([0] * len(__snake_case						))
       return [1] + ([0] * len(__snake_case						)) + [1] + ([0] * len(__snake_case						))





    def        _lowerCamelCase			(						self					: str							, __snake_case					: Union[str, Any]						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[str]							= []
       for token in re.findall(self.pat							, __snake_case						):
          UpperCAmelCase_				:   Any							= ''''''.join(
              self.byte_encoder[b] for b in token.encode('''utf-8'''						)						)  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
          bpe_tokens.extend(bpe_token for bpe_token in self.bpe(__snake_case						).split(''' '''						)						)
       return bpe_tokens





    def        _lowerCamelCase			(						self					: Optional[Any]							, __snake_case					: Optional[int]						):





       '''simple docstring'''
       return self.encoder.get(__snake_case							, self.encoder.get(self.unk_token						)						)





    def        _lowerCamelCase			(						self					: int							, __snake_case					: str						):





       '''simple docstring'''
       return self.decoder.get(__snake_case						)





    def        _lowerCamelCase			(						self					: Union[str, Any]							, __snake_case					: List[str]						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[Any]							= ''''''.join(__snake_case						)
       UpperCAmelCase_				:   str							= bytearray([self.byte_decoder[c] for c in text]						).decode('''utf-8'''							, errors=self.errors						)
       return text





    def        _lowerCamelCase			(						self					: str							, __snake_case					: str							, __snake_case					: Optional[str] = None						):





       '''simple docstring'''
       if not os.path.isdir(__snake_case						):
          logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''						)
          return
       UpperCAmelCase_				:   List[str]							= os.path.join(
           __snake_case							, (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']						)
       UpperCAmelCase_				:   List[Any]							= os.path.join(
           __snake_case							, (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']						)

       with open(__snake_case							, '''w'''							, encoding='''utf-8'''						) as f:
          f.write(json.dumps(self.encoder							, indent=2							, sort_keys=__snake_case							, ensure_ascii=__snake_case						) + '''\n'''						)

       UpperCAmelCase_				:   Union[str, Any]							= 0
       with open(__snake_case							, '''w'''							, encoding='''utf-8'''						) as writer:
          writer.write('''#version: 0.2\n'''						)
          for bpe_tokens, token_index in sorted(self.bpe_ranks.items()							, key=lambda __snake_case						: kv[1]						):
             if index != token_index:
                logger.warning(
                    f'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
                    ''' Please check that the tokenizer is not corrupted!'''						)
                UpperCAmelCase_				:   Dict							= token_index
             writer.write(''' '''.join(__snake_case						) + '''\n'''						)
             index += 1

       return vocab_file, merge_file





    def        _lowerCamelCase			(						self					: List[Any]							, __snake_case					: str							, __snake_case					: List[Any]=False							, **__snake_case					: Dict						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[Any]							= kwargs.pop('''add_prefix_space'''							, self.add_prefix_space						)
       if is_split_into_words or add_prefix_space:
          UpperCAmelCase_				:   str							= ''' ''' + text
       return (text, kwargs)





    def        _lowerCamelCase			(						self					: Union[str, Any]							, __snake_case					: "Conversation"						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[Any]							= []
       for is_user, text in conversation.iter_texts():
          input_ids.extend(self.encode(__snake_case							, add_special_tokens=__snake_case						) + [self.eos_token_id]						)
       if len(__snake_case						) > self.model_max_length:
          UpperCAmelCase_				:   str							= input_ids[-self.model_max_length :]
       return input_ids