docstr: Trains the Reformer Model
  :param epochs: The number of times you wish to loop through the dataset.
  :param train_dataloader: (torch.utils.data.DataLoader) The data to train on.
  :param eval_dataloader: (torch.utils.data.DataLoader) The data to evaluate on.
  :param log_steps: The number of steps to iterate before logging.
  :param ckpt_steps: The number of steps to iterate before checkpointing.
  :param ckpt_dir: The directory to save the checkpoints to.
  :param gradient_accumulation_steps: Optional gradient accumulation.
  :return: Total number of steps, total loss, model
docstr_tok: ["Trains","the","Reformer","Model",":","param","epochs",":","The","number","of","times","you","wish","to","loop","through","the","dataset",".",":","param","train_dataloader",":","(","torch.utils.data.DataLoader",")","The","data","to","train","on",".",":","param","eval_dataloader",":","(","torch.utils.data.DataLoader",")","The","data","to","evaluate","on",".",":","param","log_steps",":","The","number","of","steps","to","iterate","before","logging",".",":","param","ckpt_steps",":","The","number","of","steps","to","iterate","before","checkpointing",".",":","param","ckpt_dir",":","The","directory","to","save","the","checkpoints","to",".",":","param","gradient_accumulation_steps",":","Optional","gradient","accumulation",".",":","return",":","Total","number","of","steps",",","total","loss",",","model"]
code: def train(
      self,
      epochs,
      train_dataloader,
      eval_dataloader,
      log_steps,
      ckpt_steps,
      ckpt_dir=None,
      gradient_accumulation_steps=1,
  ):
      """
      Trains the Reformer Model
      :param epochs: The number of times you wish to loop through the dataset.
      :param train_dataloader: (torch.utils.data.DataLoader) The data to train on.
      :param eval_dataloader: (torch.utils.data.DataLoader) The data to evaluate on.
      :param log_steps: The number of steps to iterate before logging.
      :param ckpt_steps: The number of steps to iterate before checkpointing.
      :param ckpt_dir: The directory to save the checkpoints to.
      :param gradient_accumulation_steps: Optional gradient accumulation.
      :return: Total number of steps, total loss, model
      """
  
      optimizer = Adafactor(self.model.parameters())
      loss_fn = nn.CrossEntropyLoss()
      losses = {}
      global_steps = 0
      local_steps = 0
      step_loss = 0.0
  
      if ckpt_dir is not None:
          assert os.path.isdir(ckpt_dir)
          try:
              logging.info(
                  f"{datetime.now()} | Continuing from checkpoint..."
              )
              self.model.load_state_dict(
                  torch.load(
                      f"{ckpt_dir}/model_state_dict.pt",
                      map_location=self.device,
                  )
              )
              optimizer.load_state_dict(
                  torch.load(f"{ckpt_dir}/optimizer_state_dict.pt")
              )
  
          except Exception as e:
              logging.info(
                  f"{datetime.now()} | No checkpoint was found | {e}"
              )
  
      self.model.train()
  
      if self.n_gpu > 1:
          self.model = nn.DataParallel(self.model)
          logging.info(
              f"{datetime.now()} | Utilizing {self.n_gpu} GPUs"
          )
  
      self.model.to(self.device)
      logging.info(f"{datetime.now()} | Moved model to: {self.device}")
      logging.info(
          f"{datetime.now()} | train_batch_size: {self.train_batch_size} | eval_batch_size: {self.eval_batch_size}"
      )
      logging.info(
          f"{datetime.now()} | Epochs: {epochs} | log_steps: {log_steps} | ckpt_steps: {ckpt_steps}"
      )
      logging.info(
          f"{datetime.now()} | gradient_accumulation_steps: {gradient_accumulation_steps}"
      )
  
      for epoch in tqdm(range(epochs), desc="Epochs", position=0):
          logging.info(f"{datetime.now()} | Epoch: {epoch}")
          for step, batch in tqdm(
              enumerate(train_dataloader),
              desc="Epoch Iterator",
              position=1,
              leave=True,
              total=len(train_dataloader),
          ):
              for data in batch:
                  inputs = self._tokenize_input_ids(
                      data, pad_to_max_length=True
                  )
                  inputs, labels = self.mask_tokens(inputs)
                  inputs, labels = inputs.to(self.device), labels.to(
                      self.device
                  )
                  output = self.model(inputs)
  
                  # only calculating loss on masked tokens
                  loss_mx = labels != -100
                  output = output[loss_mx].view(
                      -1, self.tokenizer.vocab_size
                  )
                  labels = labels[loss_mx].view(-1)
  
                  loss = loss_fn(output, labels)
  
                  if gradient_accumulation_steps > 1:
                      loss /= gradient_accumulation_steps
  
                  loss.backward()
  
                  step_loss += loss.item()
                  losses[global_steps] = loss.item()
                  local_steps += 1
                  global_steps += 1
  
                  if global_steps % gradient_accumulation_steps == 0:
                      optimizer.step()
                      self.model.zero_grad()
  
                  if global_steps % log_steps == 0:
                      if self.tb_writer:
                          self.writer.add_scalar(
                              "Train/Loss",
                              step_loss / local_steps,
                              global_steps,
                          )
                          self.writer.close()
                      logging.info(
                          f"""{datetime.now()} | Train Loss: {step_loss / local_steps} | Steps: {global_steps}"""
                      )
  
                      with open(
                          f"{self.log_dir}/train_results.json", "w"
                      ) as results_file:
                          json.dump(losses, results_file)
                          results_file.close()
                      step_loss = 0.0
                      local_steps = 0
  
                  if global_steps % ckpt_steps == 0:
                      # evaluating before every checkpoint
                      self.evaluate(eval_dataloader)
                      model_to_save = (
                          self.model.module
                          if hasattr(self.model, "module")
                          else self.model
                      )
                      torch.save(
                          model_to_save.state_dict(),
                          f"{ckpt_dir}/model_state_dict.pt",
                      )
                      torch.save(
                          optimizer.state_dict(),
                          f"{ckpt_dir}/optimizer_state_dict.pt",
                      )
  
                      logging.info(
                          f"{datetime.now()} | Saved checkpoint to: {ckpt_dir}"
                      )
  
      model_to_save = (
          self.model.module
          if hasattr(self.model, "module")
          else self.model
      )
      torch.save(
          model_to_save.state_dict(), f"{ckpt_dir}/model_state_dict.pt"
      )
      torch.save(
          optimizer.state_dict(), f"{ckpt_dir}/optimizer_state_dict.pt"
      )
  
      return self.model
  
imports: import re
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.utils.data import Dataset, DataLoader, random_split
  from tqdm import tqdm
  from reformer_pytorch import Reformer, ReformerLM
  from transformers import BertTokenizer, PreTrainedTokenizer
  from fairseq.optim.adafactor import Adafactor
  import os
  import json
  import logging
  from datetime import datetime
label_desc: Use image node_id 5 for calling the ReformerTrainer obj's underlying member method code with example usage: obj.train(epochs, train_dataloader, eval_dataloader, log_steps, ckpt_steps, ckpt_dir, gradient_accumulation_steps) and returns: self