



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
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to 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)
#
# 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		:Any   	=  16
__snake_case		:Dict   	=  32





def   __snake_case			(    _UpperCAmelCase     ,				_UpperCAmelCase = 16       ):
    __a		 =       AutoTokenizer.from_pretrained('''bert-base-cased'''       )
    __a		 =       load_dataset('''glue'''     ,				'''mrpc'''       )

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

    def collate_fn(_UpperCAmelCase       ):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        __a		 =       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":
            __a		 =       16
        elif accelerator.mixed_precision != "no":
            __a		 =       8
        else:
            __a		 =       None

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

    # Instantiate dataloaders.
    __a		 =       DataLoader(
        tokenized_datasets['''train''']     ,				shuffle=_UpperCAmelCase     ,				collate_fn=_UpperCAmelCase     ,				batch_size=_UpperCAmelCase       )
    __a		 =       DataLoader(
        tokenized_datasets['''validation''']     ,				shuffle=_UpperCAmelCase     ,				collate_fn=_UpperCAmelCase     ,				batch_size=_UpperCAmelCase       )

    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		:Optional[Any]   	=  mocked_dataloaders  # noqa: F811





def   __snake_case			(    _UpperCAmelCase     ,				_UpperCAmelCase       ):
    # For testing only
    if os.environ.get('''TESTING_MOCKED_DATALOADERS'''     ,				_UpperCAmelCase       ) == "1":
        __a		 =       2
    # New Code #
    __a		 =       int(args.gradient_accumulation_steps       )
    __a		 =       int(args.local_sgd_steps       )
    # Initialize accelerator
    __a		 =       Accelerator(
        cpu=args.cpu     ,				mixed_precision=args.mixed_precision     ,				gradient_accumulation_steps=_UpperCAmelCase       )
    if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
        raise NotImplementedError('''LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'''       )
    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    __a		 =       config['''lr''']
    __a		 =       int(config['''num_epochs''']       )
    __a		 =       int(config['''seed''']       )
    __a		 =       int(config['''batch_size''']       )

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

    set_seed(_UpperCAmelCase       )
    __a ,     __a		 =       get_dataloaders(_UpperCAmelCase     ,				_UpperCAmelCase       )
    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    __a		 =       AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''     ,				return_dict=_UpperCAmelCase       )

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

    # Instantiate optimizer
    __a		 =       AdamW(params=model.parameters()     ,				lr=_UpperCAmelCase       )

    # Instantiate scheduler
    __a		 =       get_linear_schedule_with_warmup(
        optimizer=_UpperCAmelCase     ,				num_warmup_steps=100     ,				num_training_steps=(len(_UpperCAmelCase       ) * 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.
    __a ,     __a ,     __a ,     __a ,     __a		 =       accelerator.prepare(
        _UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase       )

    # Now we train the model
    for epoch in range(_UpperCAmelCase       ):
        model.train()
        with LocalSGD(
            accelerator=_UpperCAmelCase     ,				model=_UpperCAmelCase     ,				local_sgd_steps=_UpperCAmelCase     ,				enabled=local_sgd_steps is not None       ) as local_sgd:
            for step, batch in enumerate(_UpperCAmelCase       ):
                # We could avoid this line since we set the accelerator with `device_placement=True`.
                batch.to(accelerator.device       )
                # New code #
                # We use the new `accumulate` context manager to perform gradient accumulation
                # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
                with accelerator.accumulate(_UpperCAmelCase       ):
                    __a		 =       model(**_UpperCAmelCase       )
                    __a		 =       output.loss
                    accelerator.backward(_UpperCAmelCase       )
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
                    # LocalSGD-specific line
                    local_sgd.step()

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

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





def   __snake_case			(    ):
    __a		 =       argparse.ArgumentParser(description='''Simple example of training script.'''       )
    parser.add_argument(
        '''--mixed_precision'''     ,				type=_UpperCAmelCase     ,				default=_UpperCAmelCase     ,				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.'''     ,				)
    # New Code #
    parser.add_argument(
        '''--gradient_accumulation_steps'''     ,				type=_UpperCAmelCase     ,				default=1     ,				help='''The number of minibatches to be ran before gradients are accumulated.'''     ,				)
    parser.add_argument(
        '''--local_sgd_steps'''     ,				type=_UpperCAmelCase     ,				default=8     ,				help='''Number of local SGD steps or None to disable local SGD'''       )
    parser.add_argument('''--cpu'''     ,				action='''store_true'''     ,				help='''If passed, will train on the CPU.'''       )
    __a		 =       parser.parse_args()
    __a		 =       {'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
    training_function(_UpperCAmelCase     ,				_UpperCAmelCase       )


if __name__ == "__main__":
 main()










import unittest

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

from ...test_tokenization_common import TokenizerTesterMixin


__snake_case		:Dict   	=  get_tests_dir('''fixtures/test_sentencepiece_with_bytefallback.model''')

@require_sentencepiece
@require_tokenizers
class 					_A (			__UpperCAmelCase   ,unittest.TestCase					):
       UpperCamelCase__						:					List[str]		      =		GPTSwaTokenizer
       UpperCamelCase__						:					Dict		      =		False
       UpperCamelCase__						:					int		      =		True
       UpperCamelCase__						:					List[Any]		      =		False



       def 	_lowerCamelCase    (	self							:  List[Any]):


           '''simple docstring'''

           super().setUp()

           # We have a SentencePiece fixture for testing
           __a		 =       GPTSwaTokenizer(__SCREAMING_SNAKE_CASE  ,	eos_token='''<unk>'''  ,	bos_token='''<unk>'''  ,	pad_token='''<unk>''')

           tokenizer.save_pretrained(self.tmpdirname)



       def 	_lowerCamelCase    (	self							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  int):


           '''simple docstring'''

           __a		 =       '''This is a test'''
           __a		 =       '''This is a test'''
           return input_text, output_text



       def 	_lowerCamelCase    (	self							:  Dict):


           '''simple docstring'''

           __a		 =       '''<s>'''
           __a		 =       1

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



       def 	_lowerCamelCase    (	self							:  List[Any]):


           '''simple docstring'''

           __a		 =       list(self.get_tokenizer().get_vocab().keys())

           self.assertEqual(vocab_keys[0]  ,	'''<unk>''')
           self.assertEqual(vocab_keys[1]  ,	'''<s>''')
           self.assertEqual(vocab_keys[-1]  ,	'''j''')
           self.assertEqual(len(__SCREAMING_SNAKE_CASE)  ,	2_000)



       def 	_lowerCamelCase    (	self							:  Dict):


           '''simple docstring'''

           self.assertEqual(self.get_tokenizer().vocab_size  ,	2_000)



       def 	_lowerCamelCase    (	self							:  List[str]):


           '''simple docstring'''

           __a		 =       GPTSwaTokenizer(__SCREAMING_SNAKE_CASE)

           __a		 =       tokenizer.tokenize('''This is a test''')
           self.assertListEqual(__SCREAMING_SNAKE_CASE  ,	['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est'''])

           self.assertListEqual(tokenizer.convert_tokens_to_ids(__SCREAMING_SNAKE_CASE)  ,	[465, 287, 265, 631, 842])

           __a		 =       tokenizer.tokenize('''I was born in 92000, and this is falsé.''')
           # fmt: off
           self.assertListEqual(
               __SCREAMING_SNAKE_CASE  ,	['''▁I''', '''▁was''', '''▁bor''', '''n''', '''▁in''', '''▁''', '''<0x39>''', '''2''', '''0''', '''0''', '''0''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁f''', '''al''', '''s''', '''<0xC3>''', '''<0xA9>''', '''.''']  ,	)
           # fmt: on

           __a		 =       tokenizer.convert_tokens_to_ids(__SCREAMING_SNAKE_CASE)
           self.assertListEqual(
               __SCREAMING_SNAKE_CASE  ,	[262, 272, 1_525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260]  ,	)

           __a		 =       tokenizer.convert_ids_to_tokens(__SCREAMING_SNAKE_CASE)
           # fmt: off
           self.assertListEqual(
               __SCREAMING_SNAKE_CASE  ,	['''▁I''', '''▁was''', '''▁bor''', '''n''', '''▁in''', '''▁''', '''<0x39>''', '''2''', '''0''', '''0''', '''0''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁f''', '''al''', '''s''', '''<0xC3>''', '''<0xA9>''', '''.'''])
           # fmt: on



       def 	_lowerCamelCase    (	self							:  Any):


           '''simple docstring'''

           __a		 =       GPTSwaTokenizer(__SCREAMING_SNAKE_CASE)
           __a		 =       ['''This is a test''', '''I was born in 92000, and this is falsé.''']
           __a		 =       [
               [465, 287, 265, 631, 842],
               [262, 272, 1_525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260],
           ]

           # Test that encode_fast returns the same as tokenize + convert_tokens_to_ids
           for text, expected_ids in zip(__SCREAMING_SNAKE_CASE  ,	__SCREAMING_SNAKE_CASE):
               self.assertListEqual(tokenizer.encode_fast(__SCREAMING_SNAKE_CASE)  ,	__SCREAMING_SNAKE_CASE)

           # Test that decode_fast returns the input text
           for text, token_ids in zip(__SCREAMING_SNAKE_CASE  ,	__SCREAMING_SNAKE_CASE):
               self.assertEqual(tokenizer.decode_fast(__SCREAMING_SNAKE_CASE)  ,	__SCREAMING_SNAKE_CASE)





       @slow
       def 	_lowerCamelCase    (	self							:  Any):


           '''simple docstring'''

           __a		 =       [
               '''<|python|>def fibonacci(n)\n    if n < 0:\n        print(\'Incorrect input\')''',
               '''Hey there, how are you doing this fine day?''',
               '''This is a text with a trailing spaces followed by a dot     .''',
               '''Häj sväjs lillebrör! =)''',
               '''Det är inget fel på Mr. Cool''',
           ]

           # fmt: off
           __a		 =       {'''input_ids''': [[63_423, 5, 6_811, 14_954, 282, 816, 3_821, 63_466, 63_425, 63_462, 18, 63_978, 678, 301, 1_320, 63_423, 63_455, 63_458, 18, 63_982, 4_246, 3_940, 1_901, 47_789, 5_547, 18_994], [19_630, 1_100, 63_446, 1_342, 633, 544, 4_488, 593, 5_102, 2_416, 63_495, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1_652, 428, 268, 1_936, 515, 268, 58_593, 22_413, 9_106, 546, 268, 33_213, 63_979, 698, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [55_130, 63_450, 924, 63_449, 2_249, 4_062, 1_558, 318, 63_504, 21_498, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [509, 377, 2_827, 2_559, 332, 6_575, 63_443, 26_801, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], '''token_type_ids''': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], '''attention_mask''': [[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, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}
           # fmt: on
           self.tokenizer_integration_test_util(
               expected_encoding=__SCREAMING_SNAKE_CASE  ,	model_name='''AI-Sweden/gpt-sw3-126m'''  ,	sequences=__SCREAMING_SNAKE_CASE  ,	)





