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 actions and sequences=[states, actions] from state
      (i.e. input of the prompt generator model)
  
  Args:
      state (torch.Tensor): the input of the user
      state_mask (torch.Tensor): Mask for the state input (for padding)
  
  Returns:
      actions (torch.Tensor): Actions generated from the state
      sequences (torch.Tensor): Sequences generated from the
          state as [states, actions]
docstr_tok: ["Generate","actions","and","sequences=","[","states",",","actions","]","from","state","(","i.e",".","input","of","the","prompt","generator","model",")","Args",":","state","(","torch.Tensor",")",":","the","input","of","the","user","state_mask","(","torch.Tensor",")",":","Mask","for","the","state","input","(","for","padding",")","Returns",":","actions","(","torch.Tensor",")",":","Actions","generated","from","the","state","sequences","(","torch.Tensor",")",":","Sequences","generated","from","the","state","as","[","states",",","actions","]"]
code: def generate(
      self, states: torch.Tensor, state_mask: torch.Tensor
  ) -> Tuple:
      """Generate actions and sequences=[states, actions] from state
          (i.e. input of the prompt generator model)
  
      Args:
          state (torch.Tensor): the input of the user
          state_mask (torch.Tensor): Mask for the state input (for padding)
  
      Returns:
          actions (torch.Tensor): Actions generated from the state
          sequences (torch.Tensor): Sequences generated from the
              state as [states, actions]
      """
      # temperature for the actor
      temperature = self.config.temperature
  
      # max sequence length for the actor (i.e. prompt + completion)
      max_sequence_length = self.config.max_sequence_length
  
      # max and min number of tokens to generate
      max_tokens = self.config.max_tokens
      min_tokens = self.config.min_tokens
  
      # max generation possible given the state and the max sequence length
      max_generation_possible = max_sequence_length - states.shape[1]
      if max_generation_possible < min_tokens:
          raise ValueError(
              f"The prompt is too long w.r.t the "
              f"model sequence length \n"
              f"max_sequence_length={max_sequence_length}\n"
              f"state_length={states.shape[1]}\n"
              f"min_tokens={min_tokens}\n"
              f"max_tokens={max_tokens}\n"
              f"max_generation_possible={max_generation_possible}\n"
          )
  
      # take the minimum the max_tokens and the max_generation_possible
      max_completion = min(max_tokens, max_generation_possible)
  
      sequences = self.model.generate(
          input_ids=states,
          attention_mask=state_mask,
          temperature=temperature,
          max_new_tokens=max_completion,
          no_repeat_ngram_size=3,
      )
      actions = sequences[:, states.shape[1] :]  # noqa E203
      if self.config.debug:
          print(
              f"input length {states.shape[1]} \n"
              f"max sequence length {max_sequence_length} \n"
              f"max completion {max_completion} \n"
              f"generated sequence {sequences.shape[1]} \n"
          )
          print("ActorModel.generate")
          print("state", states)
          print("state shape", states.shape)
          print("sequence shape", sequences.shape)
          print("sequence", sequences)
          print("actions shape", actions.shape)
          print("actions", actions)
      return actions, sequences
  
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 7 for calling the ActorModel obj's underlying member method code with example usage: obj.generate(states, state_mask) and returns: actions,  sequences