





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


UpperCamelCase__   : Union[str, Any]					  = 16
UpperCamelCase__   : Tuple					  = 32




def    A_(			A				,      A = 16   ):
	UpperCAmelCase_								=     AutoTokenizer.from_pretrained("""bert-base-cased"""   )
	UpperCAmelCase_								=     load_dataset("""glue"""				,      """mrpc"""   )

	def tokenize_function(A   ):
		# max_length=None => use the model max length (it's actually the default)
		UpperCAmelCase_								=     tokenizer(examples["""sentence1"""]				,      examples["""sentence2"""]				,      truncation=A				,      max_length=A   )
		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(
		    A				,      batched=A				,      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(A   ):
		# 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(
		    A				,      padding="""longest"""				,      max_length=A				,      pad_to_multiple_of=A				,      return_tensors="""pt"""				,      )

	# Instantiate dataloaders.
	UpperCAmelCase_								=     DataLoader(
	    tokenized_datasets["""train"""]				,      shuffle=A				,      collate_fn=A				,      batch_size=A   )
	UpperCAmelCase_								=     DataLoader(
	    tokenized_datasets["""validation"""]				,      shuffle=A				,      collate_fn=A				,      batch_size=A   )

	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    A_(			A				,      A   ):
	# For testing only
	if os.environ.get("""TESTING_MOCKED_DATALOADERS"""				,      A   ) == "1":
		UpperCAmelCase_								=     2
	# New Code #
	UpperCAmelCase_								=     int(args.gradient_accumulation_steps   )
	# Initialize accelerator
	UpperCAmelCase_								=     Accelerator(
	    cpu=args.cpu				,      mixed_precision=args.mixed_precision				,      gradient_accumulation_steps=A   )
	if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
		raise NotImplementedError(
		    """Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`"""   )
	# 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"""   )

	set_seed(A   )
	UpperCAmelCase_						,	UpperCAmelCase_								=     get_dataloaders(A				,      A   )
	# 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=A   )

	# 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=A   )

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

	# Now we train the model
	for epoch in range(A   ):
		model.train()
		for step, batch in enumerate(A   ):
			# 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(A   ):
				UpperCAmelCase_								=     model(**A   )
				UpperCAmelCase_								=     output.loss
				accelerator.backward(A   )
				optimizer.step()
				lr_scheduler.step()
				optimizer.zero_grad()

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

		UpperCAmelCase_								=     metric.compute()
		# Use accelerator.print to print only on the main process.
		accelerator.print(f"""epoch {epoch}:"""				,      A   )




def    A_(			):
	UpperCAmelCase_								=     argparse.ArgumentParser(description="""Simple example of training script."""   )
	parser.add_argument(
	    """--mixed_precision"""				,      type=A				,      default=A				,      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=A				,      default=1				,      help="""The number of minibatches to be ran before gradients are accumulated."""				,      )
	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(A				,      A   )


if __name__ == "__main__":
				main()









import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

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


UpperCamelCase__   : List[str]					  = logging.get_logger(__name__)

UpperCamelCase__   : int					  = """▁"""

UpperCamelCase__   : Dict					  = {"""vocab_file""": """sentencepiece.bpe.model""", """monolingual_vocab_file""": """dict.txt"""}

UpperCamelCase__   : Dict					  = {
    """vocab_file""": {
        """vinai/bartpho-syllable""": """https://huggingface.co/vinai/bartpho-syllable/resolve/main/sentencepiece.bpe.model""",
    },
    """monolingual_vocab_file""": {
        """vinai/bartpho-syllable""": """https://huggingface.co/vinai/bartpho-syllable/resolve/main/dict.txt""",
    },
}

UpperCamelCase__   : Optional[int]					  = {"""vinai/bartpho-syllable""": 10_24}







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

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





					def __init__(						self    :  List[str] ,					__lowercase    :  Optional[int] ,					__lowercase    :  List[Any] ,					__lowercase    :  List[str]="<s>" ,					__lowercase    :  str="</s>" ,					__lowercase    :  Optional[Any]="</s>" ,					__lowercase    :  Any="<s>" ,					__lowercase    :  Dict="<unk>" ,					__lowercase    :  int="<pad>" ,					__lowercase    :  str="<mask>" ,					__lowercase    :  Optional[Dict[str, Any]] = None ,					**__lowercase    :  Optional[Any] ,					):






						'''simple docstring'''







						UpperCAmelCase_								=     AddedToken(__lowercase ,					lstrip=__lowercase ,					rstrip=__lowercase		) if isinstance(__lowercase ,					__lowercase		) else mask_token

						UpperCAmelCase_								=     {} if sp_model_kwargs is None else sp_model_kwargs

						super().__init__(
						    bos_token=__lowercase ,					eos_token=__lowercase ,					unk_token=__lowercase ,					sep_token=__lowercase ,					cls_token=__lowercase ,					pad_token=__lowercase ,					mask_token=__lowercase ,					sp_model_kwargs=self.sp_model_kwargs ,					**__lowercase ,					)

						UpperCAmelCase_								=     vocab_file
						UpperCAmelCase_								=     monolingual_vocab_file
						UpperCAmelCase_								=     spm.SentencePieceProcessor(**self.sp_model_kwargs		)
						self.sp_model.Load(str(__lowercase		)		)

						# Load the reduced vocab

						# Keep order of special tokens for backward compatibility
						UpperCAmelCase_								=     {}
						UpperCAmelCase_								=     0
						for token in [bos_token, pad_token, eos_token, unk_token, sep_token, cls_token]:
							if str(__lowercase		) not in self.fairseq_tokens_to_ids:
								UpperCAmelCase_								=     cnt
								cnt += 1
						with open(__lowercase ,					"""r""" ,					encoding="""utf-8"""		) as f:
							for line in f.readlines():
								UpperCAmelCase_								=     line.strip().split()[0]
								UpperCAmelCase_								=     len(self.fairseq_tokens_to_ids		)
						if str(__lowercase		) not in self.fairseq_tokens_to_ids:
							UpperCAmelCase_								=     len(self.fairseq_tokens_to_ids		)

						UpperCAmelCase_								=     {v: k for k, v in self.fairseq_tokens_to_ids.items()}





					def __getstate__(						self    :  List[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.__dict__.copy()
						UpperCAmelCase_								=     None
						UpperCAmelCase_								=     self.sp_model.serialized_model_proto()
						return state





					def __setstate__(						self    :  List[str] ,					__lowercase    :  Dict		):






						'''simple docstring'''







						UpperCAmelCase_								=     d

						# for backward compatibility
						if not hasattr(self ,					"""sp_model_kwargs"""		):
							UpperCAmelCase_								=     {}

						UpperCAmelCase_								=     spm.SentencePieceProcessor(**self.sp_model_kwargs		)
						self.sp_model.LoadFromSerializedProto(self.sp_model_proto		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any] ,					__lowercase    :  List[int] ,					__lowercase    :  Optional[List[int]] = None		):






						'''simple docstring'''







						if token_ids_a is None:
							return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
						UpperCAmelCase_								=     [self.cls_token_id]
						UpperCAmelCase_								=     [self.sep_token_id]
						return cls + token_ids_a + sep + sep + token_ids_a + sep





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any] ,					__lowercase    :  List[int] ,					__lowercase    :  Optional[List[int]] = None ,					__lowercase    :  bool = False		):






						'''simple docstring'''







						if already_has_special_tokens:
							return super().get_special_tokens_mask(
							    token_ids_a=__lowercase ,					token_ids_a=__lowercase ,					already_has_special_tokens=__lowercase		)

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





					def        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  List[int] ,					__lowercase    :  Optional[List[int]] = None		):






						'''simple docstring'''







						UpperCAmelCase_								=     [self.sep_token_id]
						UpperCAmelCase_								=     [self.cls_token_id]

						if token_ids_a is None:
							return len(cls + token_ids_a + sep		) * [0]
						return len(cls + token_ids_a + sep + sep + token_ids_a + sep		) * [0]





					@property
					def        SCREAMING_SNAKE_CASE							(						self    :  int		):






						'''simple docstring'''







						return len(self.fairseq_ids_to_tokens		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Dict		):






						'''simple docstring'''







						UpperCAmelCase_								=     {self.convert_ids_to_tokens(__lowercase		): i for i in range(self.vocab_size		)}
						vocab.update(self.added_tokens_encoder		)
						return vocab





					def        SCREAMING_SNAKE_CASE							(						self    :  str ,					__lowercase    :  str		):






						'''simple docstring'''







						return self.sp_model.encode(__lowercase ,					out_type=__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple ,					__lowercase    :  Tuple		):






						'''simple docstring'''







						if token in self.fairseq_tokens_to_ids:
							return self.fairseq_tokens_to_ids[token]
						else:
							return self.unk_token_id





					def        SCREAMING_SNAKE_CASE							(						self    :  Dict ,					__lowercase    :  Optional[Any]		):






						'''simple docstring'''







						return self.fairseq_ids_to_tokens[index]





					def        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  Dict		):






						'''simple docstring'''







						UpperCAmelCase_								=     """""".join(__lowercase		).replace(__lowercase ,					""" """		).strip()
						return out_string





					def        SCREAMING_SNAKE_CASE							(						self    :  str ,					__lowercase    :  str ,					__lowercase    :  Optional[str] = None		):






						'''simple docstring'''







						if not os.path.isdir(__lowercase		):
							logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""		)
							return
						UpperCAmelCase_								=     os.path.join(
						    __lowercase ,					(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]		)
						UpperCAmelCase_								=     os.path.join(
						    __lowercase ,					(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""monolingual_vocab_file"""] ,					)

						if os.path.abspath(self.vocab_file		) != os.path.abspath(__lowercase		) and os.path.isfile(self.vocab_file		):
							copyfile(self.vocab_file ,					__lowercase		)
						elif not os.path.isfile(self.vocab_file		):
							with open(__lowercase ,					"""wb"""		) as fi:
								UpperCAmelCase_								=     self.sp_model.serialized_model_proto()
								fi.write(__lowercase		)

						if os.path.abspath(self.monolingual_vocab_file		) != os.path.abspath(
						    __lowercase		) and os.path.isfile(self.monolingual_vocab_file		):
							copyfile(self.monolingual_vocab_file ,					__lowercase		)
						elif not os.path.isfile(self.monolingual_vocab_file		):
							with open(__lowercase ,					"""w""" ,					encoding="""utf-8"""		) as fp:
								for token in self.fairseq_tokens_to_ids:
									if token not in self.all_special_tokens:
										fp.write(F"""{str(__lowercase		)} \n"""		)

						return out_vocab_file, out_monolingual_vocab_file


