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
code: def __init__(self, config: ConfigActor) -> None:
      super().__init__()
  
      # save config
      self.config = config
  
      # initialize the self.model
      if config.model in llama_models:
          # llama module might not be present when HF models are used
          from chatllama.llama_model import (
              load_model,
              setup_model_parallel,
          )  # noqa
  
          local_rank, world_size = setup_model_parallel()
  
          # use load_model_test for testing
          self.model, self.tokenizer = load_model(
              ckpt_dir=config.model_folder,
              tokenizer_path=config.tokenizer_path,
              local_rank=local_rank,
              world_size=world_size,
              froze_embeddings=config.froze_embeddings,
              use_fairscale=config.use_fairscale,
              max_batch_size=config.batch_size,
          )
      elif config.model in hf_models_causal_lm:
          self.tokenizer = self.load_tokenizer(config)
          self.model = AutoModelForCausalLM.from_pretrained(
              config.model,
          )
  
          # Setup PEFT model
          if config.peft_enable:
              # check that the peft config exist
              if os.path.exists(config.peft_config_path):
                  # Read the peft config from yaml
                  with open(config.peft_config_path, "r") as c:
                      config_peft = yaml.safe_load(c)
              else:
                  raise ValueError(
                      f"PEFT config {config.peft_config_path} not found"
                  )
  
              print(config_peft)
              # define lora config for peft
              peft_config = LoraConfig(
                  task_type=TaskType.CAUSAL_LM, **config_peft
              )
  
              # create peft model
              self.model = get_peft_model(
                  model=self.model,
                  peft_config=peft_config,
              )
  
          self.model.to(config.device)
  
      else:
          raise ValueError(f"Model {config.model} not supported")
  
      # load the model from model_folder
      self.load()
  
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 1 to create a new ActorModel object from inherited base classes: torch.nn with example: obj = ActorModel(config)