processing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Welcome to Augmentoolkit\n",
    "\n",
    "This notebook is where you generate all your data.\n",
    "\n",
    "Augmentoolkit is meant to allow instruct-tuned models to learn from books, even using themselves to generate new data through a sort-of bootstrapping method. It is meant to stop model creators from having to work as data annotators, and not actual model trainers. It is meant to allow anyone to make their own high-quality dataset with thousands of entries.\n",
    "\n",
    "## Quickstart:\n",
    "\n",
    "- Get this notebook and the other repo code onto a machine with the power to run Airoboros-l2-70b-3.1.2.Q4_K_M\n",
    "- Put the model into ./logical_models (relative to this notebook). \n",
    "- Run all the cells below and watch as the notebook generates questions, answers, and conversations based on Principles of Chemistry, Simple Sabotage, and Introduction to Philosophy.\n",
    "\n",
    "If you want to add your own texts, follow the instructions in list item #1 above.\n",
    "\n",
    "### Note: this notebook makes roughly 1/3 characters generated to be **mildly NSFW** by default. You will need to modify the character personality code in `./generation_functions/special_instructions.py` or use \"Assistant mode\" if you want something cleaner.\n",
    "\n",
    "## Customization:\n",
    "### Here are some ways you can adapt this notebook to your use case, along with a brief description of how to do so, arranged in increasing order of difficulty (this information is also available in the README):\n",
    "1. ***Change the source texts used to generate training data.*** You can do this in the cell right below this one. **IMPORTANT** the filenames of these should be formatted in a specific way, since the filenames are used as part of the prompts and in at least one regex. You need to have them be like: `[textname], by authorname`. You can also include the publication date after the author name if you want, but note that this will tend to bias most of the characters to live in the era of the textbook, which may or may not be what you want.\n",
    "\n",
    "2. ***Change the personalities of the characters generated.*** Currently, when generating characters for the multiturn conversation step, three randomly-selected traits are appended to the \"special instructions\" set of the prompt to constrain what kind of character is generated by the model. Depending on what kind of model you want to make, or even just if your preferences vary, then you will probably want to modify this a bit. You can do so in `./generation_functions/special_instructions.py`. A more in-depth description of the trait-axis system that I (over)thought up is available in the comments of that file.\n",
    "\n",
    "3. ***Change the constants.*** There are a few constant values in this notebook, and in `./generation_functions/constant_values.py`. These constants are tested, but if your use case requires special settings (e.g., you want to make conversations from more permutations of existing questions; or you think the character counts for the \"duplicate question/answer\" validation functions are too restrictive) then feel free to change the related setting. The most intuitive and least-likely-to-break-anything settings to change are rearrangements_to_take and double_check_counter. Beyond that... you'll need to figure out what the function does before changing it if you expect it to run.\n",
    "\n",
    "4. ***Assistant Mode*** Technically this could be considered part of 3), but it's different enough that I feel it warrants separate explanation. By default, the notebook is configured to produce RP-style data; \"Assistant mode\" is something you can toggle in the settings cell immediately below this one, which skips character and scenario generation and answers every question in a chat between a user and a helpful AI assistant (with no personality). In the limited testing I have done with this, **it seems that assistant mode is simple enough to work with 13b models** such as Flatorcamaid by Ikari. So if your compute or time are very limited, or you are using this for a more professional use case, feel free to turn this on.\n",
    "\n",
    "5. ***Change the model.*** This is as simple as switching the LOGICAL_MODEL parameter out for another one, but your mileage may vary, significantly. You will also have to adjust RoPE scaling for non-llama 2 models -- e.g., if you're using Mixtral, don't leave `rope_freq_scale=0.33`, which 3xes the context (you do not need 96k context, only 12k).\n",
    "\n",
    "6. ***Change the examples.*** If you change the examples you can completely overhaul what this notebook does, but this requires a lot of prompting skill and possibly huge amounts of time to get it working again (source: most of my last three months were spent prompting, and most of this prompting was spent on the examples). Unless you want to convert this notebook from question-and-answer generation to some completely other task, I'd recommend changing only the conversation generation prompts -- they're a bit less finnicky, and if you just want to change the kind of characters generated (maybe you want a different writing style) that's where you'd find the differences.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE NOTEBOOK SETTINGS AND CONSTANTS (some script file constants are in generation_functions/constants.py)\n",
    "\n",
    "# Put your desired quant of your desired model in the relevant directories\n",
    "\n",
    "\n",
    "# \"airoboros-l2-70b-3.1.2.Q4_K_M.gguf\" <- recommended for the large logical model\n",
    "# \"flatorcamaid-13b-v0.2.Q8_0.gguf\" <- recommended for the normal logical model\n",
    "# A6000s on Vast.ai are a good choice for running this notebook\n",
    "\n",
    "LOGICAL_MODEL = \"./logical_model/flatorcamaid-13b-v0.2.Q8_0.gguf\"  # model used for decision-making and base question generation (should be \"smart\")\n",
    "LARGE_LOGICAL_MODEL = \"./logical_model/airoboros-l2-70b-3.1.2.Q4_K_M.gguf\"\n",
    "ASSISTANT_MODE = False  # change to true if you want all conversations to be with an \"AI language model\" and not characters. Useful for more professional use cases.\n",
    "DOUBLE_CHECK_COUNTER = 3  # Set to 1 to check outputs only once; set to 2 to check twice; set to 3 to check thrice, etc. Set to 0 to break everything in vet_question_loop() and elsewhere. Set to -1 and cause the universe to implode?\n",
    "\n",
    "REARRANGEMENTS_TO_TAKE = 3  # How many of the possible permutations of tuples in a group to take and make multiturn convs out of. Adjust higher to get more data out of less text, but it might be a bit repetitive. NOTE your eval loss will be basically worthless if you aren't careful with how you shuffle your dataset when you're about to train.\n",
    "\n",
    "TEXT_MANUALLY_CLEANED = False  # If you've manually cut out all the parts of the text not worthy for questions, you can skip the first LLM step. NOTE I might actually recommend doing this if your source text is small, given how permissive the filtering prompt is; it really only disallows metadata.\n",
    "\n",
    "source_texts = [\n",
    "    \"Simple Sabotage, by the Office of Strategic Services, published 1944.txt\",\n",
    "    \"Principles of Chemistry, by Demitry Mendeleev, published 1897.txt\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Below: Defines and imports functions that you will probably use no matter what cells in the notebook you choose to run:\n",
    "\n",
    "### Also: The below will absolutely flood your terminal, due to heavy gbnf grammar use. I don't think this can be configured. Be prepared to scroll to reach the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import uuid\n",
    "\n",
    "multi_turn_convs_info_dir = \"./multi_turn_convs_info\"  # we generate all the information fed to the multiturn prompt, and generate the actual multiturn prompt, separately; since every step but the last is capable of being done by a 13b\n",
    "\n",
    "\n",
    "def make_id():\n",
    "    return str(uuid.uuid4())\n",
    "\n",
    "\n",
    "def write_output_to_file(output, directory, uuid):\n",
    "    # Ensure directory exists\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "\n",
    "    # Define the file path using the directory and UUID\n",
    "    file_path = os.path.join(directory, f\"{uuid}.txt\")\n",
    "\n",
    "    # Write the output to the file\n",
    "    with open(file_path, \"w\") as file:\n",
    "        file.write(output)\n",
    "\n",
    "    print(f\"Output written to {file_path}\")\n",
    "\n",
    "    # This is in no way best practices, but all my prompts being searchable and separate files is a good way to make my life easier.\n",
    "\n",
    "\n",
    "import pkgutil\n",
    "import importlib\n",
    "import generation_functions  # This is the package directory\n",
    "import sys\n",
    "\n",
    "sys.path.append(\"./generation_functions\")\n",
    "\n",
    "# First, import all modules so they can be reloaded\n",
    "for _, module_name, _ in pkgutil.iter_modules(\n",
    "    generation_functions.__path__, generation_functions.__name__ + \".\"\n",
    "):\n",
    "    importlib.import_module(module_name)\n",
    "\n",
    "# Now, reload each module and import all callable attributes\n",
    "for _, module_name, _ in pkgutil.iter_modules(\n",
    "    generation_functions.__path__, generation_functions.__name__ + \".\"\n",
    "):\n",
    "    # Reload the module\n",
    "    module = importlib.reload(sys.modules[module_name])\n",
    "    # Iterate through each attribute in the reloaded module\n",
    "    for attribute_name in dir(module):\n",
    "        # Retrieve the attribute\n",
    "        attribute = getattr(module, attribute_name)\n",
    "        if callable(attribute):\n",
    "            # If it's callable, it's a function or class, so you set it in the globals dictionary\n",
    "            globals()[attribute_name] = attribute\n",
    "\n",
    "\n",
    "def make_multiturn_conversation(info, logic_llm):\n",
    "    conv, conv_output = multi_turn_conversation(\n",
    "        info[0], info[1], info[2], info[3], logic_llm, assistant_mode=ASSISTANT_MODE\n",
    "    )  # based on what was originally: multi_turn_conversation(qa_tuples, character, scenario, scenario_plan, logic_llm)\n",
    "    write_output_to_file(conv_output, \"./multiturn_conversation_generations\", info[4])\n",
    "\n",
    "    return conv\n",
    "\n",
    "\n",
    "def ensure_multiple_answers_are_same(\n",
    "    info, conv, scenario\n",
    "):  # why is this a whole separate function? Once upon a time, LLMs were used in validation here, too. But programmatic validation SEEMS to catch the common problems. This is here so that I can add it back in if I have to.\n",
    "    \"\"\"Loop to ensure that the answer is consistent in the conversation and in the tuple.\"\"\"\n",
    "    retries = 0\n",
    "    c = conv\n",
    "    while retries < 2:  # try twice, since multiturn is an expensive operation\n",
    "        if call_all_processors(c[0], info[0]):  # if programmatic validation passes\n",
    "            return c\n",
    "\n",
    "        retries += 1\n",
    "        if retries >= 2:\n",
    "            return None\n",
    "        # If we're here, majority of relevance checks failed\n",
    "        print(\"----------------\\n\\n\\n\\nRETRYING!!!!\\n\\n\\n\\n----------------\")\n",
    "        # Broken info is 1) rare and 2) handled by the retry limit. We don't want to waste compute on regenerating info as they take time.\n",
    "        retry = make_multiturn_conversation(info, logic_llm)\n",
    "        if retry is not None:  # Note: retry CANNOT actually be None\n",
    "            c = retry\n",
    "        else:\n",
    "            # If we failed to generate a retry, don't waste compute\n",
    "            return None\n",
    "\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "import re\n",
    "from tqdm import tqdm\n",
    "import nltk\n",
    "\n",
    "nltk.download(\"punkt\")\n",
    "from nltk.tokenize import sent_tokenize\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    \"Gryphe/MythoMax-L2-13b\"\n",
    ")  # It doesn't matter what model goes here as long as it is sentencepiece\n",
    "\n",
    "\n",
    "def sentence_chunking_algorithm(file_path, tokenizer, max_token_length=400):\n",
    "    \"\"\"\n",
    "    This function takes a plaintext file and chunks it into sentences.\n",
    "\n",
    "    :param file_path: Path to the plaintext file\n",
    "    :param tokenizer: SentencePiece tokenizer\n",
    "    :param max_token_length: The maximum token length for a chunk of sentences\n",
    "    :return: List of sentence chunks with source text information\n",
    "    \"\"\"\n",
    "    sentence_chunks_with_source = []\n",
    "    current_chunk = []\n",
    "    token_count = 0\n",
    "    source_name = file_path.replace(\".txt\", \"\")\n",
    "\n",
    "    with open(file_path, \"r\", encoding=\"utf-8\") as f:\n",
    "        content = f.read()\n",
    "\n",
    "    # Remove Gutenberg header and footer\n",
    "    content = re.sub(\n",
    "        r\"^.*?START OF (THIS|THE) PROJECT GUTENBERG EBOOK.*$\\n\",\n",
    "        \"\",\n",
    "        content,\n",
    "        flags=re.MULTILINE,\n",
    "    )\n",
    "    content = re.sub(\n",
    "        r\"^.*?END OF (THIS|THE) PROJECT GUTENBERG EBOOK.*$\\n\",\n",
    "        \"\",\n",
    "        content,\n",
    "        flags=re.MULTILINE,\n",
    "    )\n",
    "\n",
    "    sentences = sent_tokenize(content)\n",
    "\n",
    "    for sentence in tqdm(sentences, desc=f\"Processing {file_path}\"):\n",
    "        sentence_token_count = len(tokenizer.encode(sentence))\n",
    "\n",
    "        if token_count + sentence_token_count <= max_token_length:\n",
    "            current_chunk.append(sentence)\n",
    "            token_count += sentence_token_count\n",
    "        else:\n",
    "            sentence_chunks_with_source.append((\" \".join(current_chunk), source_name))\n",
    "            current_chunk = [sentence]\n",
    "            token_count = sentence_token_count\n",
    "\n",
    "    # Add the last chunk if it exists\n",
    "    if current_chunk:\n",
    "        sentence_chunks_with_source.append((\" \".join(current_chunk), source_name))\n",
    "\n",
    "    return sentence_chunks_with_source\n",
    "\n",
    "\n",
    "sentence_chunks = []\n",
    "for source_text in source_texts:\n",
    "    sentence_chunks += sentence_chunking_algorithm(source_text, tokenizer)\n",
    "\n",
    "\n",
    "def fix_text(to_replace_arr, text):\n",
    "    for startup in to_replace_arr:\n",
    "        text = text.replace(startup[0], startup[1])\n",
    "    return text\n",
    "\n",
    "\n",
    "conversions = [(\"\\n\", \" \"), (\"  \", \" \")]\n",
    "\n",
    "paragraphs_processed = [\n",
    "    (fix_text(conversions, seq[0]), seq[1]) for seq in sentence_chunks\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(paragraphs_processed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paragraphs_processed[5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model into vram\n",
    "from llama_cpp import Llama\n",
    "\n",
    "logic_llm = Llama(\n",
    "    model_path=LOGICAL_MODEL,\n",
    "    n_gqa=8,\n",
    "    offload_kqv=True,\n",
    "    n_ctx=12000,\n",
    "    rope_freq_scale=0.33,\n",
    "    n_gpu_layers=100,\n",
    "    verbose=False,\n",
    ")  # load the logical LLM and offload everything"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Create directory if it doesn't exist\n",
    "output_dir = \"./worthy_for_questions\"\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "# Determine which paragraphs are worthy of making questions from\n",
    "if not TEXT_MANUALLY_CLEANED:\n",
    "    judged_worthy_for_questions = []\n",
    "    for idx, p in tqdm(enumerate(paragraphs_processed)):\n",
    "        # for idx, p in tqdm(enumerate(paragraphs_processed[:10])):\n",
    "        file_name = f\"{idx}.json\"\n",
    "        file_path = os.path.join(output_dir, file_name)\n",
    "\n",
    "        # Check if the judgement for this paragraph already exists\n",
    "        if os.path.isfile(file_path):\n",
    "            with open(file_path, \"r\") as file:\n",
    "                data = json.load(file)\n",
    "                print(\"LOADING: \", data)\n",
    "            if isinstance(data, str):\n",
    "                judged_worthy_for_questions.append((None, data[7:]))\n",
    "            else:\n",
    "                judged_worthy_for_questions.append(\n",
    "                    (data[\"paragraph\"], data[\"metadata\"])\n",
    "                )\n",
    "        else:\n",
    "            judgement = judge_paragraph(p, logic_llm)\n",
    "            judged_worthy_for_questions.append(judgement)\n",
    "\n",
    "            # Prepare the data to be written to the file\n",
    "            if judgement[0] is not None:\n",
    "                # The paragraph passed the judgement\n",
    "                data_to_write = {\"paragraph\": judgement[0], \"metadata\": judgement[1]}\n",
    "            else:\n",
    "                # The paragraph did not pass the judgement\n",
    "                data_to_write = f\"failed|{judgement[1]}\"\n",
    "\n",
    "            # Write the judgement to a unique file as JSON\n",
    "            with open(file_path, \"w\") as file:\n",
    "                json.dump(data_to_write, file)\n",
    "\n",
    "            # Debug messages\n",
    "            try:\n",
    "                if judgement[0] is not None:\n",
    "                    print(f\"DEBUG model decided that index {idx} was suitable\")\n",
    "                else:\n",
    "                    print(f\"DEBUG model decided that index {idx} was not suitable\")\n",
    "            except:\n",
    "                print(f\"DEBUG max retries exceeded for index {idx}\")\n",
    "\n",
    "else:\n",
    "    judged_worthy_for_questions = paragraphs_processed\n",
    "    # No need to write to file since paragraph chunking is deterministic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Graphing code generated by GPT-4. May be suboptimal/ugly.\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import Counter\n",
    "\n",
    "\n",
    "def filter_and_graph(tuples):\n",
    "    # Count the occurrences of None and non-None for each source text\n",
    "    source_counts = Counter()\n",
    "    for paragraph, source in tuples:\n",
    "        if paragraph is None:\n",
    "            source_counts[source] = source_counts.get(source, [0, 0])\n",
    "            source_counts[source][0] += 1\n",
    "        else:\n",
    "            source_counts[source] = source_counts.get(source, [0, 0])\n",
    "            source_counts[source][1] += 1\n",
    "\n",
    "    # Prepare data for the graph\n",
    "    labels = list(source_counts.keys())\n",
    "    none_counts = [source_counts[source][0] for source in labels]\n",
    "    non_none_counts = [source_counts[source][1] for source in labels]\n",
    "\n",
    "    # Plotting the graph\n",
    "    x = range(len(labels))\n",
    "    plt.bar(x, none_counts, width=0.4, label=\"Not suitable\", align=\"center\")\n",
    "    plt.bar(x, non_none_counts, width=0.4, label=\"Valid Paragraphs\", align=\"edge\")\n",
    "    plt.xlabel(\"Source Text\")\n",
    "    plt.ylabel(\"Number of Paragraphs\")\n",
    "    plt.title(\"Paragraphs Suitable for Questions by Source Text\")\n",
    "    plt.xticks(x, labels, rotation=\"vertical\")\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Filter out tuples with None and return the new list\n",
    "    filtered_list = [t for t in tuples if t[0] is not None]\n",
    "    return filtered_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_worthy_for_questions = filter_and_graph(judged_worthy_for_questions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(filtered_worthy_for_questions[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Control flow helpers\n",
    "\n",
    "# Change this value to change how many times the checks must pass consecutively for a thing to be accepted\n",
    "\n",
    "import logging\n",
    "from math import ceil\n",
    "\n",
    "\n",
    "# Setup logging\n",
    "# Except I actually don't use this because I switched to print() because jupyter is annoying with logging\n",
    "logging.basicConfig(\n",
    "    filename=\"data_generation.log\",\n",
    "    filemode=\"a\",\n",
    "    format=\"%(asctime)s - %(levelname)s - %(message)s\",\n",
    "    level=logging.INFO,\n",
    ")\n",
    "\n",
    "\n",
    "def vet_answer_accuracy_loop(qa_tuple, total_retries, run_id):\n",
    "    try:\n",
    "        qtuple = qa_tuple\n",
    "        print(\n",
    "            f\"\\n\\nStarting ACCURACY loop for question: {qtuple[0]}, context: {qtuple[2]}\"\n",
    "        )\n",
    "        passed_checks = 0\n",
    "        times_checked = 0\n",
    "        dissenting_reasoning = \"\"\n",
    "        while times_checked < DOUBLE_CHECK_COUNTER:\n",
    "            print(\n",
    "                f\"\\n\\nACCURACY CALL CHECK ANSWER: {qtuple[0]}, context: {qtuple[2]}, retries: {total_retries}, dissenting reasoning: {dissenting_reasoning}\"\n",
    "            )\n",
    "            judgement, answer_accuracy_output = check_answer(qtuple, logic_llm)\n",
    "            write_output_to_file(\n",
    "                answer_accuracy_output, \"./check_answer_accuracy_generations\", run_id\n",
    "            )\n",
    "            if not judgement[0]:  # if not accurate\n",
    "                dissenting_reasoning = judgement[1]\n",
    "            else:\n",
    "                passed_checks += 1\n",
    "            times_checked += 1\n",
    "            if passed_checks >= ceil(DOUBLE_CHECK_COUNTER / 2):\n",
    "                break\n",
    "            failed_checks = times_checked - passed_checks\n",
    "            if failed_checks >= ceil(DOUBLE_CHECK_COUNTER / 2):\n",
    "                break\n",
    "\n",
    "        if passed_checks >= ceil(DOUBLE_CHECK_COUNTER / 2):  # if question checks passed\n",
    "            print(f\"\\n\\ANSWER ACCURACY CHECKS PASSED retries: {total_retries}\")\n",
    "            return qtuple\n",
    "        else:\n",
    "            # Generate new question and restart the loop\n",
    "            print(\n",
    "                f\"\\n\\nACCURACY CHECKS FAILED - SENDING BACK TO QUESTION LOOP retries: {total_retries}\"\n",
    "            )\n",
    "            total_retries += 1\n",
    "            qtuple, generate_new_q_output = generate_new_question(qtuple, logic_llm)\n",
    "            write_output_to_file(\n",
    "                generate_new_q_output, \"./regenerate_question_generations\", run_id\n",
    "            )\n",
    "            vet_question_loop(\n",
    "                qtuple, total_retries, run_id=run_id.split(\"--subquestion--\")[0]\n",
    "            )  # going to get one hell of a call stack by the end of this, but it should be fine\n",
    "    except Exception as e:\n",
    "        print(\"!!ERROR!!\")\n",
    "        print(e)\n",
    "        pass\n",
    "\n",
    "    return (None, None, None, qtuple[3])\n",
    "\n",
    "\n",
    "def vet_answer_relevance_loop(qa_tuple, total_retries, run_id):\n",
    "    try:\n",
    "        qtuple = qa_tuple\n",
    "        print(\n",
    "            f\"\\n\\nStarting RELEVANCE loop for question: {qtuple[0]}, context: {qtuple[2]}\"\n",
    "        )\n",
    "        passed_checks = 0\n",
    "        times_checked = 0\n",
    "        dissenting_reasoning = \"\"\n",
    "        while times_checked < DOUBLE_CHECK_COUNTER:\n",
    "            print(\n",
    "                f\"\\n\\nRELEVANCE CALL CHECK ANSWER: {qtuple[0]}, context: {qtuple[2]}, retries: {total_retries}, dissenting reasoning: {dissenting_reasoning}\"\n",
    "            )\n",
    "            judgement, answer_relevancy_output = check_answer_relevancy_with_text(\n",
    "                qtuple, logic_llm\n",
    "            )\n",
    "            write_output_to_file(\n",
    "                answer_relevancy_output, \"./check_answer_relevancy_generations\", run_id\n",
    "            )\n",
    "            if not judgement[0]:  # if not relevant\n",
    "                dissenting_reasoning = judgement[1]\n",
    "            else:\n",
    "                passed_checks += 1\n",
    "            times_checked += 1\n",
    "            if passed_checks >= ceil(DOUBLE_CHECK_COUNTER / 2):\n",
    "                break\n",
    "            failed_checks = times_checked - passed_checks\n",
    "            if failed_checks >= ceil(DOUBLE_CHECK_COUNTER / 2):\n",
    "                break\n",
    "\n",
    "        if passed_checks >= ceil(DOUBLE_CHECK_COUNTER / 2):\n",
    "            print(f\"\\n\\nRELEVANCE CHECKS PASSED\")\n",
    "            return vet_answer_accuracy_loop(qtuple, total_retries, run_id)\n",
    "        else:\n",
    "            print(f\"\\n\\nRELEVANCE CHECKS FAILED - SENDING BACK TO QUESTION LOOP\")\n",
    "            total_retries += 1\n",
    "            qtuple, generate_new_q_output = generate_new_question(qtuple, logic_llm)\n",
    "            write_output_to_file(\n",
    "                generate_new_q_output, \"./regenerate_question_generations\", run_id\n",
    "            )\n",
    "            return vet_question_loop(\n",
    "                qtuple, total_retries, run_id=run_id.split(\"--subquestion--\")[0]\n",
    "            )\n",
    "    except Exception as e:\n",
    "        print(\"!!ERROR!!\")\n",
    "        print(e)\n",
    "        pass\n",
    "\n",
    "    return (None, None, None, qtuple[3])\n",
    "\n",
    "\n",
    "def vet_question_loop(qa_tuple, question_group_id=None, total_retries=0):\n",
    "    try:\n",
    "        qtuple = qa_tuple\n",
    "        print(\n",
    "            f\"\\n\\nStarting QUESTION loop for question: {qtuple[0]}, context: {qtuple[2]}\"\n",
    "        )\n",
    "        while total_retries <= 4:\n",
    "            run_id = question_group_id + \"--subquestion--\" + make_id()\n",
    "            passed_checks = 0\n",
    "            times_checked = 0\n",
    "            dissenting_reasoning = \"\"\n",
    "            while times_checked < DOUBLE_CHECK_COUNTER:\n",
    "                print(\n",
    "                    f\"\\n\\nQUESTION CALL CHECK ANSWER: {qtuple[0]}, context: {qtuple[2]}, retries: {total_retries}, dissenting reasoning: {dissenting_reasoning}\"\n",
    "                )\n",
    "                judgement, check_q_output = check_question(qtuple, logic_llm)\n",
    "                write_output_to_file(\n",
    "                    check_q_output, \"./check_question_generations\", run_id\n",
    "                )\n",
    "                if not judgement[0]:  # if not relevant\n",
    "                    dissenting_reasoning = judgement[1]\n",
    "                else:\n",
    "                    passed_checks += 1\n",
    "                times_checked += 1\n",
    "                if passed_checks >= ceil(DOUBLE_CHECK_COUNTER / 2):\n",
    "                    break\n",
    "                failed_checks = times_checked - passed_checks\n",
    "                if failed_checks >= ceil(DOUBLE_CHECK_COUNTER / 2):\n",
    "                    break\n",
    "\n",
    "            if passed_checks >= ceil(\n",
    "                DOUBLE_CHECK_COUNTER / 2\n",
    "            ):  # if all question checks passed\n",
    "                print(f\"\\n\\nQUESTION CHECKS PASSED retries: {total_retries}\")\n",
    "                return vet_answer_relevance_loop(qtuple, total_retries, run_id)\n",
    "            else:\n",
    "                # Generate new question and restart the loop\n",
    "                print(\n",
    "                    f\"\\n\\nQUESTION CHECKS FAILED - GENERATING NEW QUESTION retries: {total_retries}\"\n",
    "                )\n",
    "                total_retries += 1\n",
    "                if (\n",
    "                    total_retries <= 4\n",
    "                ):  # only regen question if we're not already at max regens\n",
    "                    qtuple, generate_new_q_output = generate_new_question(\n",
    "                        qtuple, logic_llm\n",
    "                    )\n",
    "                    write_output_to_file(\n",
    "                        generate_new_q_output,\n",
    "                        \"./regenerate_question_generations\",\n",
    "                        run_id,\n",
    "                    )\n",
    "                    print(\"New question: \", qtuple)\n",
    "                # no calling of vet_question_loop, since we're already in a while loop\n",
    "    except Exception as e:\n",
    "        print(\"!!ERROR!!\")\n",
    "        print(e)\n",
    "\n",
    "    return (None, None, None, qtuple[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# control flow\n",
    "import json\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "\n",
    "# Directory for QA tuples\n",
    "qa_tuples_dir = \"./qatuples_raw\"\n",
    "if not os.path.exists(qa_tuples_dir):\n",
    "    os.makedirs(qa_tuples_dir)\n",
    "\n",
    "# Initialize vetted_qa_tuples\n",
    "vetted_qa_tuples = []  # tuple list of qa tuples that have been judged good\n",
    "\n",
    "# Attempt to initialize filtered_worthy_for_questions\n",
    "try:\n",
    "    _ = filtered_worthy_for_questions\n",
    "except NameError:\n",
    "    filtered_worthy_for_questions = []\n",
    "\n",
    "if not filtered_worthy_for_questions:\n",
    "    # Load all files in the qa_tuples_dir if filtered_worthy_for_questions is not initialized\n",
    "    existing_files = glob.glob(os.path.join(qa_tuples_dir, \"*.json\"))\n",
    "    for file_path in existing_files:\n",
    "        with open(file_path, \"r\") as file:\n",
    "            qa_tuple = tuple(json.load(file))\n",
    "        vetted_qa_tuples.append(qa_tuple)\n",
    "\n",
    "\n",
    "else:\n",
    "    for idx, para in enumerate(tqdm(filtered_worthy_for_questions)):\n",
    "        # for idx, para in enumerate(tqdm(filtered_worthy_for_questions[:10])): # Use this instead if you are just testing all steps of the notebook\n",
    "        try:\n",
    "            existing_files = glob.glob(\n",
    "                os.path.join(qa_tuples_dir, f\"para_{idx}_*.json\")\n",
    "            )  # check if qs already exist\n",
    "\n",
    "            if len(existing_files) > 0:  # If files exist, skip this paragraph entirely\n",
    "                print(f\"Skipping para_{idx} as files already exist; loading said files\")\n",
    "                for file_path in existing_files:\n",
    "                    with open(file_path, \"r\") as file:\n",
    "                        qa_tuple = tuple(json.load(file))\n",
    "                    vetted_qa_tuples.append(qa_tuple)\n",
    "                continue\n",
    "\n",
    "            question_group_id = make_id()\n",
    "            print(f\"\\n\\n\\nOUTER LOOP CALL GENERATE QPLAN para: {para}, \\n\\n idx: {idx}\")\n",
    "            plan, questions_plan_output = generate_questions_plan(para, logic_llm)\n",
    "            write_output_to_file(\n",
    "                questions_plan_output, \"./question_plan_generations\", question_group_id\n",
    "            )\n",
    "            print(\n",
    "                f\"\\n\\n\\nOUTER LOOP CALL GENERATE Q: {para}, \\n\\n idx: {idx} \\n\\n plan: {plan}\"\n",
    "            )\n",
    "            question_answer_tuples, question_generation_output = generate_questions(\n",
    "                para, plan, logic_llm\n",
    "            )\n",
    "            write_output_to_file(\n",
    "                question_generation_output,\n",
    "                \"./question_generation_generations\",\n",
    "                question_group_id,\n",
    "            )\n",
    "            for qnum, question_answer_tuple in enumerate(question_answer_tuples):\n",
    "                print(\n",
    "                    f\"\\n\\n=======!!=BEGIN VETTING QA TUPLE {idx}_{qnum}=!!=======\\n\\n\"\n",
    "                )\n",
    "                good_qa_tuple = vet_question_loop(\n",
    "                    question_answer_tuple, question_group_id=question_group_id\n",
    "                )\n",
    "\n",
    "                # Write resulting question file if the tuple is not None\n",
    "                if good_qa_tuple[0] is not None:\n",
    "                    file_path = os.path.join(qa_tuples_dir, f\"para_{idx}_q_{qnum}.json\")\n",
    "                    with open(file_path, \"w\") as file:\n",
    "                        json.dump(good_qa_tuple, file, indent=4)\n",
    "\n",
    "                vetted_qa_tuples.append(\n",
    "                    good_qa_tuple\n",
    "                )  # We must filter out all None values at the end; but appending Nones lets us know where things went wrong, and how often.\n",
    "        except Exception as e:\n",
    "            print(f\"Q ERROR: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    \"-------------- QUESTIONS CREATED ------------- STATS SO FAR (may be wrong if run was continued from interruption):\"\n",
    ")\n",
    "nones = list(filter(lambda x: x[0] is None, vetted_qa_tuples))\n",
    "print(f\"Nones: {len(nones)}\")\n",
    "print(f\"Non-nones: {len(vetted_qa_tuples) - len(nones)}\")\n",
    "print(f\"Total: {len(vetted_qa_tuples)}\")\n",
    "# filter out all None values\n",
    "vetted_qa_tuples = [qa for qa in vetted_qa_tuples if qa[0] is not None]\n",
    "print(\"---------------- ONTO EXAMPLES GENERATION-------------------\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check for and fix the common mistake: mentioning \"the text\".\n",
    "# TODO refactor to be continuable, should take like 30 mins at most\n",
    "\n",
    "writepath = \"./qatuples_revised\"\n",
    "\n",
    "\n",
    "# Assuming vetted_qa_tuples is a list that might or might not exist\n",
    "try:\n",
    "    _ = vetted_qa_tuples\n",
    "except NameError:\n",
    "    vetted_qa_tuples = []\n",
    "\n",
    "\n",
    "# Load all files at the start if vetted_qa_tuples is empty\n",
    "if not vetted_qa_tuples:\n",
    "    # Check if the directory exists\n",
    "    if os.path.exists(writepath):\n",
    "        # List all files in directory\n",
    "        for file_name in os.listdir(writepath):\n",
    "            file_path = os.path.join(writepath, file_name)\n",
    "            try:\n",
    "                with open(file_path, \"r\", encoding=\"utf-8\") as f:\n",
    "                    content = f.read()\n",
    "                    print(f\"Loading file: {file_path}\")\n",
    "                    if content == \"failed\":\n",
    "                        vetted_qa_tuples.append(None)\n",
    "                    else:\n",
    "                        try:\n",
    "                            data = json.loads(content)\n",
    "                            vetted_qa_tuples.append(\n",
    "                                (data[0], data[1], data[2], data[3])\n",
    "                            )\n",
    "                        except json.JSONDecodeError:\n",
    "                            print(\"JSON decode error with the contents:\", content)\n",
    "                            vetted_qa_tuples.append(None)\n",
    "            except Exception as e:\n",
    "                print(f\"Error reading {file_path}: {e}\")\n",
    "\n",
    "else:\n",
    "    old_tuples = vetted_qa_tuples.copy()\n",
    "\n",
    "    for idx, tup in enumerate(vetted_qa_tuples):\n",
    "        file_path = os.path.join(writepath, f\"revised_{idx}.json\")\n",
    "        if os.path.exists(file_path):\n",
    "            with open(file_path, \"r\", encoding=\"utf-8\") as f:\n",
    "                content = f.read()  # Read the file once and store its content\n",
    "                print(file_path)\n",
    "                if content == \"failed\":\n",
    "                    print(\"Loaded failed file\")\n",
    "                    vetted_qa_tuples[idx] = None\n",
    "                    continue\n",
    "                print(\"Loaded file:\")\n",
    "                print(content)\n",
    "                # Reset the file pointer to the beginning if you need to read again or convert string back to JSON\n",
    "                try:\n",
    "                    data = json.loads(content)  # Convert the string back to JSON\n",
    "                    vetted_qa_tuples[idx] = (data[0], data[1], data[2], data[3])\n",
    "                    continue\n",
    "                except json.JSONDecodeError:\n",
    "                    print(\"JSON decode error with the contents:\", content)\n",
    "                    # Handle the error appropriately\n",
    "\n",
    "        try:\n",
    "            revision_id = make_id()\n",
    "            revision, revision_output = check_qatuple_context(tup, logic_llm)\n",
    "            write_output_to_file(\n",
    "                revision_output, \"./question_context_revision_generations\", revision_id\n",
    "            )  # incidentally, identifying the problem and fixing it in the same step (without another planning step) works a lot better than identifying it and then trying to fix it in the next step.\n",
    "            if isinstance(revision[0], str):  # if the thing was reworded\n",
    "                vetted_qa_tuples[idx] = revision\n",
    "            elif not revision[0]:\n",
    "                vetted_qa_tuples[\n",
    "                    idx\n",
    "                ] = None  # prepare item for deletion later; right now we just store it as None because indexes\n",
    "            # else, if it passed, we just leave it be.\n",
    "\n",
    "            # Write in-progress\n",
    "            if not os.path.exists(writepath):\n",
    "                os.makedirs(writepath)\n",
    "\n",
    "            if vetted_qa_tuples[idx]:\n",
    "                with open(file_path, \"w\") as file:\n",
    "                    json.dump(vetted_qa_tuples[idx], file, indent=4)\n",
    "            else:\n",
    "                with open(file_path, \"w\") as file:\n",
    "                    file.write(\"failed\")\n",
    "\n",
    "        except Exception as e:\n",
    "            print(\"!!! ERROR!\", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print stats related to revised qatuples, and filter out nones (questions that were unanswerable due to lack of context).\n",
    "import json\n",
    "import os\n",
    "\n",
    "print(\"-------------- QUESTIONS REVISED ------------- STATS SO FAR:\")\n",
    "nones = list(filter(lambda x: x is None, vetted_qa_tuples))\n",
    "print(f\"Nones: {len(nones)}\")\n",
    "print(f\"Non-nones: {len(vetted_qa_tuples) - len(nones)}\")\n",
    "print(f\"Total: {len(vetted_qa_tuples)}\")\n",
    "# filter out all None values\n",
    "vetted_qa_tuples = [qa for qa in vetted_qa_tuples if qa is not None]\n",
    "print(\"---------------- ONTO EXAMPLES GENERATION-------------------\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Group tuples for multiturn example generation (by chunk of source text) and then run that helper (so that we can make multiturn conversations from questions based on the same paragraphs)\n",
    "def group_by_text(tuples_list):\n",
    "    # Dictionary to hold the groups with text as the key\n",
    "    groups = {}\n",
    "\n",
    "    # Iterate over each tuple in the list\n",
    "    for question, answer, text, textname in tuples_list:\n",
    "        # If the text is not yet a key in the dictionary, add it with an empty list\n",
    "        if text not in groups:\n",
    "            groups[text] = []\n",
    "\n",
    "        # Append the current tuple to the appropriate list\n",
    "        groups[text].append((question, answer, text, textname))\n",
    "\n",
    "    # Return the values of the dictionary, which are the lists of tuples grouped by text; also remove duplicates\n",
    "    return [identify_duplicates(group) for group in list(groups.values())]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_tuples_by_paragraph = group_by_text(vetted_qa_tuples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import ceil\n",
    "\n",
    "# multiturn helpers\n",
    "# These will probably be used for multiturn rapid-fire answering.\n",
    "\n",
    "\n",
    "# Idea: use multiple short answers to train the task of answering multiple questions in one response. Two-three short answers per response should be enough.\n",
    "def make_multiturn_character(qa_tuples, conv_id):\n",
    "    if (\n",
    "        ASSISTANT_MODE\n",
    "    ):  # If assistant mode is on, multiturn convs will have hardcoded information in its prompt file; but we still need to put something in the file\n",
    "        return \"will_be_replaced\", \"will_be_replaced\"\n",
    "    plan, instructions, card_plan_output = create_character_card_plan_many_tuples(\n",
    "        qa_tuples, logic_llm\n",
    "    )  # I will reuse the many tuples function for short question-answers, there's a lot of prompting in here already\n",
    "    write_output_to_file(card_plan_output, \"./multiturn_card_plan_generations\", conv_id)\n",
    "    char, char_output = create_character_card_many_tuples(\n",
    "        qa_tuples, plan, instructions, logic_llm\n",
    "    )  # creates a character card\n",
    "    write_output_to_file(char_output, \"./multiturn_card_generations\", conv_id)\n",
    "    return char, instructions\n",
    "\n",
    "\n",
    "def make_multiturn_scenario(qa_tuples, character, conv_id):\n",
    "    if (\n",
    "        ASSISTANT_MODE\n",
    "    ):  # If assistant mode is on, multiturn convs will have hardcoded information in its prompt file; but we still need to put something in the file\n",
    "        return \"will_be_replaced\", \"will_be_replaced\"\n",
    "    plan, scenario_plan_output = create_scenario_plan_many_tuples(\n",
    "        qa_tuples, character, logic_llm\n",
    "    )\n",
    "    write_output_to_file(\n",
    "        scenario_plan_output, \"./multiturn_scenario_plan_generations\", conv_id\n",
    "    )\n",
    "    scenario, scenario_output = create_scenario_many_tuples(\n",
    "        qa_tuples, character, plan, logic_llm\n",
    "    )  # creates a scenario based on a character card and question/answer tuple\n",
    "    write_output_to_file(scenario_output, \"./multiturn_scenario_generations\", conv_id)\n",
    "    return scenario, plan\n",
    "\n",
    "\n",
    "def make_multiturn_conversation_info(qa_tuples, logic_llm):\n",
    "    conv_id = make_id()\n",
    "    if (\n",
    "        ASSISTANT_MODE\n",
    "    ):  # If assistant mode is on, multiturn convs will have hardcoded information in its prompt file; but we still need to put something in the file\n",
    "        return (qa_tuples, \"will\", \"be\", \"replaced\", conv_id)\n",
    "    # thought_plan = create_thought_plan_many_tuples(qa_tuples,character,scenario,logic_llm) # There IS a way to make multiturn chain of thought answering work: generate each pair of messages using a separate prompt or a separate function, each of which has only the thought plan for that question/answer pair. But simply cramming in all the step-by-step things will confuse the hell out of the poor model. So for the first release version we're skipping it and just giving the response, with no reasoning, in the multiturn convs.\n",
    "    character, instructions = make_multiturn_character(qa_tuples, conv_id)\n",
    "    scenario, scenario_plan = make_multiturn_scenario(qa_tuples, character, conv_id)\n",
    "\n",
    "    return (qa_tuples, character, scenario, scenario_plan, conv_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "import itertools\n",
    "\n",
    "\n",
    "if not os.path.exists(multi_turn_convs_info_dir):\n",
    "    os.makedirs(multi_turn_convs_info_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_turn_convs_info = []\n",
    "for idx, group in enumerate(qa_tuples_by_paragraph):\n",
    "    all_permutations = list(itertools.permutations(group))\n",
    "\n",
    "    sample_size = min(REARRANGEMENTS_TO_TAKE, len(all_permutations))\n",
    "    sampled_permutations = random.sample(all_permutations, sample_size)\n",
    "\n",
    "    group_convs_info = []\n",
    "\n",
    "    for iter, perm in enumerate(sampled_permutations):\n",
    "        file_path = os.path.join(multi_turn_convs_info_dir, f\"info_{idx}_{iter}.json\")\n",
    "\n",
    "        # Skip if file already exists\n",
    "        if not os.path.exists(file_path):\n",
    "            info = make_multiturn_conversation_info(perm, logic_llm)\n",
    "\n",
    "            if info is not None:\n",
    "                with open(file_path, \"w\") as file:\n",
    "                    json.dump(info, file, indent=4)\n",
    "\n",
    "            group_convs_info.append(info)\n",
    "        else:\n",
    "            print(f\"Skipped generating {file_path} as it already exists\")\n",
    "\n",
    "    multi_turn_convs_info.append(group_convs_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del logic_llm  # Apparently frees the model, haven't tested this extensively"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stop Here, Restart the Notebook, and Reimport Everything IF you are doing 2-step Generation (where you do the easy bits with a small model, and the hard bit with a large one)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model into vram\n",
    "from llama_cpp import Llama\n",
    "\n",
    "logic_llm = Llama(\n",
    "    model_path=LARGE_LOGICAL_MODEL,\n",
    "    n_gqa=8,\n",
    "    offload_kqv=True,\n",
    "    n_ctx=12000,\n",
    "    rope_freq_scale=0.33,\n",
    "    n_gpu_layers=100,\n",
    "    verbose=False,\n",
    ")  # load the logical LLM and offload everything"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "\n",
    "def read_json_files_info(directory):\n",
    "    # Create a list to hold the tuples\n",
    "    tuple_list = []\n",
    "\n",
    "    # Get all the .json files in the directory, sorted\n",
    "    json_files = sorted([f for f in os.listdir(directory) if f.endswith(\".json\")])\n",
    "\n",
    "    # Read each file and convert the contents\n",
    "    for file in json_files:\n",
    "        with open(os.path.join(directory, file), \"r\") as f:\n",
    "            data = json.load(f)\n",
    "            # Ensure the data is in the correct format before converting to tuple\n",
    "            if (\n",
    "                isinstance(data, list)\n",
    "                and len(data) == 5\n",
    "                and isinstance(data[0], list)\n",
    "                and all(len(item) == 4 for item in data[0])\n",
    "                and all(isinstance(i, str) for i in data[1:])\n",
    "            ):\n",
    "                tuple_list.append((data[0], data[1], data[2], data[3], data[4]))\n",
    "\n",
    "    return tuple_list\n",
    "\n",
    "\n",
    "convs_info = read_json_files_info(multi_turn_convs_info_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "import itertools\n",
    "\n",
    "multi_turn_convs_dir = \"./multi_turn_convs\"\n",
    "if not os.path.exists(multi_turn_convs_dir):\n",
    "    os.makedirs(multi_turn_convs_dir)\n",
    "\n",
    "multi_turn_convs = []\n",
    "for idx, info in enumerate(convs_info):\n",
    "    file_path = os.path.join(multi_turn_convs_dir, f\"conv_{idx}.json\")\n",
    "\n",
    "    # Skip if file already exists\n",
    "    if not os.path.exists(file_path):\n",
    "        conv = make_multiturn_conversation(info, logic_llm)\n",
    "        final_conv = ensure_multiple_answers_are_same(info, conv, logic_llm)\n",
    "\n",
    "        if final_conv is not None:\n",
    "            with open(file_path, \"w\") as file:\n",
    "                json.dump(final_conv, file, indent=4)\n",
    "\n",
    "        multi_turn_convs.append(final_conv)\n",
    "    else:\n",
    "        with open(file_path, \"r\", encoding=\"utf-8\") as f:\n",
    "            data = json.load(f)\n",
    "            multi_turn_convs.append(data)\n",
    "        print(f\"Skipped generating {file_path} as it already exists\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Yay! Now you have a dataset!\n",
    "### GPT wrote the cell below. I think it successfully converts things to ShareGPT format for use with axolotl, but I am not sure because I don't know that format very well and haven't used Axolotl. However, the json produced by the second function looks fine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "\n",
    "def convert_directory_to_list(directory_path):\n",
    "    master_list = []\n",
    "    simplified_list = []\n",
    "\n",
    "    for filename in os.listdir(directory_path):\n",
    "        if filename.endswith(\".json\"):\n",
    "            filepath = os.path.join(directory_path, filename)\n",
    "            with open(filepath, \"r\") as file:\n",
    "                data = json.load(file)\n",
    "                if isinstance(data, list) and all(\n",
    "                    isinstance(item, (list, str)) for item in data\n",
    "                ):\n",
    "                    master_list.append(data)\n",
    "\n",
    "                    # Extract and process conversation\n",
    "                    conversation, primary_char_desc = data[0], data[1]\n",
    "                    primary_char_name = extract_name(primary_char_desc)\n",
    "                    dialogues = extract_conversation(conversation)\n",
    "\n",
    "                    # Convert to simplified format\n",
    "                    simplified_conversations = []\n",
    "                    for i, (charname, message) in enumerate(\n",
    "                        dialogues\n",
    "                    ):  # Skipping the first message\n",
    "                        from_person = (\n",
    "                            \"human\" if charname == primary_char_name else \"gpt\"\n",
    "                        )\n",
    "                        simplified_conversations.append(\n",
    "                            {\"from\": from_person, \"value\": f\"{charname}: {message}\"}\n",
    "                        )\n",
    "\n",
    "                    if simplified_conversations:  # If there are any conversations\n",
    "                        simplified_list.append(\n",
    "                            {\"conversations\": simplified_conversations}\n",
    "                        )\n",
    "\n",
    "    # Write the master list to a new .jsonl file\n",
    "    with open(\"master_list.jsonl\", \"w\") as file:\n",
    "        for item in master_list:\n",
    "            file.write(json.dumps(item) + \"\\n\")\n",
    "\n",
    "    # Write the simplified data to a different .jsonl file\n",
    "    with open(\"simplified_data.jsonl\", \"w\") as file:\n",
    "        for item in simplified_list:\n",
    "            file.write(json.dumps(item) + \"\\n\")\n",
    "\n",
    "    print(\n",
    "        \"Conversion complete. Master list written to 'master_list.json'. Simplified data written to 'simplified_data.json'.\"\n",
    "    )\n",
    "\n",
    "\n",
    "convert_directory_to_list(\"./multi_turn_convs/\")\n",
    "\n",
    "# Write version where each line in the conversation is separate, so that you do not need to write an extract_convs function if you want a custom format.\n",
    "\n",
    "\n",
    "def convert_directory_and_process_conversations(directory_path):\n",
    "    master_list = []\n",
    "\n",
    "    for filename in os.listdir(directory_path):\n",
    "        if filename.endswith(\".json\"):\n",
    "            filepath = os.path.join(directory_path, filename)\n",
    "            with open(filepath, \"r\") as file:\n",
    "                data = json.load(file)\n",
    "\n",
    "                if isinstance(data, list) and all(\n",
    "                    isinstance(item, (list, str)) for item in data\n",
    "                ):\n",
    "                    # Extract and process the conversation part\n",
    "                    conversations = extract_conversation(data[0])\n",
    "                    # Convert tuples back to the formatted string as required\n",
    "                    data[0] = [\n",
    "                        f\"{charname}: {message}\" for charname, message in conversations\n",
    "                    ]\n",
    "                    master_list.append(data)\n",
    "                else:\n",
    "                    print(f\"File {filename} is not in the expected format.\")\n",
    "\n",
    "    # Write the master list to a new file\n",
    "    with open(\"processed_master_list.json\", \"w\") as file:\n",
    "        json.dump(master_list, file)\n",
    "\n",
    "    print(\n",
    "        \"Conversion complete. The processed master list is written to 'processed_master_list.json'.\"\n",
    "    )\n",
    "\n",
    "\n",
    "# Usage: replace 'path_to_directory' with the actual directory path\n",
    "convert_directory_and_process_conversations(\"./multi_turn_convs/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./processed_master_list.json\") as f:\n",
    "    first = f.read()\n",
    "    data = json.loads(first)\n",
    "\n",
    "\n",
    "# For curiosity's sake, you can find out how many lines of dialogue you generated\n",
    "def filter_and_flatten(lst):\n",
    "    # Initialize an empty list to hold the flattened elements\n",
    "    flat_list = []\n",
    "\n",
    "    # Loop through each sublist in the main list\n",
    "    for sublst in lst:\n",
    "        # Check if the first element of the sublist is itself a list (subsublist1)\n",
    "        if isinstance(sublst[0], list):\n",
    "            # Extend the flat_list with the elements from subsublist1\n",
    "            flat_list.extend(sublst[0])\n",
    "\n",
    "    return flat_list\n",
    "\n",
    "\n",
    "len(filter_and_flatten(data))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.15 ('mlp')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "4976e0179d97dd6d59b1329a76e601e17b789c2571b41c8b57f5fd69821c0dd3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}