



'''simple docstring'''





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											=						16
_UpperCamelCase											=						32

def        a_     (							_lowerCAmelCase       ,_lowerCAmelCase = 16       )      ->      Any:
  __lowerCamelCase       :       int           = AutoTokenizer.from_pretrained('bert-base-cased'       )
  __lowerCamelCase       :       Optional[int]           = load_dataset('glue'       ,'mrpc'       )

  def tokenize_function(_lowerCAmelCase       ):
    # max_length=None => use the model max length (it's actually the default)
    __lowerCamelCase       :       Optional[int]           = 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       :       List[Any]           = 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       :       Union[str, Any]           = tokenized_datasets.rename_column('label'       ,'labels'       )

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

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

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

  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											=						mocked_dataloaders  # noqa: F811

def        a_     (							_lowerCAmelCase       ,_lowerCAmelCase       )      ->      Union[str, Any]:
  # For testing only
  if os.environ.get('TESTING_MOCKED_DATALOADERS'       ,_lowerCAmelCase       ) == "1":
    __lowerCamelCase       :       Optional[Any]           = 2
  # New Code #
  __lowerCamelCase       :       Dict           = int(args.gradient_accumulation_steps       )
  # Initialize accelerator
  __lowerCamelCase       :       List[Any]           = Accelerator(
      cpu=args.cpu       ,mixed_precision=args.mixed_precision       ,gradient_accumulation_steps=_lowerCAmelCase       )
  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
  __lowerCamelCase       :       List[Any]           = config['lr']
  __lowerCamelCase       :       Any           = int(config['num_epochs']       )
  __lowerCamelCase       :       Tuple           = int(config['seed']       )
  __lowerCamelCase       :       Any           = int(config['batch_size']       )

  __lowerCamelCase       :       int           = evaluate.load('glue'       ,'mrpc'       )

  set_seed(_lowerCAmelCase       )
  __lowerCamelCase       ,__lowerCamelCase       :       int           = get_dataloaders(_lowerCAmelCase       ,_lowerCAmelCase       )
  # Instantiate the model (we build the model here so that the seed also control new weights initialization)
  __lowerCamelCase       :       Optional[Any]           = AutoModelForSequenceClassification.from_pretrained('bert-base-cased'       ,return_dict=_lowerCAmelCase       )

  # 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       :       Dict           = AdamW(params=model.parameters()       ,lr=_lowerCAmelCase       )

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

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

    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       :       Tuple           = model(**_lowerCAmelCase       )
      __lowerCamelCase       :       Any           = outputs.logits.argmax(dim=-1       )
      __lowerCamelCase       ,__lowerCamelCase       :       int           = accelerator.gather_for_metrics((predictions, batch['labels'])       )
      metric.add_batch(
          predictions=_lowerCAmelCase       ,references=_lowerCAmelCase       ,)

    __lowerCamelCase       :       Any           = metric.compute()
    # Use accelerator.print to print only on the main process.
    accelerator.print(F'epoch {epoch}:'       ,_lowerCAmelCase       )

def        a_     (							)      ->      int:
  __lowerCamelCase       :       Any           = 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.'       ,)
  # New Code #
  parser.add_argument(
      '--gradient_accumulation_steps'       ,type=_lowerCAmelCase       ,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.'       )
  __lowerCamelCase       :       Tuple           = parser.parse_args()
  __lowerCamelCase       :       Dict           = {'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
  training_function(_lowerCAmelCase       ,_lowerCAmelCase       )


if __name__ == "__main__":
   main()











'''simple docstring'''





import argparse
from pathlib import Path

from transformers import AutoConfig, AutoTokenizer, RagConfig, RagSequenceForGeneration, RagTokenForGeneration

def        a_     (							_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase = None       ,_lowerCAmelCase = None       ,_lowerCAmelCase = None       ,)      ->      Optional[int]:
  if config_name_or_path is None:
    __lowerCamelCase       :       List[Any]           = 'facebook/rag-token-base' if model_type == 'rag_token' else 'facebook/rag-sequence-base'

  if generator_tokenizer_name_or_path is None:
    __lowerCamelCase       :       Optional[int]           = generator_name_or_path

  if question_encoder_tokenizer_name_or_path is None:
    __lowerCamelCase       :       Tuple           = question_encoder_name_or_path

  __lowerCamelCase       :       Union[str, Any]           = RagTokenForGeneration if model_type == 'rag_token' else RagSequenceForGeneration

  # Save model.
  __lowerCamelCase       :       Tuple           = RagConfig.from_pretrained(_lowerCAmelCase       )
  __lowerCamelCase       :       List[Any]           = AutoConfig.from_pretrained(_lowerCAmelCase       )
  __lowerCamelCase       :       Union[str, Any]           = AutoConfig.from_pretrained(_lowerCAmelCase       )

  __lowerCamelCase       :       Tuple           = gen_config
  __lowerCamelCase       :       List[Any]           = question_encoder_config

  __lowerCamelCase       :       str           = model_class.from_pretrained_question_encoder_generator(
      _lowerCAmelCase       ,_lowerCAmelCase       ,config=_lowerCAmelCase       )
  rag_model.save_pretrained(_lowerCAmelCase       )

  # Sanity check.
  model_class.from_pretrained(_lowerCAmelCase       )

  # Save tokenizers.
  __lowerCamelCase       :       List[str]           = AutoTokenizer.from_pretrained(_lowerCAmelCase       )
  gen_tokenizer.save_pretrained(dest_dir / 'generator_tokenizer/'       )
  __lowerCamelCase       :       Optional[int]           = AutoTokenizer.from_pretrained(_lowerCAmelCase       )
  question_encoder_tokenizer.save_pretrained(dest_dir / 'question_encoder_tokenizer/'       )


if __name__ == "__main__":
   _UpperCamelCase											=						argparse.ArgumentParser()
   parser.add_argument(
       '--model_type',
       choices=['rag_sequence', 'rag_token'],
       required=True,
       type=str,
       help='RAG model type: rag_sequence, rag_token',
   )
   parser.add_argument('--dest', type=str, required=True, help='Path to the output checkpoint directory.')
   parser.add_argument('--generator_name_or_path', type=str, required=True, help='Generator model identifier')
   parser.add_argument(
       '--question_encoder_name_or_path', type=str, required=True, help='Question encoder model identifier'
   )

   parser.add_argument(
       '--generator_tokenizer_name_or_path',
       type=str,
       help='Generator tokenizer identifier, if not specified, resolves to ``generator_name_or_path``',
   )
   parser.add_argument(
       '--question_encoder_tokenizer_name_or_path',
       type=str,
       help='Question encoder tokenizer identifier, if not specified, resolves to ``question_encoder_name_or_path``',
   )
   parser.add_argument(
       '--config_name_or_path',
       type=str,
       help=(
           'Identifier of the model config to use, if not provided, resolves to a base config for a given'
           ' ``model_type``'
       ),
   )

   _UpperCamelCase											=						parser.parse_args()

   _UpperCamelCase											=						Path(args.dest)
   dest_dir.mkdir(exist_ok=True)

   consolidate(
       args.model_type,
       args.generator_name_or_path,
       args.question_encoder_name_or_path,
       dest_dir,
       args.config_name_or_path,
       args.generator_tokenizer_name_or_path,
       args.question_encoder_tokenizer_name_or_path,
   )






