class_docstr: Trainer is training and eval loop for adversarial training.
  
  It is designed to work with PyTorch and Transformers models.
  
  Args:
      model_wrapper (:class:`~textattack.models.wrappers.ModelWrapper`):
          Model wrapper containing both the model and the tokenizer.
      task_type (:obj:`str`, `optional`, defaults to :obj:`"classification"`):
          The task that the model is trained to perform.
          Currently, :class:`~textattack.Trainer` supports two tasks: (1) :obj:`"classification"`, (2) :obj:`"regression"`.
      attack (:class:`~textattack.Attack`):
          :class:`~textattack.Attack` used to generate adversarial examples for training.
      train_dataset (:class:`~textattack.datasets.Dataset`):
          Dataset for training.
      eval_dataset (:class:`~textattack.datasets.Dataset`):
          Dataset for evaluation
      training_args (:class:`~textattack.TrainingArgs`):
          Arguments for training.
  
  Example::
  
      >>> import textattack
      >>> import transformers
  
      >>> model = transformers.AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
      >>> tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased")
      >>> model_wrapper = textattack.models.wrappers.HuggingFaceModelWrapper(model, tokenizer)
  
      >>> # We only use DeepWordBugGao2018 to demonstration purposes.
      >>> attack = textattack.attack_recipes.DeepWordBugGao2018.build(model_wrapper)
      >>> train_dataset = textattack.datasets.HuggingFaceDataset("imdb", split="train")
      >>> eval_dataset = textattack.datasets.HuggingFaceDataset("imdb", split="test")
  
      >>> # Train for 3 epochs with 1 initial clean epochs, 1000 adversarial examples per epoch, learning rate of 5e-5, and effective batch size of 32 (8x4).
      >>> training_args = textattack.TrainingArgs(
      ...     num_epochs=3,
      ...     num_clean_epochs=1,
      ...     num_train_adv_examples=1000,
      ...     learning_rate=5e-5,
      ...     per_device_train_batch_size=8,
      ...     gradient_accumulation_steps=4,
      ...     log_to_tb=True,
      ... )
  
      >>> trainer = textattack.Trainer(
      ...     model_wrapper,
      ...     "classification",
      ...     attack,
      ...     train_dataset,
      ...     eval_dataset,
      ...     training_args
      ... )
      >>> trainer.train()
  
  .. note::
      When using :class:`~textattack.Trainer` with `parallel=True` in :class:`~textattack.TrainingArgs`,
      make sure to protect the “entry point” of the program by using :obj:`if __name__ == '__main__':`.
      If not, each worker process used for generating adversarial examples will execute the training code again.
class_docstr_tok: ["Trainer","is","training","and","eval","loop","for","adversarial","training",".","It","is","designed","to","work","with","PyTorch","and","Transformers","models",".","Args",":","model_wrapper","(",":","class",":","`","~textattack.models.wrappers.ModelWrapper","`",")",":","Model","wrapper","containing","both","the","model","and","the","tokenizer",".","task_type","(",":","obj",":","`","str","`",",","`","optional","`",",","defaults","to",":","obj",":","`","``","classification","''","`",")",":","The","task","that","the","model","is","trained","to","perform",".","Currently",",",":","class",":","`","~textattack.Trainer","`","supports","two","tasks",":","(","1",")",":","obj",":","`","``","classification","''","`",",","(","2",")",":","obj",":","`","``","regression","''","`",".","attack","(",":","class",":","`","~textattack.Attack","`",")",":",":","class",":","`","~textattack.Attack","`","used","to","generate","adversarial","examples","for","training",".","train_dataset","(",":","class",":","`","~textattack.datasets.Dataset","`",")",":","Dataset","for","training",".","eval_dataset","(",":","class",":","`","~textattack.datasets.Dataset","`",")",":","Dataset","for","evaluation","training_args","(",":","class",":","`","~textattack.TrainingArgs","`",")",":","Arguments","for","training",".","Example",":",":",">",">",">","import","textattack",">",">",">","import","transformers",">",">",">","model","=","transformers.AutoModelForSequenceClassification.from_pretrained","(","``","bert-base-uncased","''",")",">",">",">","tokenizer","=","transformers.AutoTokenizer.from_pretrained","(","``","bert-base-uncased","''",")",">",">",">","model_wrapper","=","textattack.models.wrappers.HuggingFaceModelWrapper","(","model",",","tokenizer",")",">",">",">","#","We","only","use","DeepWordBugGao2018","to","demonstration","purposes",".",">",">",">","attack","=","textattack.attack_recipes.DeepWordBugGao2018.build","(","model_wrapper",")",">",">",">","train_dataset","=","textattack.datasets.HuggingFaceDataset","(","``","imdb","''",",","split=","''","train","''",")",">",">",">","eval_dataset","=","textattack.datasets.HuggingFaceDataset","(","``","imdb","''",",","split=","''","test","''",")",">",">",">","#","Train","for","3","epochs","with","1","initial","clean","epochs",",","1000","adversarial","examples","per","epoch",",","learning","rate","of","5e-5",",","and","effective","batch","size","of","32","(","8x4",")",".",">",">",">","training_args","=","textattack.TrainingArgs","(","...","num_epochs=3",",","...","num_clean_epochs=1",",","...","num_train_adv_examples=1000",",","...","learning_rate=5e-5",",","...","per_device_train_batch_size=8",",","...","gradient_accumulation_steps=4",",","...","log_to_tb=True",",","...",")",">",">",">","trainer","=","textattack.Trainer","(","...","model_wrapper",",","...","``","classification","''",",","...","attack",",","...","train_dataset",",","...","eval_dataset",",","...","training_args","...",")",">",">",">","trainer.train","(",")","..","note",":",":","When","using",":","class",":","`","~textattack.Trainer","`","with","`","parallel=True","`","in",":","class",":","`","~textattack.TrainingArgs","`",",","make","sure","to","protect","the","\u201c","entry","point","\u201d","of","the","program","by","using",":","obj",":","`","if","__name__","==","'__main__","'",":","`",".","If","not",",","each","worker","process","used","for","generating","adversarial","examples","will","execute","the","training","code","again","."]
docstr: Train the model on given training dataset.
docstr_tok: ["Train","the","model","on","given","training","dataset","."]
code: def train(self):
      """Train the model on given training dataset."""
      if not self.train_dataset:
          raise ValueError("No `train_dataset` available for training.")
  
      textattack.shared.utils.set_seed(self.training_args.random_seed)
      if not os.path.exists(self.training_args.output_dir):
          os.makedirs(self.training_args.output_dir)
  
      # Save logger writes to file
      log_txt_path = os.path.join(
          self.training_args.output_dir, "train_log.txt"
      )
      fh = logging.FileHandler(log_txt_path)
      fh.setLevel(logging.DEBUG)
      logger.addHandler(fh)
      logger.info(f"Writing logs to {log_txt_path}.")
  
      # Save original self.training_args to file
      args_save_path = os.path.join(
          self.training_args.output_dir, "training_args.json"
      )
      with open(args_save_path, "w", encoding="utf-8") as f:
          json.dump(self.training_args.__dict__, f)
      logger.info(f"Wrote original training args to {args_save_path}.")
  
      num_gpus = torch.cuda.device_count()
      tokenizer = self.model_wrapper.tokenizer
      model = self.model_wrapper.model
  
      if self.training_args.parallel and num_gpus > 1:
          # TODO: torch.nn.parallel.DistributedDataParallel
          # Supposedly faster than DataParallel, but requires more work to setup properly.
          model = torch.nn.DataParallel(model)
          logger.info(
              f"Training on {num_gpus} GPUs via `torch.nn.DataParallel`."
          )
          train_batch_size = (
              self.training_args.per_device_train_batch_size * num_gpus
          )
      else:
          train_batch_size = (
              self.training_args.per_device_train_batch_size
          )
  
      if self.attack is None:
          num_clean_epochs = self.training_args.num_epochs
      else:
          num_clean_epochs = self.training_args.num_clean_epochs
  
      total_clean_training_steps = (
          math.ceil(
              len(self.train_dataset)
              / (
                  train_batch_size
                  * self.training_args.gradient_accumulation_steps
              )
          )
          * num_clean_epochs
      )
  
      # calculate total_adv_training_data_length based on type of
      # num_train_adv_examples.
      # if num_train_adv_examples is float , num_train_adv_examples is a portion of train_dataset.
      if isinstance(self.training_args.num_train_adv_examples, float):
          total_adv_training_data_length = (
              len(self.train_dataset)
              * self.training_args.num_train_adv_examples
          )
  
      # if num_train_adv_examples is int and >=0 then it is taken as value.
      elif (
          isinstance(self.training_args.num_train_adv_examples, int)
          and self.training_args.num_train_adv_examples >= 0
      ):
          total_adv_training_data_length = (
              self.training_args.num_train_adv_examples
          )
  
      # if num_train_adv_examples is = -1 , we generate all possible adv examples.
      # Max number of all possible adv examples would be equal to train_dataset.
      else:
          total_adv_training_data_length = len(self.train_dataset)
  
      # Based on total_adv_training_data_length calculation , find total total_adv_training_steps
      total_adv_training_steps = math.ceil(
          (len(self.train_dataset) + total_adv_training_data_length)
          / (
              train_batch_size
              * self.training_args.gradient_accumulation_steps
          )
      ) * (self.training_args.num_epochs - num_clean_epochs)
  
      total_training_steps = (
          total_clean_training_steps + total_adv_training_steps
      )
  
      optimizer, scheduler = self.get_optimizer_and_scheduler(
          model, total_training_steps
      )
  
      self._print_training_args(
          total_training_steps, train_batch_size, num_clean_epochs
      )
  
      model.to(textattack.shared.utils.device)
  
      # Variables across epochs
      self._total_loss = 0.0
      self._current_loss = 0.0
      self._last_log_step = 0
  
      # `best_score` is used to keep track of the best model across training.
      # Could be loss, accuracy, or other metrics.
      best_eval_score = 0.0
      best_eval_score_epoch = 0
      best_model_path = None
      epochs_since_best_eval_score = 0
  
      for epoch in range(1, self.training_args.num_epochs + 1):
          logger.info(
              "=========================================================="
          )
          logger.info(f"Epoch {epoch}")
  
          if self.attack and epoch > num_clean_epochs:
              if (
                  epoch - num_clean_epochs - 1
              ) % self.training_args.attack_epoch_interval == 0:
                  # only generate a new adversarial training set every self.training_args.attack_period epochs after the clean epochs
                  # adv_dataset is instance of `textattack.datasets.Dataset`
                  model.eval()
                  adv_dataset = self._generate_adversarial_examples(
                      epoch
                  )
                  model.train()
                  model.to(textattack.shared.utils.device)
              else:
                  adv_dataset = None
          else:
              logger.info(
                  f"Running clean epoch {epoch}/{num_clean_epochs}"
              )
              adv_dataset = None
  
          train_dataloader = self.get_train_dataloader(
              self.train_dataset, adv_dataset, train_batch_size
          )
          model.train()
          # Epoch variables
          all_preds = []
          all_targets = []
          prog_bar = tqdm.tqdm(
              train_dataloader,
              desc="Iteration",
              position=0,
              leave=True,
              dynamic_ncols=True,
          )
          for step, batch in enumerate(prog_bar):
              loss, preds, targets = self.training_step(
                  model, tokenizer, batch
              )
  
              if isinstance(model, torch.nn.DataParallel):
                  loss = loss.mean()
  
              loss = (
                  loss / self.training_args.gradient_accumulation_steps
              )
              loss.backward()
              loss = loss.item()
              self._total_loss += loss
              self._current_loss += loss
  
              all_preds.append(preds)
              all_targets.append(targets)
  
              if (
                  (step + 1)
                  % self.training_args.gradient_accumulation_steps
                  == 0
              ):
                  optimizer.step()
                  if scheduler:
                      scheduler.step()
                  optimizer.zero_grad()
                  self._global_step += 1
  
              if self._global_step > 0:
                  prog_bar.set_description(
                      f"Loss {self._total_loss/self._global_step:.5f}"
                  )
  
              # TODO: Better way to handle TB and Wandb logging
              if (self._global_step > 0) and (
                  self._global_step
                  % self.training_args.logging_interval_step
                  == 0
              ):
                  lr_to_log = (
                      scheduler.get_last_lr()[0]
                      if scheduler
                      else self.training_args.learning_rate
                  )
                  if self._global_step - self._last_log_step >= 1:
                      loss_to_log = round(
                          self._current_loss
                          / (self._global_step - self._last_log_step),
                          4,
                      )
                  else:
                      loss_to_log = round(self._current_loss, 4)
  
                  log = {
                      "train/loss": loss_to_log,
                      "train/learning_rate": lr_to_log,
                  }
                  if self.training_args.log_to_tb:
                      self._tb_log(log, self._global_step)
  
                  if self.training_args.log_to_wandb:
                      self._wandb_log(log, self._global_step)
  
                  self._current_loss = 0.0
                  self._last_log_step = self._global_step
  
              # Save model checkpoint to file.
              if self.training_args.checkpoint_interval_steps:
                  if (
                      self._global_step > 0
                      and (
                          self._global_step
                          % self.training_args.checkpoint_interval_steps
                      )
                      == 0
                  ):
                      self._save_model_checkpoint(
                          model, tokenizer, step=self._global_step
                      )
  
          preds = torch.cat(all_preds)
          targets = torch.cat(all_targets)
          if self._metric_name == "accuracy":
              correct_predictions = (preds == targets).sum().item()
              accuracy = correct_predictions / len(targets)
              metric_log = {"train/train_accuracy": accuracy}
              logger.info(f"Train accuracy: {accuracy*100:.2f}%")
          else:
              (
                  pearson_correlation,
                  pearson_pvalue,
              ) = scipy.stats.pearsonr(preds, targets)
              metric_log = {
                  "train/pearson_correlation": pearson_correlation,
                  "train/pearson_pvalue": pearson_pvalue,
              }
              logger.info(
                  f"Train Pearson correlation: {pearson_correlation:.4f}%"
              )
  
          if len(targets) > 0:
              if self.training_args.log_to_tb:
                  self._tb_log(metric_log, epoch)
              if self.training_args.log_to_wandb:
                  metric_log["epoch"] = epoch
                  self._wandb_log(metric_log, self._global_step)
  
          # Evaluate after each epoch.
          eval_score = self.evaluate()
  
          if self.training_args.log_to_tb:
              self._tb_log(
                  {f"eval/{self._metric_name}": eval_score}, epoch
              )
          if self.training_args.log_to_wandb:
              self._wandb_log(
                  {
                      f"eval/{self._metric_name}": eval_score,
                      "epoch": epoch,
                  },
                  self._global_step,
              )
  
          if (
              self.training_args.checkpoint_interval_epochs
              and (
                  epoch % self.training_args.checkpoint_interval_epochs
              )
              == 0
          ):
              self._save_model_checkpoint(model, tokenizer, epoch=epoch)
  
          if eval_score > best_eval_score:
              best_eval_score = eval_score
              best_eval_score_epoch = epoch
              epochs_since_best_eval_score = 0
              self._save_model_checkpoint(model, tokenizer, best=True)
              logger.info(
                  f"Best score found. Saved model to {self.training_args.output_dir}/best_model/"
              )
          else:
              epochs_since_best_eval_score += 1
              if self.training_args.early_stopping_epochs and (
                  epochs_since_best_eval_score
                  > self.training_args.early_stopping_epochs
              ):
                  logger.info(
                      f"Stopping early since it's been {self.training_args.early_stopping_epochs} steps since validation score increased."
                  )
                  break
  
      if self.training_args.log_to_tb:
          self._tb_writer.flush()
  
      # Finish training
      if isinstance(model, torch.nn.DataParallel):
          model = model.module
  
      if self.training_args.load_best_model_at_end:
          best_model_path = os.path.join(
              self.training_args.output_dir, "best_model"
          )
          if hasattr(model, "from_pretrained"):
              model = model.__class__.from_pretrained(best_model_path)
          else:
              model = model.load_state_dict(
                  torch.load(
                      os.path.join(best_model_path, "pytorch_model.bin")
                  )
              )
  
      if self.training_args.save_last:
          self._save_model_checkpoint(model, tokenizer, last=True)
  
      self.model_wrapper.model = model
      self._write_readme(
          best_eval_score, best_eval_score_epoch, train_batch_size
      )
  
imports: import collections
  import json
  import logging
  import math
  import os
  import scipy
  import torch
  import tqdm
  import transformers
  import textattack
  from textattack.shared.utils import logger
  from .attack import Attack
  from .attack_args import AttackArgs
  from .attack_results import MaximizedAttackResult, SuccessfulAttackResult
  from .attacker import Attacker
  from .model_args import HUGGINGFACE_MODELS
  from .models.helpers import LSTMForClassification, WordCNNForClassification
  from .models.wrappers import ModelWrapper
  from .training_args import CommandLineTrainingArgs, TrainingArgs
label_desc: Use image node_id 12 for calling the Trainer obj's underlying member method code with example usage: obj.train() without return types