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


UpperCAmelCase__          = 16
UpperCAmelCase__          = 32
def 					_UpperCAmelCase      ( __lowerCamelCase   :       Accelerator ,					__lowerCamelCase   :       int = 16		)     ->					Optional[int]:
  _snake_case					     =       AutoTokenizer.from_pretrained('''bert-base-cased'''		)
  _snake_case					     =       load_dataset('''glue''' ,					'''mrpc'''		)

  def tokenize_function(__lowerCamelCase   :       str		):
    # max_length=None => use the model max length (it's actually the default)
    _snake_case					     =       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():
    _snake_case					     =       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
  _snake_case					     =       tokenized_datasets.rename_column('''label''' ,					'''labels'''		)

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

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

  # Instantiate dataloaders.
  _snake_case					     =       DataLoader(
      tokenized_datasets['''train'''] ,					shuffle=__lowerCamelCase ,					collate_fn=__lowerCamelCase ,					batch_size=__lowerCamelCase		)
  _snake_case					     =       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 					_UpperCAmelCase      ( __lowerCamelCase   :       Optional[int] ,					__lowerCamelCase   :       Union[str, Any]		)     ->					List[str]:
  # For testing only
  if os.environ.get('''TESTING_MOCKED_DATALOADERS''' ,					__lowerCamelCase		) == "1":
    _snake_case					     =       2
  # New Code #
  _snake_case					     =       int(args.gradient_accumulation_steps		)
  _snake_case					     =       int(args.local_sgd_steps		)
  # Initialize accelerator
  _snake_case					     =       Accelerator(
      cpu=args.cpu ,					mixed_precision=args.mixed_precision ,					gradient_accumulation_steps=__lowerCamelCase		)
  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					     =       config['''lr''']
  _snake_case					     =       int(config['''num_epochs''']		)
  _snake_case					     =       int(config['''seed''']		)
  _snake_case					     =       int(config['''batch_size''']		)

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

  set_seed(__lowerCamelCase		)
  _snake_case     ,					_snake_case					     =       get_dataloaders(__lowerCamelCase ,					__lowerCamelCase		)
  # Instantiate the model (we build the model here so that the seed also control new weights initialization)
  _snake_case					     =       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).
  _snake_case					     =       model.to(accelerator.device		)

  # Instantiate optimizer
  _snake_case					     =       AdamW(params=model.parameters() ,					lr=__lowerCamelCase		)

  # Instantiate scheduler
  _snake_case					     =       get_linear_schedule_with_warmup(
      optimizer=__lowerCamelCase ,					num_warmup_steps=1_00 ,					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.
  _snake_case     ,					_snake_case     ,					_snake_case     ,					_snake_case     ,					_snake_case					     =       accelerator.prepare(
      __lowerCamelCase ,					__lowerCamelCase ,					__lowerCamelCase ,					__lowerCamelCase ,					__lowerCamelCase		)

  # Now we train the model
  for epoch in range(__lowerCamelCase		):
    model.train()
    with LocalSGD(
        accelerator=__lowerCamelCase ,					model=__lowerCamelCase ,					local_sgd_steps=__lowerCamelCase ,					enabled=local_sgd_steps is not None		) as local_sgd:
      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		):
          _snake_case					     =       model(**__lowerCamelCase		)
          _snake_case					     =       output.loss
          accelerator.backward(__lowerCamelCase		)
          optimizer.step()
          lr_scheduler.step()
          optimizer.zero_grad()
          # LocalSGD-specific line
          local_sgd.step()

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

    _snake_case					     =       metric.compute()
    # Use accelerator.print to print only on the main process.
    accelerator.print(f'''epoch {epoch}:''' ,					__lowerCamelCase		)
def 					_UpperCAmelCase      ( )     ->					Tuple:
  _snake_case					     =       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(
      '''--local_sgd_steps''' ,					type=__lowerCamelCase ,					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					     =       parser.parse_args()
  _snake_case					     =       {'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
  training_function(__lowerCamelCase ,					__lowerCamelCase		)


if __name__ == "__main__":
 main()





"""simple docstring"""



import os
import time

import numpy as np
import onnxruntime as ort


UpperCAmelCase__          = '1'
UpperCAmelCase__          = '0'
UpperCAmelCase__          = '1'

UpperCAmelCase__          = ort.SessionOptions()
UpperCAmelCase__          = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
print('Create inference session...')
UpperCAmelCase__          = ['TensorrtExecutionProvider', 'CUDAExecutionProvider']
UpperCAmelCase__          = ort.InferenceSession('model.onnx', sess_options=sess_opt, providers=execution_provider)
UpperCAmelCase__          = ort.RunOptions()

UpperCAmelCase__          = 128
UpperCAmelCase__          = 1
UpperCAmelCase__          = np.ones((batch, sequence), dtype=np.intaa)
UpperCAmelCase__          = np.ones((batch, sequence), dtype=np.intaa)
UpperCAmelCase__          = np.ones((batch, sequence), dtype=np.intaa)

print('Warm up phase...')
sess.run(
    None,
    {
        sess.get_inputs()[0].name: input_ids,
        sess.get_inputs()[1].name: attention_mask,
        sess.get_inputs()[2].name: token_type_ids,
    },
    run_options=run_opt,
)

print('Start inference...')
UpperCAmelCase__          = time.time()
UpperCAmelCase__          = 2000
UpperCAmelCase__          = {}
for iter in range(max_iters):
 UpperCAmelCase__          = sess.run(
     None,
     {
         sess.get_inputs()[0].name: input_ids,
         sess.get_inputs()[1].name: attention_mask,
         sess.get_inputs()[2].name: token_type_ids,
     },
     run_options=run_opt,
 )
print('Average Inference Time = {:.3f} ms'.format((time.time() - start_time) * 1000 / max_iters))



