class_docstr: Used to pre-train the actor model to generate better prompts.
  
  Args:
      config (ConfigActor): Configuration for the actor model
  
  Attributes:
      config (ConfigActor): Configuration for the actor model
      model (ActorModel): Actor model
      loss_function (torch.nn.CrossEntropyLoss): Loss function
      optimizer (torch.optim.Adam): Optimizer
      validation_flag (bool): Flag to indicate if the validation dataset
          is provided
      train_dataset (ActorDataset): Training dataset
      train_dataloader (DataLoader): Training dataloader
      validation_dataset (ActorDataset): Validation dataset
      validation_dataloader (DataLoader): Validation dataloader
      scheduler (torch.optim.lr_scheduler): Learning rate scheduler
      training_stats (TrainingStats): Training statistics
      model_engine (ModelEngine): Model engine for deepspeed training
      accelerator (Accelerator): Accelerator for accelerate training
  
  Methods:
      train: Train the actor model
      load_checkpoint: Load a checkpoint
      save_checkpoint: Save a checkpoint
docstr: Train the model
docstr_tok: ["Train","the","model"]
code: def train(
      self,
  ) -> None:
      """Train the model"""
      print("Start Actor Model Pretraining")
  
      # get config parameters
      if self.config.deepspeed_enable:
          batch_size = self.train_dataloader.batch_size
      else:
          batch_size = self.config.batch_size
      epochs = self.config.epochs
      device = self.config.device
      checkpoint_steps = self.config.checkpoint_steps
  
      # compute the number of iterations
      n_iter = int(len(self.train_dataset) / batch_size)
  
      # load model_checkpoint
      start_epoch, start_step = self.load_checkpoint()
  
      if start_epoch == 0 and start_step == 0:
          self.training_stats.clear()
  
      # counter for the checkpoint
      cnt_checkpoint = 1
  
      # traing loop
      for epoch in range(start_epoch, epochs):
          self.actor.train()
          for i, input_text in enumerate(self.train_dataloader):
              # skip the first steps if we are resuming training
              if i < start_step:
                  continue
  
              # tokenize input
              with torch.no_grad():
                  input_tokenized = self.actor.tokenizer(
                      input_text,
                      return_tensors="pt",
                      truncation=True,
                      padding=True,
                  )
  
                  # split tokens and mask
                  input_tokenized_id = input_tokenized["input_ids"]
                  input_tokenized_mask = input_tokenized[
                      "attention_mask"
                  ]
  
                  # add eos token
                  (
                      input_tokenized_id,
                      input_tokenized_mask,
                  ) = self.add_eos_token(
                      input_tokenized_id,
                      input_tokenized_mask,
                  )
  
                  # split into input and output
                  training_output = input_tokenized_id[:, 1:]
                  training_input = input_tokenized_id[:, :-1]
                  attention_mask = input_tokenized_mask[:, :-1]
  
                  # move to device
                  training_output = training_output.to(device)
                  training_input = training_input.to(device)
                  attention_mask = attention_mask.to(device)
  
              # forward pass
              if self.config.deepspeed_enable:
                  est_output = self.model_engine(
                      training_input, attention_mask
                  )
              else:
                  est_output = self.actor(
                      training_input, attention_mask
                  )
  
              # compute loss
              est_output = rearrange(est_output, "b s v -> (b s) v")
              training_output = rearrange(
                  training_output, "b s -> (b s)"
              )
              loss = self.loss_function(est_output, training_output)
              self.training_stats.training_loss.append(loss.item())
  
              # backward pass
              if self.config.deepspeed_enable:
                  self.model_engine.backward(loss)
                  self.model_engine.step()
              elif self.config.accelerate_enable:
                  self.optimizer.zero_grad()
                  self.accelerator.backward(loss)
                  self.optimizer.step()
                  self.scheduler.step()
              else:
                  self.optimizer.zero_grad()
                  loss.backward()
                  self.optimizer.step()
                  self.scheduler.step()
  
              # print progress
              if i % self.config.iteration_per_print == 0:
                  print(
                      f"Epoch: {epoch+1}/{epochs}, "
                      f"Iteration: {i+1}/{n_iter}, "
                      f"Training Loss: {loss}"
                  )
  
              # save checkpoint periodically
              if cnt_checkpoint % checkpoint_steps == 0:
                  self.save_checkpoint(epoch, i, epochs, n_iter)
                  self.training_stats.save()
                  cnt_checkpoint = 1
              else:
                  cnt_checkpoint += 1
  
          # Validation
          if self.validation_flag:
              self.actor.eval()
              with torch.no_grad():
                  for i, input_text in enumerate(
                      self.validation_dataloader
                  ):
                      # tokenize input
                      input_tokenized = self.actor.tokenizer(
                          input_text, return_tensors="pt", padding=True
                      )
                      validation_output = input_tokenized["input_ids"][
                          :, 1:
                      ]
                      validation_input = input_tokenized["input_ids"][
                          :, :-1
                      ]
                      attention_mask = input_tokenized[
                          "attention_mask"
                      ][:, :-1]
  
                      # forward pass
                      est_output = self.actor.forward(
                          validation_input, attention_mask
                      )
                      validation_output = rearrange(
                          validation_output, "b s -> (b s)"
                      )
  
                      # compute loss
                      est_output = rearrange(
                          est_output, "b s v -> (b s) v"
                      )
                      loss = self.loss_function(
                          est_output, validation_output
                      )
                      self.training_stats.validation_loss.append(
                          loss.item()
                      )
  
                      # print progress
                      if i % self.config.iteration_per_print == 0:
                          print(
                              f"Epoch: {epoch+1}/{epochs}, "
                              f"Iteration: {i+1}/{n_iter}, "
                              f"Validation Loss: {loss}"
                          )
          # reset start_step after training is resumed
          start_step = 0
  
      # save the model
      self.actor.save()
      print("Training Finished ")
  
imports: import json
  import yaml
  import os
  import shutil
  import deepspeed
  import torch
  from accelerate import Accelerator
  from beartype import beartype
  from beartype.typing import Tuple
  from einops import rearrange
  from peft import get_peft_model, LoraConfig, TaskType
  from torch.utils.data import DataLoader, Dataset
  from transformers import AutoModelForCausalLM, AutoTokenizer
  from chatllama.rlhf.config import ConfigActor
  from chatllama.rlhf.model_list import hf_models_causal_lm, llama_models, hf_models
  from chatllama.rlhf.model_loader import ModelLoader
  from chatllama.rlhf.utils import TrainingStats
label_desc: Use image node_id 5 for calling the ActorTrainer obj's underlying member method code with example usage: obj.train() without return types