code: def main():
      import json
      from argparse import ArgumentParser
  
      parser = ArgumentParser()
      parser.add_argument(
          "--model",
          type=str,
          help="Model name.",
          default="google/flan-t5-xl",
      )
      parser.add_argument(
          "--templates", type=str, help="Path to templates."
      )
      parser.add_argument("--num_prompts", type=int, default=1000)
      parser.add_argument(
          "--data_dir", type=str, help="Path where data are stored"
      )
  
      args = parser.parse_args()
      model_name = args.model
      templates_path = args.templates
      data_dir = args.data_dir
  
      with open(
          os.path.join(data_dir, "rlhf_training_data.json"), "r"
      ) as f:
          examples = json.load(f)
  
      with open(templates_path, "r") as f:
          templates = json.load(f)
      user_prompt = templates.get("rlhf")
      if user_prompt is None:
          raise ValueError("No rlhs template found.")
  
      if "davinci" in model_name:
          chain = use_langchain_model(
              user_prompt, model_name, with_examples=True
          )
      else:
          if "t5" not in model_name:
              raise ValueError(
                  "Only Flan-t5 models are supported for HF."
              )
          chain = use_huggingface_model(
              user_prompt, model_name, with_examples=True
          )
  
      for i in range(args.num_prompts):
          example = np.random.choice(examples)
          new_example = chain.run(example=example["user_input"])
          example_dict = {"user_input": new_example}
          examples.append(example_dict)
  
      with open(
          os.path.join(data_dir, "rlhf_training_data.json"), "w"
      ) as f:
          json.dump(examples, f)
  
file_path: nebuly/optimization/chatllama/artifacts/extend_rlhf_dataset.py
imports: import os.path
  import numpy
  from langchain import OpenAI, LLMChain, PromptTemplate
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
label_desc: Use image node_id 4 for calling a global function with example usage: main() without return types