




import argparse
import gc
import os
import threading

import evaluate
import psutil
import torch
from datasets import load_dataset
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
from torch.utils.data import DataLoader
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
    set_seed,
)

from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin


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


__SCREAMING_SNAKE_CASE			:      Dict      			=16
__SCREAMING_SNAKE_CASE			:      List[str]      			=32

def 	lowerCAmelCase_	(     lowerCamelCase:				Dict  )    ->	List[Any]:


      '''simple docstring'''




      return int(x / 2**20  )






class 	A	:
     """simple docstring"""

     def __enter__(			self     :      str	) ->				Dict:
           gc.collect()
           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
           lowerCamelCase__					     = torch.cuda.memory_allocated()
           lowerCamelCase__					     = psutil.Process()

           lowerCamelCase__					     = self.cpu_mem_used()
           lowerCamelCase__					     = True
           lowerCamelCase__					     = threading.Thread(target=self.peak_monitor_func	)
           lowerCamelCase__					     = True
           peak_monitor_thread.start()
           return self




     def        __snake_case      (			self     :      Optional[Any]	) ->				Optional[Any]:
           return self.process.memory_info().rss




     def        __snake_case      (			self     :      Dict	) ->				Tuple:
           lowerCamelCase__					     = -1

           while True:
                 lowerCamelCase__					     = max(self.cpu_mem_used()		,     self.cpu_peak	)

                 # can't sleep or will not catch the peak right (this comment is here on purpose)
                 # time.sleep(0.001) # 1msec

                 if not self.peak_monitoring:
                       break




     def __exit__(			self     :      Tuple		,     *A__     :      List[str]	) ->				Any:
           lowerCamelCase__					     = False

           gc.collect()
           torch.cuda.empty_cache()
           lowerCamelCase__					     = torch.cuda.memory_allocated()
           lowerCamelCase__					     = torch.cuda.max_memory_allocated()
           lowerCamelCase__					     = bamb(self.end - self.begin	)
           lowerCamelCase__					     = bamb(self.peak - self.begin	)

           lowerCamelCase__					     = self.cpu_mem_used()
           lowerCamelCase__					     = bamb(self.cpu_end - self.cpu_begin	)
           lowerCamelCase__					     = bamb(self.cpu_peak - self.cpu_begin	)
           # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


# For testing only
if os.environ.get("""TESTING_MOCKED_DATALOADERS""", None) == "1":
      from accelerate.test_utils.training import mocked_dataloaders

      __SCREAMING_SNAKE_CASE			:      List[Any]      			=mocked_dataloaders  # noqa: F811

def 	lowerCAmelCase_	(     lowerCamelCase:				Union[str, Any]   , lowerCamelCase:				Optional[int]  )    ->	Any:


      '''simple docstring'''




      if os.environ.get('TESTING_MOCKED_DATALOADERS'   , lowerCamelCase  ) == "1":
            lowerCamelCase__					     = 2

      # New Code #
      # Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
      lowerCamelCase__					     = FullyShardedDataParallelPlugin(
          state_dict_config=FullStateDictConfig(offload_to_cpu=lowerCamelCase   , ranka_only=lowerCamelCase  )   , optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=lowerCamelCase   , ranka_only=lowerCamelCase  )   , )

      # Initialize accelerator
      if args.with_tracking:
            lowerCamelCase__					     = Accelerator(
                cpu=args.cpu   , mixed_precision=args.mixed_precision   , log_with='wandb'   , project_dir=args.logging_dir   , fsdp_plugin=lowerCamelCase   , )
      else:
            lowerCamelCase__					     = Accelerator(fsdp_plugin=lowerCamelCase  )
      accelerator.print(accelerator.distributed_type  )

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

      # We need to initialize the trackers we use, and also store our configuration
      if args.with_tracking:
            lowerCamelCase__					     = vars(lowerCamelCase  )
            accelerator.init_trackers('fsdp_glue_no_trainer'   , lowerCamelCase  )

      lowerCamelCase__					     = AutoTokenizer.from_pretrained(args.model_name_or_path  )
      lowerCamelCase__					     = load_dataset('glue'   , 'mrpc'  )
      lowerCamelCase__					     = evaluate.load('glue'   , 'mrpc'  )

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

      # If the batch size is too big we use gradient accumulation
      lowerCamelCase__					     = 1
      if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
            lowerCamelCase__					     = batch_size // MAX_GPU_BATCH_SIZE
            lowerCamelCase__					     = MAX_GPU_BATCH_SIZE

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

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

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

      set_seed(lowerCamelCase  )

      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      lowerCamelCase__					     = AutoModelForSequenceClassification.from_pretrained(
          args.model_name_or_path   , return_dict=lowerCamelCase   , low_cpu_mem_usage=lowerCamelCase  )

      # New Code #
      # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
      lowerCamelCase__					     = accelerator.prepare(lowerCamelCase  )
      accelerator.print(lowerCamelCase  )

      # Instantiate optimizer
      # New Code #
      # For FSDP feature, at present it doesn't support multiple parameter groups,
      # so we need to create a single parameter group for the whole model
      lowerCamelCase__					     = torch.optim.AdamW(params=model.parameters()   , lr=lowerCamelCase   , weight_decay=2e-4  )

      # Instantiate scheduler
      lowerCamelCase__					     = get_linear_schedule_with_warmup(
          optimizer=lowerCamelCase   , num_warmup_steps=10   , num_training_steps=(len(lowerCamelCase  ) * num_epochs) // gradient_accumulation_steps   , )

      # New Code #
      # For FSDP feature, prepare everything except the model as we have already prepared the model
      # before creating the optimizer
      # 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__					     = accelerator.prepare(
          lowerCamelCase   , lowerCamelCase   , lowerCamelCase   , lowerCamelCase  )

      lowerCamelCase__					     = 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__					     = os.path.basename(args.resume_from_checkpoint  )
            else:
                  # Get the most recent checkpoint
                  lowerCamelCase__					     = [f.name for f in os.scandir(os.getcwd()  ) if f.is_dir()]
                  dirs.sort(key=os.path.getctime  )
                  lowerCamelCase__					     = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
            # Extract `epoch_{i}` or `step_{i}`
            lowerCamelCase__					     = os.path.splitext(lowerCamelCase  )[0]

            if "epoch" in training_difference:
                  num_epochs -= int(training_difference.replace('epoch_'   , ''  )  )
                  lowerCamelCase__					     = None
            else:
                  lowerCamelCase__					     = int(training_difference.replace('step_'   , ''  )  )
                  num_epochs -= resume_step // len(lowerCamelCase  )
                  # If resuming by step, we also need to know exactly how far into the DataLoader we went
                  lowerCamelCase__					     = (num_epochs * len(lowerCamelCase  )) - resume_step

    # Now we train the model
      for epoch in range(lowerCamelCase  ):
            # New Code #
            # context manager to track the peak memory usage during the training epoch
            with TorchTracemalloc() as tracemalloc:
                  model.train()
                  if args.with_tracking:
                        lowerCamelCase__					     = 0
                  for step, batch in enumerate(lowerCamelCase  ):
                        # We need to skip steps until we reach the resumed step
                        if args.resume_from_checkpoint and epoch == 0:
                              if resume_step is not None and step < resume_step:
                                    pass
                # We could avoid this line since we set the accelerator with `device_placement=True`.
                        batch.to(accelerator.device  )
                        lowerCamelCase__					     = model(**lowerCamelCase  )
                        lowerCamelCase__					     = outputs.loss
                        lowerCamelCase__					     = loss / gradient_accumulation_steps
                        # We keep track of the loss at each epoch
                        if args.with_tracking:
                              total_loss += loss.detach().float()
                        accelerator.backward(lowerCamelCase  )
                        if step % gradient_accumulation_steps == 0:
                              optimizer.step()
                              lr_scheduler.step()
                              optimizer.zero_grad()
                              # accelerator.print(lr_scheduler.get_lr())

                        overall_step += 1

                        if isinstance(lowerCamelCase   , lowerCamelCase  ):
                              lowerCamelCase__					     = f"""step_{overall_step}"""
                              if overall_step % checkpointing_steps == 0:
                                    if args.output_dir is not None:
                                          lowerCamelCase__					     = os.path.join(args.output_dir   , lowerCamelCase  )
                                    accelerator.save_state(lowerCamelCase  )
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
            accelerator.print('Memory before entering the train : {}'.format(bamb(tracemalloc.begin  )  )  )
            accelerator.print('Memory consumed at the end of the train (end-begin): {}'.format(tracemalloc.used  )  )
            accelerator.print('Peak Memory consumed during the train (max-begin): {}'.format(tracemalloc.peaked  )  )
            accelerator.print(
                'Total Peak Memory consumed during the train (max): {}'.format(
                    tracemalloc.peaked + bamb(tracemalloc.begin  )  )  )
            # Logging the peak memory usage of the GPU to the tracker
            if args.with_tracking:
                  accelerator.log(
                      {
                          'train_total_peak_memory': tracemalloc.peaked + bamb(tracemalloc.begin  ),
                      }   , step=lowerCamelCase   , )

            # New Code #
            # context manager to track the peak memory usage during the evaluation
            with TorchTracemalloc() as tracemalloc:
                  model.eval()
                  for step, batch in enumerate(lowerCamelCase  ):
                        # We could avoid this line since we set the accelerator with `device_placement=True`.
                        batch.to(accelerator.device  )
                        with torch.no_grad():
                              lowerCamelCase__					     = model(**lowerCamelCase  )
                        lowerCamelCase__					     = outputs.logits.argmax(dim=-1  )
                        lowerCamelCase__     ,						lowerCamelCase__					     = accelerator.gather_for_metrics((predictions, batch['labels'])  )
                        metric.add_batch(
                            predictions=lowerCamelCase   , references=lowerCamelCase   , )

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

                  if checkpointing_steps == "epoch":
                        lowerCamelCase__					     = f"""epoch_{epoch}"""
                        if args.output_dir is not None:
                              lowerCamelCase__					     = os.path.join(args.output_dir   , lowerCamelCase  )
                        accelerator.save_state(lowerCamelCase  )
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
            accelerator.print('Memory before entering the eval : {}'.format(bamb(tracemalloc.begin  )  )  )
            accelerator.print('Memory consumed at the end of the eval (end-begin): {}'.format(tracemalloc.used  )  )
            accelerator.print('Peak Memory consumed during the eval (max-begin): {}'.format(tracemalloc.peaked  )  )
            accelerator.print(
                'Total Peak Memory consumed during the eval (max): {}'.format(tracemalloc.peaked + bamb(tracemalloc.begin  )  )  )
            # Logging the peak memory usage of the GPU to the tracker
            if args.with_tracking:
                  accelerator.log(
                      {
                          'eval_total_peak_memory': tracemalloc.peaked + bamb(tracemalloc.begin  ),
                      }   , step=lowerCamelCase   , )

      if args.with_tracking:
            accelerator.end_training()

def 	lowerCAmelCase_	(     )    ->	int:


      '''simple docstring'''




      lowerCamelCase__					     = argparse.ArgumentParser(description='Simple example of training script.'  )
      parser.add_argument(
          '--mixed_precision'   , type=lowerCamelCase   , default=lowerCamelCase   , 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=lowerCamelCase   , default=lowerCamelCase   , 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=lowerCamelCase   , default=lowerCamelCase   , 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=lowerCamelCase   , default='.'   , help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'   , )
      parser.add_argument(
          '--logging_dir'   , type=lowerCamelCase   , default='logs'   , help='Location on where to store experiment tracking logs`'   , )
      parser.add_argument(
          '--model_name_or_path'   , type=lowerCamelCase   , help='Path to pretrained model or model identifier from huggingface.co/models.'   , required=lowerCamelCase   , )
      lowerCamelCase__					     = parser.parse_args()
      lowerCamelCase__					     = {'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
      training_function(lowerCamelCase   , lowerCamelCase  )


if __name__ == "__main__":
      main()



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

from ...tokenization_utils_fast import AddedToken, PreTrainedTokenizerFast
from ...utils import logging


snake_case_  :	Tuple							  =							logging.get_logger(__name__)

snake_case_  :	int							  =							{'''vocab_file''': '''spiece.model''', '''tokenizer_file''': '''tokenizer.json'''}

snake_case_  :	Dict							  =							{
    '''vocab_file''': {
        '''TsinghuaAI/CPM-Generate''': '''https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/spiece.model''',
    },
    '''tokenizer_file''': {
        '''TsinghuaAI/CPM-Generate''': '''https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/tokenizer.json''',
    },
}

class   A__							(						UpperCamelCase__     ):




						def __init__(				self    :					Optional[Any] ,			_a    :					Tuple=None ,			_a    :					Union[str, Any]=None ,			_a    :					Dict=False ,			_a    :					int=True ,			_a    :					str=False ,			_a    :					int="<s>" ,			_a    :					int="</s>" ,			_a    :					Optional[int]="<unk>" ,			_a    :					Dict="<sep>" ,			_a    :					Optional[int]="<pad>" ,			_a    :					str="<cls>" ,			_a    :					List[Any]="<mask>" ,			_a    :					Optional[Any]=["<eop>", "<eod>"] ,			**_a    :					Any ,			)							->   Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =AddedToken(_a ,			lstrip=_a ,			rstrip=_a     ) if isinstance(_a ,			_a     ) else mask_token

										super().__init__(
										    vocab_file=_a ,			tokenizer_file=_a ,			do_lower_case=_a ,			remove_space=_a ,			keep_accents=_a ,			bos_token=_a ,			eos_token=_a ,			unk_token=_a ,			sep_token=_a ,			pad_token=_a ,			cls_token=_a ,			mask_token=_a ,			additional_special_tokens=_a ,			**_a ,			)

										_SCREAMING_SNAKE_CASE	     =3
										_SCREAMING_SNAKE_CASE	     =do_lower_case
										_SCREAMING_SNAKE_CASE	     =remove_space
										_SCREAMING_SNAKE_CASE	     =keep_accents
										_SCREAMING_SNAKE_CASE	     =vocab_file
										_SCREAMING_SNAKE_CASE	     =False if not self.vocab_file else True

										try:
														import jieba
										except ModuleNotFoundError as error:
														raise error.__class__(
														    '''You need to install jieba to use CpmTokenizer or CpmTokenizerFast. '''
														    '''See https://pypi.org/project/jieba/ for installation.'''     )
										_SCREAMING_SNAKE_CASE	     =jieba
										_SCREAMING_SNAKE_CASE	     =str.maketrans(''' \n''' ,			'''\u2582\u2583'''     )




						def 							__UpperCamelCase (				self    :					Any ,			_a    :					List[int] ,			_a    :					Optional[List[int]] = None     )							->   List[int]:


										"""simple docstring"""





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




						def 							__UpperCamelCase (				self    :					Dict ,			_a    :					List[int] ,			_a    :					Optional[List[int]] = None     )							->   List[int]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =[self.sep_token_id]
										_SCREAMING_SNAKE_CASE	     =[2]

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




						def 							__UpperCamelCase (				self    :					str ,			_a    :					str ,			_a    :					Optional[str] = None     )							->   Tuple[str]:


										"""simple docstring"""





										if not self.can_save_slow_tokenizer:
														raise ValueError(
														    '''Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '''
														    '''tokenizer.'''     )

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

										if os.path.abspath(self.vocab_file     ) != os.path.abspath(_a     ):
														copyfile(self.vocab_file ,			_a     )

										return (out_vocab_file,)




						def 							__UpperCamelCase (				self    :					Dict ,			_a    :					str ,			*_a    :					Optional[int] ,			**_a    :					Any     )							->   int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =[
										    ''' '''.join([x.translate(self.translator     ) for x in self.jieba.cut(_a ,			cut_all=_a     )]     )
										    for text in batch_text_or_text_pairs
										]
										return super()._batch_encode_plus(_a ,			*_a ,			**_a     )




						def 							__UpperCamelCase (				self    :					List[Any] ,			*_a    :					Union[str, Any] ,			**_a    :					Any     )							->   Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =super()._decode(*_a ,			**_a     )
										_SCREAMING_SNAKE_CASE	     =text.replace(''' ''' ,			''''''     ).replace('''\u2582''' ,			''' '''     ).replace('''\u2583''' ,			'''\n'''     )
										return text