





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


lowerCAmelCase_    =    1_6
lowerCAmelCase_    =    3_2






def 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase = 16			)		->						Tuple:






				"""simple docstring"""


				snake_case_						:   Dict           =  AutoTokenizer.from_pretrained('''bert-base-cased'''			)
				snake_case_						:   Any           =  load_dataset('''glue'''			,    '''mrpc'''			)

				def tokenize_function(_UpperCamelCase			):
								# max_length=None => use the model max length (it's actually the default)
								snake_case_						:   Optional[int]           =  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():
								snake_case_						:   List[str]           =  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
				snake_case_						:   Optional[Any]           =  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.
								snake_case_						:   str           =  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":
												snake_case_						:   int           =  16
								elif accelerator.mixed_precision != "no":
												snake_case_						:   Tuple           =  8
								else:
												snake_case_						:   Optional[Any]           =  None

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

				# Instantiate dataloaders.
				snake_case_						:   Tuple           =  DataLoader(
				    tokenized_datasets['''train''']			,    shuffle=_UpperCamelCase			,    collate_fn=_UpperCamelCase			,    batch_size=_UpperCamelCase			)
				snake_case_						:   List[str]           =  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

		lowerCAmelCase_    =    mocked_dataloaders  # noqa: F811






def 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase			)		->						Tuple:






				"""simple docstring"""


				if os.environ.get('''TESTING_MOCKED_DATALOADERS'''			,    _UpperCamelCase			) == "1":
								snake_case_						:   Tuple           =  2
				# New Code #
				snake_case_						:   Tuple           =  int(args.gradient_accumulation_steps			)
				snake_case_						:   Union[str, Any]           =  int(args.local_sgd_steps			)
				# Initialize accelerator
				snake_case_						:   List[Any]           =  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
				snake_case_						:   Any           =  config['''lr''']
				snake_case_						:   Any           =  int(config['''num_epochs''']			)
				snake_case_						:   Tuple           =  int(config['''seed''']			)
				snake_case_						:   Dict           =  int(config['''batch_size''']			)

				snake_case_						:   List[str]           =  evaluate.load('''glue'''			,    '''mrpc'''			)

				set_seed(_UpperCamelCase			)
				snake_case_						:   int           =  get_dataloaders(_UpperCamelCase			,    _UpperCamelCase			)
				# Instantiate the model (we build the model here so that the seed also control new weights initialization)
				snake_case_						:   Union[str, Any]           =  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).
				snake_case_						:   str           =  model.to(accelerator.device			)

				# Instantiate optimizer
				snake_case_						:   Tuple           =  AdamW(params=model.parameters()			,    lr=_UpperCamelCase			)

				# Instantiate scheduler
				snake_case_						:   str           =  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.
				snake_case_						:   Dict           =  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			):
																				snake_case_						:   List[str]           =  model(**_UpperCamelCase			)
																				snake_case_						:   int           =  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():
																snake_case_						:   int           =  model(**_UpperCamelCase			)
												snake_case_						:   str           =  outputs.logits.argmax(dim=-1			)
												snake_case_						:   Tuple           =  accelerator.gather_for_metrics((predictions, batch['''labels'''])			)
												metric.add_batch(
												    predictions=_UpperCamelCase			,    references=_UpperCamelCase			,    )

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






def 					lowerCamelCase_						(    )		->						Tuple:






				"""simple docstring"""


				snake_case_						:   List[str]           =  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.'''			)
				snake_case_						:   Optional[int]           =  parser.parse_args()
				snake_case_						:   List[str]           =  {'''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


lowerCAmelCase_    =    get_tests_dir('''fixtures/test_sentencepiece_with_bytefallback.model''')



@require_sentencepiece
@require_tokenizers
class  __lowerCAmelCase	(      _a,					unittest.TestCase					):
			lowerCamelCase_					:					List[str]	      =				GPTSwaTokenizer
			lowerCamelCase_					:					str	      =				False
			lowerCamelCase_					:					str	      =				True
			lowerCamelCase_					:					List[Any]	      =				False






			def       lowerCamelCase  (self   )					->       List[str]:


							'''simple docstring'''



							super().setUp()

							# We have a SentencePiece fixture for testing
							snake_case_						:   Union[str, Any]           =  GPTSwaTokenizer(__magic_name__ ,							eos_token='''<unk>''' ,							bos_token='''<unk>''' ,							pad_token='''<unk>'''   )

							tokenizer.save_pretrained(self.tmpdirname   )






			def       lowerCamelCase  (self ,							__magic_name__   )					->       int:


							'''simple docstring'''



							snake_case_						:   Any           =  '''This is a test'''
							snake_case_						:   str           =  '''This is a test'''
							return input_text, output_text






			def       lowerCamelCase  (self   )					->       Any:


							'''simple docstring'''



							snake_case_						:   int           =  '''<s>'''
							snake_case_						:   Any           =  1

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






			def       lowerCamelCase  (self   )					->       List[Any]:


							'''simple docstring'''



							snake_case_						:   Union[str, Any]           =  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(__magic_name__   ) ,							2000   )






			def       lowerCamelCase  (self   )					->       List[str]:


							'''simple docstring'''



							self.assertEqual(self.get_tokenizer().vocab_size ,							2000   )






			def       lowerCamelCase  (self   )					->       List[str]:


							'''simple docstring'''



							snake_case_						:   str           =  GPTSwaTokenizer(__magic_name__   )

							snake_case_						:   List[str]           =  tokenizer.tokenize('''This is a test'''   )
							self.assertListEqual(__magic_name__ ,							['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']   )

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

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

							snake_case_						:   int           =  tokenizer.convert_tokens_to_ids(__magic_name__   )
							self.assertListEqual(
							    __magic_name__ ,							[262, 272, 1525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260] ,							)

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






			def       lowerCamelCase  (self   )					->       Union[str, Any]:


							'''simple docstring'''



							snake_case_						:   Optional[Any]           =  GPTSwaTokenizer(__magic_name__   )
							snake_case_						:   Tuple           =  ['''This is a test''', '''I was born in 92000, and this is falsé.''']
							snake_case_						:   Optional[int]           =  [
							    [465, 287, 265, 631, 842],
							    [262, 272, 1525, 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(__magic_name__ ,							__magic_name__   ):
											self.assertListEqual(tokenizer.encode_fast(__magic_name__   ) ,							__magic_name__   )

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






			@slow
			def       lowerCamelCase  (self   )					->       Tuple:


							'''simple docstring'''



							snake_case_						:   str           =  [
							    '''<|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
							snake_case_						:   str           =  {'''input_ids''': [[6_3423, 5, 6811, 1_4954, 282, 816, 3821, 6_3466, 6_3425, 6_3462, 18, 6_3978, 678, 301, 1320, 6_3423, 6_3455, 6_3458, 18, 6_3982, 4246, 3940, 1901, 4_7789, 5547, 1_8994], [1_9630, 1100, 6_3446, 1342, 633, 544, 4488, 593, 5102, 2416, 6_3495, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1652, 428, 268, 1936, 515, 268, 5_8593, 2_2413, 9106, 546, 268, 3_3213, 6_3979, 698, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [5_5130, 6_3450, 924, 6_3449, 2249, 4062, 1558, 318, 6_3504, 2_1498, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [509, 377, 2827, 2559, 332, 6575, 6_3443, 2_6801, 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=__magic_name__ ,							model_name='''AI-Sweden/gpt-sw3-126m''' ,							sequences=__magic_name__ ,							)
