


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
#
# 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)
#
# This example also demonstrates the checkpointing and sharding capabilities
#
# 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_           =	1_6
lowercase_           =	3_2
def       UpperCamelCase__     (  SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__						):
 # Initialize accelerator
 if args.with_tracking:
  __lowerCamelCase		:							int							=						Accelerator(
      cpu=args.cpu			,       mixed_precision=args.mixed_precision			,       log_with='all'			,       project_dir=args.project_dir						)
 else:
  __lowerCamelCase		:							Union[str, Any]							=						Accelerator(cpu=args.cpu			,       mixed_precision=args.mixed_precision						)

 if hasattr(args.checkpointing_steps			,       'isdigit'						):
  if args.checkpointing_steps == "epoch":
   __lowerCamelCase		:							Union[str, Any]							=						args.checkpointing_steps
  elif args.checkpointing_steps.isdigit():
   __lowerCamelCase		:							Union[str, Any]							=						int(args.checkpointing_steps						)
  else:
   raise ValueError(
       f'Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'						)
 else:
  __lowerCamelCase		:							Any							=						None
 # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
 __lowerCamelCase		:							Optional[int]							=						config['lr']
 __lowerCamelCase		:							Tuple							=						int(config['num_epochs']						)
 __lowerCamelCase		:							Any							=						int(config['seed']						)
 __lowerCamelCase		:							Union[str, Any]							=						int(config['batch_size']						)

 # We need to initialize the trackers we use, and also store our configuration
 if args.with_tracking:
  __lowerCamelCase		:							Optional[int]							=						os.path.split(SCREAMING_SNAKE_CASE__						)[-1].split('.'						)[0]
  accelerator.init_trackers(SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__						)

 __lowerCamelCase		:							Tuple							=						AutoTokenizer.from_pretrained('bert-base-cased'						)
 __lowerCamelCase		:							Union[str, Any]							=						load_dataset('glue'			,       'mrpc'						)
 __lowerCamelCase		:							str							=						evaluate.load('glue'			,       'mrpc'						)

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

 # If the batch size is too big we use gradient accumulation
 __lowerCamelCase		:							int							=						1
 if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
  __lowerCamelCase		:							Optional[Any]							=						batch_size // MAX_GPU_BATCH_SIZE
  __lowerCamelCase		:							Union[str, Any]							=						MAX_GPU_BATCH_SIZE

 def collate_fn(SCREAMING_SNAKE_CASE__						):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  __lowerCamelCase		:							Dict							=						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":
   __lowerCamelCase		:							Union[str, Any]							=						16
  elif accelerator.mixed_precision != "no":
   __lowerCamelCase		:							int							=						8
  else:
   __lowerCamelCase		:							Tuple							=						None

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

 # Instantiate dataloaders.
 __lowerCamelCase		:							Any							=						DataLoader(
     tokenized_datasets['train']			,       shuffle=SCREAMING_SNAKE_CASE__			,       collate_fn=SCREAMING_SNAKE_CASE__			,       batch_size=SCREAMING_SNAKE_CASE__						)
 __lowerCamelCase		:							Optional[Any]							=						DataLoader(
     tokenized_datasets['validation']			,       shuffle=SCREAMING_SNAKE_CASE__			,       collate_fn=SCREAMING_SNAKE_CASE__			,       batch_size=SCREAMING_SNAKE_CASE__						)

 set_seed(SCREAMING_SNAKE_CASE__						)

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

 # 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).
 __lowerCamelCase		:							Tuple							=						model.to(accelerator.device						)

 # Instantiate optimizer
 __lowerCamelCase		:							Optional[Any]							=						AdamW(params=model.parameters()			,       lr=SCREAMING_SNAKE_CASE__						)

 # Instantiate scheduler
 __lowerCamelCase		:							List[Any]							=						get_linear_schedule_with_warmup(
     optimizer=SCREAMING_SNAKE_CASE__			,       num_warmup_steps=100			,       num_training_steps=(len(SCREAMING_SNAKE_CASE__						) * 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.
 __lowerCamelCase					,     __lowerCamelCase					,     __lowerCamelCase					,     __lowerCamelCase					,     __lowerCamelCase		:							Tuple							=						accelerator.prepare(
     SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__						)

 # We need to keep track of how many total steps we have iterated over
 __lowerCamelCase		:							Any							=						0
 # We also need to keep track of the stating epoch so files are named properly
 __lowerCamelCase		:							Dict							=						0

 # Potentially load in the weights and states from a previous save
 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						)
   __lowerCamelCase		:							Tuple							=						os.path.basename(args.resume_from_checkpoint						)
  else:
   # Get the most recent checkpoint
   __lowerCamelCase		:							List[Any]							=						[f.name for f in os.scandir(os.getcwd()						) if f.is_dir()]
   dirs.sort(key=os.path.getctime						)
   __lowerCamelCase		:							int							=						dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
  # Extract `epoch_{i}` or `step_{i}`
  __lowerCamelCase		:							List[Any]							=						os.path.splitext(SCREAMING_SNAKE_CASE__						)[0]

  if "epoch" in training_difference:
   __lowerCamelCase		:							int							=						int(training_difference.replace('epoch_'			,       ''						)						) + 1
   __lowerCamelCase		:							Optional[int]							=						None
  else:
   __lowerCamelCase		:							List[str]							=						int(training_difference.replace('step_'			,       ''						)						)
   __lowerCamelCase		:							List[str]							=						resume_step // len(SCREAMING_SNAKE_CASE__						)
   resume_step -= starting_epoch * len(SCREAMING_SNAKE_CASE__						)

    # Now we train the model
 for epoch in range(SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__						):
  model.train()
  if args.with_tracking:
   __lowerCamelCase		:							List[str]							=						0
  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
   __lowerCamelCase		:							Union[str, Any]							=						accelerator.skip_first_batches(SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__						)
   overall_step += resume_step
  else:
   # After the first iteration though, we need to go back to the original dataloader
   __lowerCamelCase		:							Any							=						train_dataloader
  for step, batch in enumerate(SCREAMING_SNAKE_CASE__						):
   # We could avoid this line since we set the accelerator with `device_placement=True`.
   batch.to(accelerator.device						)
   __lowerCamelCase		:							Optional[Any]							=						model(**SCREAMING_SNAKE_CASE__						)
   __lowerCamelCase		:							List[str]							=						outputs.loss
   __lowerCamelCase		:							Any							=						loss / gradient_accumulation_steps
   # We keep track of the loss at each epoch
   if args.with_tracking:
    total_loss += loss.detach().float()
   accelerator.backward(SCREAMING_SNAKE_CASE__						)
   if step % gradient_accumulation_steps == 0:
    optimizer.step()
    lr_scheduler.step()
    optimizer.zero_grad()

   overall_step += 1

   if isinstance(SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__						):
    __lowerCamelCase		:							Union[str, Any]							=						f'step_{overall_step}'
    if overall_step % checkpointing_steps == 0:
     if args.output_dir is not None:
      __lowerCamelCase		:							List[str]							=						os.path.join(args.output_dir			,       SCREAMING_SNAKE_CASE__						)
     accelerator.save_state(SCREAMING_SNAKE_CASE__						)

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

  __lowerCamelCase		:							Tuple							=						metric.compute()
  # Use accelerator.print to print only on the main process.
  accelerator.print(f'epoch {epoch}:'			,       SCREAMING_SNAKE_CASE__						)
  if args.with_tracking:
   accelerator.log(
       {
           'accuracy': eval_metric['accuracy'],
           'f1': eval_metric['f1'],
           'train_loss': total_loss.item() / len(SCREAMING_SNAKE_CASE__						),
           'epoch': epoch,
       }			,       step=SCREAMING_SNAKE_CASE__			,       )

  if checkpointing_steps == "epoch":
   __lowerCamelCase		:							str							=						f'epoch_{epoch}'
   if args.output_dir is not None:
    __lowerCamelCase		:							int							=						os.path.join(args.output_dir			,       SCREAMING_SNAKE_CASE__						)
   accelerator.save_state(SCREAMING_SNAKE_CASE__						)

 if args.with_tracking:
  accelerator.end_training()
def       UpperCamelCase__     (  ):
 __lowerCamelCase		:							List[Any]							=						argparse.ArgumentParser(description='Simple example of training script.'						)
 parser.add_argument(
     '--mixed_precision'			,       type=SCREAMING_SNAKE_CASE__			,       default=SCREAMING_SNAKE_CASE__			,       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=SCREAMING_SNAKE_CASE__			,       default=SCREAMING_SNAKE_CASE__			,       help='Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch.'			,       )
 parser.add_argument(
     '--resume_from_checkpoint'			,       type=SCREAMING_SNAKE_CASE__			,       default=SCREAMING_SNAKE_CASE__			,       help='If the training should continue from a checkpoint folder.'			,       )
 parser.add_argument(
     '--with_tracking'			,       action='store_true'			,       help='Whether to load in all available experiment trackers from the environment and use them for logging.'			,       )
 parser.add_argument(
     '--output_dir'			,       type=SCREAMING_SNAKE_CASE__			,       default='.'			,       help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'			,       )
 parser.add_argument(
     '--project_dir'			,       type=SCREAMING_SNAKE_CASE__			,       default='logs'			,       help='Location on where to store experiment tracking logs` and relevent project information'			,       )
 __lowerCamelCase		:							Optional[int]							=						parser.parse_args()
 __lowerCamelCase		:							Any							=						{'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
 training_function(SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__						)


if __name__ == "__main__":
       main()





import unittest

from transformers import ErnieMTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


lowercase_           =	get_tests_dir('fixtures/spiece.model')


@require_sentencepiece
@require_tokenizers
class        A_					(       __UpperCamelCase					,      unittest.TestCase    ):




      '''simple docstring'''


      __snake_case					       =   ErnieMTokenizer
      __snake_case					       =   False
      __snake_case					       =   True
      __snake_case					       =   False
      __snake_case					       =   False




      def      _snake_case	(				self:       Optional[int]				):
       super().setUp()

       # We have a SentencePiece fixture for testing
       __lowerCamelCase		:							List[str]							=						ErnieMTokenizer(a						,      unk_token='<unk>'						,      pad_token='<pad>'				)
       tokenizer.save_pretrained(self.tmpdirname				)




      def      _snake_case	(				self:       Dict						,      a:       List[Any]				):
       __lowerCamelCase		:							Optional[Any]							=						'this is a test'
       __lowerCamelCase		:							Optional[int]							=						'this is a test'
       return input_text, output_text




      def      _snake_case	(				self:       Tuple				):
       __lowerCamelCase		:							List[str]							=						'<pad>'
       __lowerCamelCase		:							Union[str, Any]							=						0

       self.assertEqual(self.get_tokenizer()._convert_token_to_id(a				)						,      a				)
       self.assertEqual(self.get_tokenizer()._convert_id_to_token(a				)						,      a				)




      def      _snake_case	(				self:       List[Any]				):
       __lowerCamelCase		:							List[str]							=						list(self.get_tokenizer().get_vocab().keys()				)

       self.assertEqual(vocab_keys[0]						,      '<pad>'				)
       self.assertEqual(vocab_keys[1]						,      '<unk>'				)
       self.assertEqual(vocab_keys[-1]						,      '▁eloquent'				)
       self.assertEqual(len(a				)						,      3_0000				)




      def      _snake_case	(				self:       Tuple				):
       self.assertEqual(self.get_tokenizer().vocab_size						,      3_0000				)




      def      _snake_case	(				self:       int				):
       if not self.test_rust_tokenizer:
        return

       __lowerCamelCase		:							Any							=						self.get_tokenizer()
       __lowerCamelCase		:							Optional[Any]							=						self.get_rust_tokenizer()

       __lowerCamelCase		:							Optional[int]							=						'I was born in 92000, and this is falsé.'

       __lowerCamelCase		:							Optional[Any]							=						tokenizer.tokenize(a				)
       __lowerCamelCase		:							Dict							=						rust_tokenizer.tokenize(a				)
       self.assertListEqual(a						,      a				)

       __lowerCamelCase		:							List[str]							=						tokenizer.encode(a						,      add_special_tokens=a				)
       __lowerCamelCase		:							List[Any]							=						rust_tokenizer.encode(a						,      add_special_tokens=a				)
       self.assertListEqual(a						,      a				)

       __lowerCamelCase		:							str							=						self.get_rust_tokenizer()
       __lowerCamelCase		:							str							=						tokenizer.encode(a				)
       __lowerCamelCase		:							List[str]							=						rust_tokenizer.encode(a				)
       self.assertListEqual(a						,      a				)




      def      _snake_case	(				self:       Optional[Any]				):
       __lowerCamelCase		:							Union[str, Any]							=						ErnieMTokenizer(a						,      do_lower_case=a						,      unk_token='<unk>'						,      pad_token='<pad>'				)

       __lowerCamelCase		:							List[str]							=						tokenizer.tokenize('This is a test'				)
       self.assertListEqual(a						,      ['▁this', '▁is', '▁a', '▁test']				)

       self.assertListEqual(tokenizer.convert_tokens_to_ids(a				)						,      [48, 25, 21, 1289]				)

       __lowerCamelCase		:							Optional[Any]							=						tokenizer.tokenize('I was born in 92000, and this is falsé.'				)
       # ErnieMTokenizer(paddlenlp implementation) outputs '9' instead of '_9' so to mimic that '_9' is changed to '9'
       self.assertListEqual(
           a						,      ['▁i', '▁was', '▁born', '▁in', '9', '2000', ',', '▁and', '▁this', '▁is', '▁fal', 's', 'é', '.']				)
       __lowerCamelCase		:							Any							=						tokenizer.convert_tokens_to_ids(a				)
       self.assertListEqual(a						,      [31, 23, 386, 19, 518, 3050, 15, 17, 48, 25, 8256, 18, 1, 9]				)

       __lowerCamelCase		:							int							=						tokenizer.convert_ids_to_tokens(a				)
       self.assertListEqual(
           a						,      ['▁i', '▁was', '▁born', '▁in', '9', '2000', ',', '▁and', '▁this', '▁is', '▁fal', 's', '<unk>', '.']						,      )




      def      _snake_case	(				self:       str				):
       __lowerCamelCase		:							Tuple							=						ErnieMTokenizer(a						,      unk_token='<unk>'						,      pad_token='<pad>'				)

       __lowerCamelCase		:							Tuple							=						tokenizer.encode('sequence builders'				)
       __lowerCamelCase		:							List[Any]							=						tokenizer.encode('multi-sequence build'				)

       __lowerCamelCase		:							Optional[Any]							=						tokenizer.build_inputs_with_special_tokens(a				)
       __lowerCamelCase		:							Tuple							=						tokenizer.build_inputs_with_special_tokens(a						,      a				)

       assert encoded_sentence == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id]
       assert encoded_pair == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id] + [
           tokenizer.sep_token_id
       ] + text_a + [tokenizer.sep_token_id]




      @slow
      def      _snake_case	(				self:       List[Any]				):
       # fmt: off
       __lowerCamelCase		:							int							=						{'input_ids': [[0, 1_1062, 8_2772, 7, 15, 8_2772, 538, 5_1529, 237, 1_7198, 1290, 206, 9, 21_5175, 1314, 136, 1_7198, 1290, 206, 9, 5_6359, 42, 12_2009, 9, 1_6466, 16, 8_7344, 4537, 9, 4717, 7_8381, 6, 15_9958, 7, 15, 2_4480, 618, 4, 527, 2_2693, 9, 304, 4, 2777, 2_4480, 9874, 4, 4_3523, 594, 4, 803, 1_8392, 3_3189, 18, 4, 4_3523, 2_4447, 5, 5, 5, 16, 100, 2_4955, 8_3658, 9626, 14_4057, 15, 839, 2_2335, 16, 136, 2_4955, 8_3658, 8_3479, 15, 3_9102, 724, 16, 678, 645, 6460, 1328, 4589, 42, 12_2009, 11_5774, 23, 3559, 1328, 4_6876, 7, 136, 5_3894, 1940, 4_2227, 4_1159, 1_7721, 823, 425, 4, 2_7512, 9_8722, 206, 136, 5531, 4970, 919, 1_7336, 5, 2], [0, 2_0080, 618, 83, 8_2775, 47, 479, 9, 1517, 73, 5_3894, 333, 8_0581, 11_0117, 1_8811, 5256, 1295, 51, 15_2526, 297, 7986, 390, 12_4416, 538, 3_5431, 214, 98, 1_5044, 2_5737, 136, 7108, 4_3701, 23, 756, 13_5355, 7, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 581, 6_3773, 11_9455, 6, 14_7797, 8_8203, 7, 645, 70, 21, 3285, 1_0269, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}
       # fmt: on

       self.tokenizer_integration_test_util(
           expected_encoding=a						,      model_name='susnato/ernie-m-base_pytorch'						,      sequences=[
               'Transformers (formerly known as pytorch-transformers and pytorch-pretrained-bert) provides '
               'general-purpose architectures (BERT, GPT-2, RoBERTa, XLM, DistilBert, XLNet...) for Natural '
               'Language Understanding (NLU) and Natural Language Generation (NLG) with over32+ pretrained '
               'models in100+ languages and deep interoperability between Jax, PyTorch and TensorFlow.',
               'BERT is designed to pre-train deep bidirectional representations from unlabeled text by jointly '
               'conditioning on both left and right context in all layers.',
               'The quick brown fox jumps over the lazy dog.',
           ]						,      )

