code: def test_add_open_assistant(
      fixup_personality,
      only_personality,
      deberta_grading,
      prompt_type,
      save_json=True,
  ):
      """
      Flatten tree structure into one row per path from root to leaf
      Also turn into human_bot prompting format:
          <human>: question\n<bot>: answer <human>: question2\n<bot>: answer2 Etc.
      Also saves a .json locally as side-effect
      returns list of dicts, containing intput, prompt_type and source
      """
      from datasets import load_dataset
  
      data_file = "OpenAssistant/oasst1"
      ds = load_dataset(data_file)
      df = pd.concat(
          [ds["train"].to_pandas(), ds["validation"].to_pandas()],
          axis=0,
      )
      rows = {}
      message_ids = df["message_id"].values.tolist()
      message_tree_ids = df["message_tree_id"].values.tolist()
      parent_ids = df["parent_id"].values.tolist()
      texts = df["text"].values.tolist()
      roles = df["role"].values.tolist()
      deleteds = df["deleted"].values.tolist()
      for i in range(df.shape[0]):
          # collect all trees
          message_id = message_ids[i]
          message_tree_id = message_tree_ids[i]
          parent_id = parent_ids[i]
          text = texts[i]
          deleted = deleteds[i]
          if deleted:
              continue
          if fixup_personality:
              text = text.replace("Open Assistant", "h2oGPT")
              text = text.replace("Open-Assistant", "h2oGPT")
              text = text.replace("open-assistant", "h2oGPT")
              text = text.replace("OpenAssistant", "h2oGPT")
              text = text.replace("open assistant", "h2oGPT")
              text = text.replace("Open Assistand", "h2oGPT")
              text = text.replace("Open Assitant", "h2oGPT")
              text = text.replace("Open Assistent", "h2oGPT")
              text = text.replace("Open Assisstant", "h2oGPT")
              text = text.replace("Open Assitent", "h2oGPT")
              text = text.replace("Open Assitiant", "h2oGPT")
              text = text.replace("Open Assistiant", "h2oGPT")
              text = text.replace("Open Assitan ", "h2oGPT ")
              text = text.replace("Open Assistan ", "h2oGPT ")
              text = text.replace("Open Asistant", "h2oGPT")
              text = text.replace("Open Assiant", "h2oGPT")
              text = text.replace("Assistant", "h2oGPT")
              text = text.replace("LAION AI", "H2O.ai")
              text = text.replace("LAION-AI", "H2O.ai")
              text = text.replace("LAION,", "H2O.ai,")
              text = text.replace("LAION.ai", "H2O.ai")
              text = text.replace("LAION.", "H2O.ai.")
              text = text.replace("LAION", "H2O.ai")
  
          role = roles[i]
          if prompt_type == "llama2":
              new_data = (
                  "[INST] " if role == "prompter" else " [/INST] "
              ) + text
              if parent_id and role == "prompter":
                  new_data = " " + new_data
          elif prompt_type == "human_bot":
              new_data = (
                  "<human>: " if role == "prompter" else "<bot>: "
              ) + text
          else:
              raise NotImplementedError("prompt_type not supported")
          entry = dict(
              message_id=message_id, parent_id=parent_id, text=new_data
          )
          if message_tree_id not in rows:
              rows[message_tree_id] = [entry]
          else:
              rows[message_tree_id].append(entry)
  
      all_rows = []
  
      for node_id in rows:
          # order responses in tree, based on message/parent relationship
          conversations = []
  
          list_msgs = rows[node_id]
          # find start
          while len(list_msgs):
              for i, leaf in enumerate(list_msgs):
                  found = False
                  parent_id = leaf["parent_id"]
                  if parent_id is None:
                      # conversation starter
                      conversations.append(leaf)
                      found = True
                  else:
                      for conv in conversations:
                          # find all conversations to add my message to
                          if (
                              parent_id in conv["message_id"]
                              and parent_id
                              != conv["message_id"][-len(parent_id) :]
                          ):
                              # my message doesn't follow conversation
                              continue
                          if (
                              parent_id
                              == conv["message_id"][-len(parent_id) :]
                          ):
                              # my message follows conversation, but fork first, so another follow-on message can do same
                              conversations.append(conv.copy())
                              if prompt_type == "llama2":
                                  conv["text"] += f"""{leaf['text']}"""
                              elif prompt_type == "human_bot":
                                  conv[
                                      "text"
                                  ] += f"""
  {leaf['text']}
  """
                              else:
                                  raise NotImplementedError
                              conv["message_id"] += leaf["message_id"]
                              found = True
                              break
                  if found:
                      # my content was used, so nuke from list
                      del list_msgs[i]
                      break
  
          # now reduce down to final conversations, find the longest chains of message ids
          for i, conv in enumerate(conversations):
              for j, conv2 in enumerate(conversations):
                  if i == j:
                      continue
                  if conv["message_id"] and conv2["message_id"]:
                      assert conv["message_id"] != conv2["message_id"]
                      # delete the shorter conversation, if one contains the other
                      if conv["message_id"] in conv2["message_id"]:
                          conv["message_id"] = None
                      if conv2["message_id"] in conv["message_id"]:
                          conv2["message_id"] = None
          conversations = [c for c in conversations if c["message_id"]]
          if only_personality:
              if prompt_type == "human_bot":
                  all_rows.extend(
                      [
                          dict(
                              input=c["text"] + "\n<human>:",
                              output="",
                              prompt_type="plain",
                              source=data_file,
                          )
                          for c in conversations
                          if "h2oGPT" in c["text"]
                      ]
                  )
              elif prompt_type == "llama2":
                  all_rows.extend(
                      [
                          dict(
                              input=c["text"]
                              + (
                                  ""
                                  if c["text"].rfind("[/INST]")
                                  > c["text"].rfind("[INST]")
                                  else " [/INST]"
                              ),
                              output="",
                              prompt_type="plain",
                              source=data_file,
                          )
                          for c in conversations
                          if "h2oGPT" in c["text"]
                      ]
                  )
              else:
                  raise NotImplementedError
          else:
              if prompt_type == "human_bot":
                  all_rows.extend(
                      [
                          dict(
                              input=c["text"] + "\n<human>:",
                              output="",
                              prompt_type="plain",
                              source=data_file,
                          )
                          for c in conversations
                          if "What is H2O.ai" not in c["text"]
                      ]
                  )
              elif prompt_type == "llama2":
                  all_rows.extend(
                      [
                          dict(
                              input=c["text"]
                              + (
                                  " "
                                  if c["text"].rfind("[/INST]")
                                  > c["text"].rfind("[INST]")
                                  else " [/INST]"
                              ),
                              output="",
                              prompt_type="plain",
                              source=data_file,
                          )
                          for c in conversations
                          if "What is H2O.ai" not in c["text"]
                      ]
                  )
              else:
                  raise NotImplementedError
  
      unhelpful = get_unhelpful_list()
      all_rows = [
          x
          for x in all_rows
          if not any(u in x["input"] for u in unhelpful)
      ]
      personality = create_personality_data(prompt_type=prompt_type)
      all_rows.extend(personality * 10)
      np.random.seed(123)
      np.random.shuffle(all_rows)
      print(len(all_rows))
      if deberta_grading:
          df = pd.DataFrame(all_rows)
          df = df.rename(columns={"input": "text"})
          df = add_deberta_grade(df)
          df = df.rename(columns={"text": "input"})
          drop = True
          if drop:
              min_grade = 0.3
              max_grade = np.inf
              before_rows = df.shape[0]
              df = df[df["grade_deberta"] >= min_grade]
              df = df[df["grade_deberta"] <= max_grade]
              after_rows = df.shape[0]
              print(
                  "Dropped %d rows out of %d due to deberta grade"
                  % (before_rows - after_rows, before_rows)
              )
              print("After DeBERTa grade")
          print(df.describe())
          all_rows = []
          for i in range(df.shape[0]):
              all_rows.append(
                  dict(
                      input=df["input"].iloc[i],
                      output=df["output"].iloc[i],
                      source=df["source"].iloc[i],
                      prompt_type=df["prompt_type"].iloc[i],
                      grade_deberta=df["grade_deberta"].iloc[i],
                  )
              )
      if save_json:
          data_file = (
              data_file
              + ("_h2ogpt" if fixup_personality else "")
              + ("_only" if only_personality else "")
              + ("_graded" if deberta_grading else "")
              + ("_llama2_chat" if prompt_type == "llama2" else "")
          )
          for i in range(len(all_rows)):
              all_rows[i]["id"] = i
          with open(
              data_file.lower().replace("/", "_") + ".json", "w"
          ) as f:
              f.write(json.dumps(all_rows, indent=2))
      return all_rows
  
imports: import ast
  import concurrent.futures
  import contextlib
  import hashlib
  import json
  import os
  import shutil
  import signal
  import sys
  import traceback
  from concurrent.futures import ProcessPoolExecutor
  import psutil
  import pytest
  import pandas
  import numpy
  from tqdm import tqdm
  from utils import flatten_list, remove
  from joblib import Parallel, delayed, effective_n_jobs
  from sklearn.utils import gen_even_slices
  from sklearn.utils.validation import _num_samples
label_desc: Use image node_id 32 for calling a global function with example usage: test_add_open_assistant(fixup_personality, only_personality, deberta_grading, prompt_type, save_json) and returns: all_rows