class_docstr: ORTTrainer is a simple but feature-complete training and eval loop for ONNX Runtime, optimized for 🤗 Transformers.
  
  Args:
      model ([`~transformers.PreTrainedModel`] or `torch.nn.Module`, *optional*):
          The model to train, evaluate or use for predictions. If not provided, a `model_init` must be passed.
  
          <Tip>
  
          [`ORTTrainer`] is optimized to work with the [`~transformers.PreTrainedModel`] provided by the transformers library.
          You can still use your own models defined as `torch.nn.Module` for training with ONNX Runtime backend
          and inference with PyTorch backend as long as they work the same way as the 🤗 Transformers models.
  
          </Tip>
  
      args ([`ORTTrainingArguments`], *optional*):
          The arguments to tweak for training. Will default to a basic instance of [`ORTTrainingArguments`] with the
          `output_dir` set to a directory named *tmp_trainer* in the current directory if not provided.
      data_collator ([`~transformers.DataCollator`], *optional*):
          The function to use to form a batch from a list of elements of `train_dataset` or `eval_dataset`. Will
          default to [`~transformers.default_data_collator`] if no `tokenizer` is provided, an instance of
          [`~transformers.DataCollatorWithPadding`] otherwise.
      train_dataset (`torch.utils.data.Dataset` or `torch.utils.data.IterableDataset`, *optional*):
          The dataset to use for training. If it is a [`~datasets.Dataset`], columns not accepted by the
          `model.forward()` method are automatically removed.
          Note that if it's a `torch.utils.data.IterableDataset` with some randomization and you are training in a
          distributed fashion, your iterable dataset should either use a internal attribute `generator` that is a
          `torch.Generator` for the randomization that must be identical on all processes (and the ORTTrainer will
          manually set the seed of this `generator` at each epoch) or have a `set_epoch()` method that internally
          sets the seed of the RNGs used.
      eval_dataset (Union[`torch.utils.data.Dataset`, Dict[str, `torch.utils.data.Dataset`]), *optional*):
          The dataset to use for evaluation. If it is a [`~datasets.Dataset`], columns not accepted by the
          `model.forward()` method are automatically removed. If it is a dictionary, it will evaluate on each
          dataset prepending the dictionary key to the metric name.
      tokenizer ([`~transformers.PreTrainedTokenizerBase`], *optional*):
          The tokenizer used to preprocess the data. If provided, will be used to automatically pad the inputs the
          maximum length when batching inputs, and it will be saved along the model to make it easier to rerun an
          interrupted training or reuse the fine-tuned model.
      model_init (`Callable[[], PreTrainedModel]`, *optional*):
          A function that instantiates the model to be used. If provided, each call to [`ORTTrainer.train`] will start
          from a new instance of the model as given by this function.
          The function may have zero argument, or a single one containing the optuna/Ray Tune/SigOpt trial object, to
          be able to choose different architectures according to hyper parameters (such as layer count, sizes of
          inner layers, dropout probabilities etc).
      compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):
          The function that will be used to compute metrics at evaluation. Must take a [`EvalPrediction`] and return
          a dictionary string to metric values.
      callbacks (List of [`TrainerCallback`], *optional*):
          A list of callbacks to customize the training loop. Will add those to the list of default callbacks
          detailed in [here](callback).
          If you want to remove one of the default callbacks used, use the [`ORTTrainer.remove_callback`] method.
      optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*): A tuple
          containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your model
          and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`.
      preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*):
          A function that preprocess the logits right before caching them at each evaluation step. Must take two
          tensors, the logits and the labels, and return the logits once processed as desired. The modifications made
          by this function will be reflected in the predictions received by `compute_metrics`.
          Note that the labels (second parameter) will be `None` if the dataset does not have them.
  Important attributes:
      - **model** -- Always points to the core model. If using a transformers model, it will be a [`~transformers.PreTrainedModel`]
      subclass.
      - **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the
      original model. This is the model that should be used for the forward pass. For example, under `DeepSpeed`,
      the inner model is first wrapped in `ORTModule` and then in `DeepSpeed` and then again in
      `torch.nn.DistributedDataParallel`. If the inner model hasn't been wrapped, then `self.model_wrapped` is the
      same as `self.model`.
      - **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from
      data parallelism, this means some of the model layers are split on different GPUs).
      - **place_model_on_device** -- Whether or not to automatically place the model on the device - it will be set
      to `False` if model parallel or deepspeed is used, or if the default
      `ORTTrainingArguments.place_model_on_device` is overridden to return `False` .
      - **is_in_train** -- Whether or not a model is currently running `train` (e.g. when `evaluate` is called while
      in `train`)
class_docstr_tok: ["ORTTrainer","is","a","simple","but","feature-complete","training","and","eval","loop","for","ONNX","Runtime",",","optimized","for","\ud83e\udd17","Transformers",".","Args",":","model","(","[","`","~transformers.PreTrainedModel","`","]","or","`","torch.nn.Module","`",",","*","optional","*",")",":","The","model","to","train",",","evaluate","or","use","for","predictions",".","If","not","provided",",","a","`","model_init","`","must","be","passed",".","<","Tip",">","[","`","ORTTrainer","`","]","is","optimized","to","work","with","the","[","`","~transformers.PreTrainedModel","`","]","provided","by","the","transformers","library",".","You","can","still","use","your","own","models","defined","as","`","torch.nn.Module","`","for","training","with","ONNX","Runtime","backend","and","inference","with","PyTorch","backend","as","long","as","they","work","the","same","way","as","the","\ud83e\udd17","Transformers","models",".","<","\/Tip",">","args","(","[","`","ORTTrainingArguments","`","]",",","*","optional","*",")",":","The","arguments","to","tweak","for","training",".","Will","default","to","a","basic","instance","of","[","`","ORTTrainingArguments","`","]","with","the","`","output_dir","`","set","to","a","directory","named","*","tmp_trainer","*","in","the","current","directory","if","not","provided",".","data_collator","(","[","`","~transformers.DataCollator","`","]",",","*","optional","*",")",":","The","function","to","use","to","form","a","batch","from","a","list","of","elements","of","`","train_dataset","`","or","`","eval_dataset","`",".","Will","default","to","[","`","~transformers.default_data_collator","`","]","if","no","`","tokenizer","`","is","provided",",","an","instance","of","[","`","~transformers.DataCollatorWithPadding","`","]","otherwise",".","train_dataset","(","`","torch.utils.data.Dataset","`","or","`","torch.utils.data.IterableDataset","`",",","*","optional","*",")",":","The","dataset","to","use","for","training",".","If","it","is","a","[","`","~datasets.Dataset","`","]",",","columns","not","accepted","by","the","`","model.forward","(",")","`","method","are","automatically","removed",".","Note","that","if","it","'s","a","`","torch.utils.data.IterableDataset","`","with","some","randomization","and","you","are","training","in","a","distributed","fashion",",","your","iterable","dataset","should","either","use","a","internal","attribute","`","generator","`","that","is","a","`","torch.Generator","`","for","the","randomization","that","must","be","identical","on","all","processes","(","and","the","ORTTrainer","will","manually","set","the","seed","of","this","`","generator","`","at","each","epoch",")","or","have","a","`","set_epoch","(",")","`","method","that","internally","sets","the","seed","of","the","RNGs","used",".","eval_dataset","(","Union","[","`","torch.utils.data.Dataset","`",",","Dict","[","str",",","`","torch.utils.data.Dataset","`","]",")",",","*","optional","*",")",":","The","dataset","to","use","for","evaluation",".","If","it","is","a","[","`","~datasets.Dataset","`","]",",","columns","not","accepted","by","the","`","model.forward","(",")","`","method","are","automatically","removed",".","If","it","is","a","dictionary",",","it","will","evaluate","on","each","dataset","prepending","the","dictionary","key","to","the","metric","name",".","tokenizer","(","[","`","~transformers.PreTrainedTokenizerBase","`","]",",","*","optional","*",")",":","The","tokenizer","used","to","preprocess","the","data",".","If","provided",",","will","be","used","to","automatically","pad","the","inputs","the","maximum","length","when","batching","inputs",",","and","it","will","be","saved","along","the","model","to","make","it","easier","to","rerun","an","interrupted","training","or","reuse","the","fine-tuned","model",".","model_init","(","`","Callable","[","[","]",",","PreTrainedModel","]","`",",","*","optional","*",")",":","A","function","that","instantiates","the","model","to","be","used",".","If","provided",",","each","call","to","[","`","ORTTrainer.train","`","]","will","start","from","a","new","instance","of","the","model","as","given","by","this","function",".","The","function","may","have","zero","argument",",","or","a","single","one","containing","the","optuna\/Ray","Tune\/SigOpt","trial","object",",","to","be","able","to","choose","different","architectures","according","to","hyper","parameters","(","such","as","layer","count",",","sizes","of","inner","layers",",","dropout","probabilities","etc",")",".","compute_metrics","(","`","Callable","[","[","EvalPrediction","]",",","Dict","]","`",",","*","optional","*",")",":","The","function","that","will","be","used","to","compute","metrics","at","evaluation",".","Must","take","a","[","`","EvalPrediction","`","]","and","return","a","dictionary","string","to","metric","values",".","callbacks","(","List","of","[","`","TrainerCallback","`","]",",","*","optional","*",")",":","A","list","of","callbacks","to","customize","the","training","loop",".","Will","add","those","to","the","list","of","default","callbacks","detailed","in","[","here","]","(","callback",")",".","If","you","want","to","remove","one","of","the","default","callbacks","used",",","use","the","[","`","ORTTrainer.remove_callback","`","]","method",".","optimizers","(","`","Tuple","[","torch.optim.Optimizer",",","torch.optim.lr_scheduler.LambdaLR","]","`",",","*","optional","*",")",":","A","tuple","containing","the","optimizer","and","the","scheduler","to","use",".","Will","default","to","an","instance","of","[","`","AdamW","`","]","on","your","model","and","a","scheduler","given","by","[","`","get_linear_schedule_with_warmup","`","]","controlled","by","`","args","`",".","preprocess_logits_for_metrics","(","`","Callable","[","[","torch.Tensor",",","torch.Tensor","]",",","torch.Tensor","]","`",",","*","optional","*",")",":","A","function","that","preprocess","the","logits","right","before","caching","them","at","each","evaluation","step",".","Must","take","two","tensors",",","the","logits","and","the","labels",",","and","return","the","logits","once","processed","as","desired",".","The","modifications","made","by","this","function","will","be","reflected","in","the","predictions","received","by","`","compute_metrics","`",".","Note","that","the","labels","(","second","parameter",")","will","be","`","None","`","if","the","dataset","does","not","have","them",".","Important","attributes",":","-","*","*","model","*","*","--","Always","points","to","the","core","model",".","If","using","a","transformers","model",",","it","will","be","a","[","`","~transformers.PreTrainedModel","`","]","subclass",".","-","*","*","model_wrapped","*","*","--","Always","points","to","the","most","external","model","in","case","one","or","more","other","modules","wrap","the","original","model",".","This","is","the","model","that","should","be","used","for","the","forward","pass",".","For","example",",","under","`","DeepSpeed","`",",","the","inner","model","is","first","wrapped","in","`","ORTModule","`","and","then","in","`","DeepSpeed","`","and","then","again","in","`","torch.nn.DistributedDataParallel","`",".","If","the","inner","model","has","n't","been","wrapped",",","then","`","self.model_wrapped","`","is","the","same","as","`","self.model","`",".","-","*","*","is_model_parallel","*","*","--","Whether","or","not","a","model","has","been","switched","to","a","model","parallel","mode","(","different","from","data","parallelism",",","this","means","some","of","the","model","layers","are","split","on","different","GPUs",")",".","-","*","*","place_model_on_device","*","*","--","Whether","or","not","to","automatically","place","the","model","on","the","device","-","it","will","be","set","to","`","False","`","if","model","parallel","or","deepspeed","is","used",",","or","if","the","default","`","ORTTrainingArguments.place_model_on_device","`","is","overridden","to","return","`","False","`",".","-","*","*","is_in_train","*","*","--","Whether","or","not","a","model","is","currently","running","`","train","`","(","e.g",".","when","`","evaluate","`","is","called","while","in","`","train","`",")"]
docstr: Setup the optimizer.
  
  We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
  ORTTrainer's init through `optimizers`, or subclass and override this method in a subclass.
docstr_tok: ["Setup","the","optimizer",".","We","provide","a","reasonable","default","that","works","well",".","If","you","want","to","use","something","else",",","you","can","pass","a","tuple","in","the","ORTTrainer","'s","init","through","`","optimizers","`",",","or","subclass","and","override","this","method","in","a","subclass","."]
code: def create_optimizer(self):
      """
      Setup the optimizer.
  
      We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
      ORTTrainer's init through `optimizers`, or subclass and override this method in a subclass.
      """
      opt_model = (
          self.model_wrapped
          if is_sagemaker_mp_enabled()
          else self.model
      )
  
      if self.optimizer is None:
          decay_parameters = get_parameter_names(
              opt_model, [nn.LayerNorm]
          )
          decay_parameters = [
              name for name in decay_parameters if "bias" not in name
          ]
          optimizer_grouped_parameters = [
              {
                  "params": [
                      p
                      for n, p in opt_model.named_parameters()
                      if n in decay_parameters
                  ],
                  "weight_decay": self.args.weight_decay,
              },
              {
                  "params": [
                      p
                      for n, p in opt_model.named_parameters()
                      if n not in decay_parameters
                  ],
                  "weight_decay": 0.0,
              },
          ]
  
          if self.args.optim in ORTOptimizerNames:
              (
                  optimizer_cls,
                  optimizer_kwargs,
              ) = ORTTrainer.get_ort_optimizer_cls_and_kwargs(self.args)
          else:
              (
                  optimizer_cls,
                  optimizer_kwargs,
              ) = Trainer.get_optimizer_cls_and_kwargs(self.args)
  
          self.optimizer = optimizer_cls(
              optimizer_grouped_parameters, **optimizer_kwargs
          )
          if optimizer_cls.__name__ == "Adam8bit":
              import bitsandbytes
  
              manager = (
                  bitsandbytes.optim.GlobalOptimManager.get_instance()
              )
  
              skipped = 0
              for module in opt_model.modules():
                  if isinstance(module, nn.Embedding):
                      skipped += sum(
                          {
                              p.data_ptr(): p.numel()
                              for p in module.parameters()
                          }.values()
                      )
                      logger.info(
                          f"skipped {module}: {skipped/2**20}M params"
                      )
                      manager.register_module_override(
                          module, "weight", {"optim_bits": 32}
                      )
                      logger.debug(
                          f"bitsandbytes: will optimize {module} in fp32"
                      )
              logger.info(f"skipped: {skipped/2**20}M params")
  
      if is_sagemaker_mp_enabled():
          raise NotImplementedError(
              "Sagemaker's distributed data parallel features are not supported by `ORTTrainer` yet."
          )
  
      return self.optimizer
  
imports: import functools
  import math
  import os
  import shutil
  import sys
  import time
  import types
  import warnings
  from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
  from transformers.integrations import hp_params
  from transformers.utils import is_accelerate_available
  from packaging import version
  import huggingface_hub.utils
  import torch
  import torch.distributed
  from torch import nn
  from torch.utils.data import Dataset, RandomSampler
  from transformers.data.data_collator import DataCollator
  from transformers.debug_utils import DebugOption, DebugUnderflowOverflow
  from transformers.deepspeed import deepspeed_init, deepspeed_load_checkpoint, is_deepspeed_zero3_enabled
  from transformers.modeling_utils import PreTrainedModel, unwrap_model
  from transformers.tokenization_utils_base import PreTrainedTokenizerBase
  from transformers.trainer import Trainer
  from transformers.trainer_callback import TrainerCallback, TrainerState
  from transformers.trainer_pt_utils import get_model_param_count, get_module_class_from_name, get_parameter_names
  from transformers.trainer_utils import EvalPrediction, HPSearchBackend, TrainOutput, enable_full_determinism, find_executable_batch_size, get_last_checkpoint, has_length, set_seed, speed_metrics
  from transformers.training_args import ParallelMode
  from transformers.utils import is_apex_available, is_sagemaker_dp_enabled, is_sagemaker_mp_enabled, is_torch_tpu_available
  from ..utils import logging
  from .training_args import ORTOptimizerNames, ORTTrainingArguments
  from .utils import is_onnxruntime_training_available
label_desc: Use image node_id 8 for calling the ORTTrainer obj's underlying member method code with example usage: obj.create_optimizer() and returns: self