class_docstr: Actor model that generates the augmented prompt from the initial
  user_input. The aim is to train this model to generate better prompts.
  
  Attributes:
      model: The model from LLaMA to be used
      tokenizer: The LLaMA tokenizer
      config (ConfigActor): Configuration for the actor model
  
  Methods:
      load: Load the model from a path
      save: Save the model to a path
      forward: Compute the action logits for a given sequence.
      generate: Generate a sequence from a given prompt
docstr: Generate logits to have probability distribution over the vocabulary
      of the actions
  
  Args:
      sequences (torch.Tensor): Sequences of states and actions used to
              compute token logits for the whole list of sequences
      attention_mask (torch.Tensor): Mask for the sequences attention
  
  Returns:
      logits (torch.Tensor): Logits for the actions taken
docstr_tok: ["Generate","logits","to","have","probability","distribution","over","the","vocabulary","of","the","actions","Args",":","sequences","(","torch.Tensor",")",":","Sequences","of","states","and","actions","used","to","compute","token","logits","for","the","whole","list","of","sequences","attention_mask","(","torch.Tensor",")",":","Mask","for","the","sequences","attention","Returns",":","logits","(","torch.Tensor",")",":","Logits","for","the","actions","taken"]
code: def forward(
      self, sequences: torch.Tensor, sequences_mask: torch.Tensor
  ) -> torch.Tensor:
      """Generate logits to have probability distribution over the vocabulary
          of the actions
  
      Args:
          sequences (torch.Tensor): Sequences of states and actions used to
                  compute token logits for the whole list of sequences
          attention_mask (torch.Tensor): Mask for the sequences attention
  
      Returns:
          logits (torch.Tensor): Logits for the actions taken
      """
      model_output = self.model.forward(
          sequences, attention_mask=sequences_mask
      )
      # need to return logits for the actions
      if self.config.model in hf_models_causal_lm:
          model_output = model_output.logits
      if self.config.debug:
          print("ActorModel.forward")
          print("model_output_logits shape", model_output.shape)
          print("model_output logits", model_output)
      return model_output
  
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 6 for calling the ActorModel obj's underlying member method code with example usage: obj.forward(sequences, sequences_mask) and returns: model_output