FineTuning_v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "235d821b-1ff8-4ef6-8f0b-559c95254479",
   "metadata": {},
   "source": [
    "# Slackbot Example\n",
    "\n",
    "In this notebook, you’ll see how to train BeepGPT on your Slack history using only OpenAI’s API’s and open-source Python libraries - Data Science PhD not required.\n",
    "\n",
    "We'll train BeepGPT in three steps:\n",
    "1. Build an initial set of training examples from a Slack Export\n",
    "2. Use few-shot learning with ChatGPT to clean our training examples\n",
    "3. Send our training data to OpenAI and create a fine-tuned model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d550080",
   "metadata": {},
   "source": [
    "\n",
    "First lets install the libraries we will use below. And initialize our OpenAI session with an API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70440303",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q backoff pandas pyarrow openai scikit-learn kaskada==0.6.0a4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19f2cb2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import getpass\n",
    "\n",
    "# Initialize OpenAI\n",
    "openai.api_key = getpass.getpass('OpenAI: API Key')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bfd95ea",
   "metadata": {},
   "source": [
    "## 1. Build an initial set of training examples from a Slack Export\n",
    "\n",
    "The steps involved here include:\n",
    "* 1.1 Convert the Slack Export into a format that can be consumed by Kaskada\n",
    "* 1.2 Use Kaskada to break the Slack Export into a set of conversations\n",
    "* 1.3 Generate the initial training examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d140e2e4",
   "metadata": {},
   "source": [
    "### 1.1 Convert the Slack Export into a format that can be consumed by Kaskada\n",
    "\n",
    "Historical slack messages can be exported by following the instructions in Slack's [Export your workspace data](https://slack.com/help/articles/201658943-Export-your-workspace-data) web page. We'll use these messages to teach BeepGPT about the members of your workspace.\n",
    "\n",
    "The export from Slack contains a zip of numererous folders and files. After uncompressing the archive, there are folders for each public channel in your Slack workspace. Inside each folder are json files for each day, which each contain all the events from the day.\n",
    "\n",
    "We execute a short python script (utilizing pandas), to concatenate all the data files together into a single parquet file.\n",
    "\n",
    "Parquet files store data in columns instead of rows. Some benefits of Parquet include:\n",
    "* Fast queries that can fetch specific column values without reading full row data\n",
    "* Highly efficient column-wise compression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed117a7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "def get_file_df(json_path):\n",
    "    df = pd.read_json(json_path, precise_float=True)\n",
    "    # drop rows where subType is not null\n",
    "    if \"subtype\" in df.columns:\n",
    "        df = df[df[\"subtype\"].isnull()]\n",
    "    # only keep these columns\n",
    "    df = df[df.columns.intersection([\"ts\", \"user\", \"text\", \"thread_ts\"])]\n",
    "    return df\n",
    "\n",
    "def get_channel_df(channel_path):\n",
    "    dfs = []\n",
    "    for root, dirs, files in os.walk(channel_path):\n",
    "        for file in files:\n",
    "            dfs.append(get_file_df(os.path.join(root, file)))\n",
    "    return pd.concat(dfs, ignore_index=True)\n",
    "\n",
    "def get_export_df(export_path):\n",
    "    dfs = []\n",
    "    for root, dirs, files in os.walk(export_path):\n",
    "        for dir in dirs:\n",
    "            df = get_channel_df(os.path.join(root, dir))\n",
    "            # add channel column\n",
    "            df[\"channel\"] = dir\n",
    "            dfs.append(df)\n",
    "    return pd.concat(dfs, ignore_index=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c2ce5c7",
   "metadata": {},
   "source": [
    "Be sure to set the path to your slack export before running the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "366353df",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_slack_export = \"slack-export\"\n",
    "\n",
    "get_export_df(path_to_slack_export).to_parquet(\"messages.parquet\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4e83c06",
   "metadata": {},
   "source": [
    "### 1.2 Use Kaskada to break the Slack Export into a set of conversations\n",
    "\n",
    "To do this, we will:\n",
    "* 1.2.1 Start a Kaskada session, and load in the data\n",
    "* 1.2.2 Break the data into *threads* and *non_threads*\n",
    "* 1.2.3 Convert the *non_threads* into *threads*\n",
    "* 1.2.4 Rejoin all the messages into a single timestream\n",
    "* 1.2.5 Split the messages into conversations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fbf4346",
   "metadata": {},
   "source": [
    "#### 1.2.1 Start a Kaskada session, and load in the data\n",
    "\n",
    "The following block should only be run once per session:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61ea2e95-6d9d-4068-ab98-8cf94bc4d9d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import kaskada as kd\n",
    "\n",
    "# Initialize Kaskada with a local execution context.\n",
    "kd.init_session()\n",
    "\n",
    "# set pandas to display all floats with 6 decimal places\n",
    "pd.options.display.float_format = '{:.6f}'.format"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61850395",
   "metadata": {},
   "source": [
    "This block can be re-run as often as you would like to restart the process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b61aa253",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load events from a Parquet file\n",
    "#\n",
    "# if you wan to load in your own slack data, change this to the path of your output file from 1.1 above\n",
    "# otherwise continue with `slack-generation.parquet`, which contains generated slack data for\n",
    "# example purposes. See the `slack-generation/notebook.ipynb` notebook for more info.\n",
    "input_file = \"slack-generation.parquet\"\n",
    "\n",
    "# Use the \"ts\" column as the time associated with each row,\n",
    "# and the \"channel\" column as the entity associated with each row.\n",
    "messages = await kd.sources.Parquet.create(\n",
    "    input_file,\n",
    "    time_column = \"ts\",\n",
    "    key_column = \"channel\",\n",
    "    time_unit = \"s\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf7ba58f",
   "metadata": {},
   "source": [
    "The `messages` object is a Kaskada Timestream.  We can use the `preview()` method to export the first few rows as a pandas dataframe:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa1d0ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the first 5 events\n",
    "messages.preview(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "443449d3",
   "metadata": {},
   "source": [
    "#### 1.2.2 Break the data into *threads* and *non_threads*\n",
    "\n",
    "Before generating examples to fine-tune a LLM model, we need to break up the Slack data into *\"conversations\"*. For our purposes, we define *\"conversations\"* as either:\n",
    "* All the messages in a thread\n",
    "* A group of messages outside a thread, that have at least a 10 minute gap bewteen the next message group.\n",
    "\n",
    "In the Slack Export data, messages in a thread all have the same `thread_ts` value, which matches the `ts` of the first message in the thread. Messages outside a thread (in the root of the channel), do not have a `thread_ts` set. Therefore, we can use this field to filter our data into 2 sets: *threads* and *non-threads*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6735d8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "threads = messages.filter(messages.col(\"thread_ts\").is_not_null())\n",
    "threads.preview(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a26e01a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "non_threads = messages.filter(messages.col(\"thread_ts\").is_null())\n",
    "non_threads.preview(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4ea290c",
   "metadata": {},
   "source": [
    "#### 1.2.3 Convert the *non_threads* into *threads*\n",
    "\n",
    "Using the definition of a *\"conversation\"* for a non-thread from above, we can convert the *non-threads* data into *threads* by:\n",
    "* separating the messages into groups, where there is at least 10 mintues between each group\n",
    "* setting the `thread_ts` field on all messages in the group equal to the `ts` of the first message in the group\n",
    "\n",
    "Note that all data in Kaskada is sorted by time. Therefore you can know that the first message in a group is always the first one that occured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "197d24b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow as pa\n",
    "from datetime import timedelta\n",
    "\n",
    "ts = non_threads.col(\"ts\")\n",
    "ts_since = ts.seconds_since_previous()\n",
    "is_new = ts_since > timedelta(minutes=10)\n",
    "\n",
    "# Eventually this will just be: `thread_ts = ts.first(window=kd.windows.Since(is_new, start=\"inclusive\"))`\n",
    "#\n",
    "# However, the `Since()` window currently collects all the messages until the predicate is True,\n",
    "# then outputs, and starts re-collecting. But we need the opposite: when the predicate is True,\n",
    "# clear the output, then start collecting, and output.\n",
    "#\n",
    "# In other words, `Since()` is currently exclusive on the start of the window, inclusive on the end.\n",
    "# But we need inclusive on the start and exclusive on the end.\n",
    "#\n",
    "# The hack below does what we need until the `Since()` provides additional options for inclusivity\n",
    "shifted_non_threads = non_threads.shift_by(timedelta(microseconds=0.001))\n",
    "shifted_ts = shifted_non_threads.lag(1).col(\"ts\").first(window=kd.windows.Since(is_new))\n",
    "thread_ts = ts.if_(is_new).else_(shifted_ts)\n",
    "\n",
    "# replace the `thread_ts` in the data with our generated value and filter out messages we no longer need\n",
    "non_threads_threads = non_threads.extend({\"thread_ts\": thread_ts}).filter(ts.is_not_null().and_(thread_ts.is_not_null()))\n",
    "non_threads_threads.preview(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41dc1307",
   "metadata": {},
   "source": [
    "#### 1.2.4 Rejoin all the messages into a single timestream\n",
    "\n",
    "Now that both threads and non-threads have the `threads_ts` set, we merge the data back into a single timestream."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "169b1d1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "joined = threads.else_(non_threads_threads)\n",
    "joined.preview(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12a7c538",
   "metadata": {},
   "source": [
    "#### 1.2.5 Split the messages into conversations\n",
    "\n",
    "Now we create a new timestream from the `joined`, using the `with_key()` method. This updates the entity key for the timestream to be based on the combined value of the `channel` and the `thread_ts`. The effect of this is that each entity is now a unique conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad202c51",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = joined.with_key(kd.record({\n",
    "        \"channel\": joined.col(\"channel\"),\n",
    "        \"thread\": joined.col(\"thread_ts\"),\n",
    "    }))\n",
    "\n",
    "messages.preview(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0a40d25",
   "metadata": {},
   "source": [
    "### 1.3 Generate the initial training examples\n",
    "\n",
    "Here we:\n",
    "* 1.3.1 Collect messages into groups on a per-conversation basis\n",
    "* 1.3.2 Create single-token labels for all of the users\n",
    "* 1.3.3 Format, clean, and output the initial examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ad42f3d",
   "metadata": {},
   "source": [
    "#### 1.3.1 Collect messages into groups on a per-conversation basis\n",
    "\n",
    "We collect up lines from the conversation for outputting. On each row, we want the previous set of messages from the conversation, limiting to at most 5 messages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dd4639c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# collect the previous 1 to 5 messages and the associated user for each message\n",
    "conversation = messages.select(\"user\", \"text\").collect(max=5, min=1).lag(1)\n",
    "\n",
    "# add the conversation to the current row\n",
    "examples = messages.extend({\"conversation\":conversation}).filter(conversation.is_not_null())\n",
    "\n",
    "examples.preview(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13c8961e",
   "metadata": {},
   "source": [
    "#### 1.3.2 Create single-token labels for all of the users\n",
    "\n",
    "The following script initializes a Scikit-Learn LabelEncoder, which we can use to ensure that each user is represented by a single \"token\", and that our training examples are formatted in a way that is easier for model fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ce6ba81",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn import preprocessing\n",
    "import json\n",
    "\n",
    "# Encode user ID labels\n",
    "le = preprocessing.LabelEncoder()\n",
    "le.fit(examples.to_pandas()[\"user\"])\n",
    "with open('labels_v2.json', 'w') as f:\n",
    "    json.dump(le.classes_.tolist(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b16166",
   "metadata": {},
   "source": [
    "#### 1.3.3 Format, clean, and output the initial examples\n",
    "\n",
    "The following script interates over the full results, and outputs the examples to a jsonl file: `examples.jsonl`\n",
    "\n",
    "The examples will be used to teach the model the specific users who are interested in a given conversation. Each example consists of a \"prompt\" containing the state of a conversation at a point in time and a \"completion\" containing the user that responded to the previous set of messagees."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4716e77",
   "metadata": {},
   "source": [
    "We will soon be releasing an update with UDF support, so that the following cleanup can be done inside Kaskada"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adafc251",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, re\n",
    "\n",
    "def strip_links_and_users(line):\n",
    "    return re.sub(r\"<.*?>\", '', line)\n",
    "\n",
    "def strip_emoji(line):\n",
    "    return re.sub(r\":.*?:\", '', line)\n",
    "\n",
    "def clean_messages(messages):\n",
    "    cleaned = []\n",
    "    for msg in messages:\n",
    "        text = strip_links_and_users(msg)\n",
    "        text = strip_emoji(text)\n",
    "        text = text.strip()\n",
    "        if text == \"\" or text.find(\"```\") >= 0:\n",
    "            continue\n",
    "        cleaned.append(text)\n",
    "    return cleaned\n",
    "\n",
    "prompt_suffix = \"\\n\\n###\\n\\n\"\n",
    "max_prompt_len = 5000\n",
    "\n",
    "# Format prompt for the OpenAI API\n",
    "def format_prompt(messages):\n",
    "    cleaned = clean_messages(messages)\n",
    "    if len(cleaned) == 0:\n",
    "        return None\n",
    "    cleaned.reverse()\n",
    "    prompt = \"\\n\\n\".join(cleaned)\n",
    "    if len(prompt) > max_prompt_len:\n",
    "        prompt = prompt[0:max_prompt_len]\n",
    "    return prompt+prompt_suffix\n",
    "\n",
    "\n",
    "# use the label mapping to transform the userId to a single token\n",
    "def format_completion(user):\n",
    "    return \" \" + le.transform([user])[0].astype(str) + \" end\"\n",
    "\n",
    "with open('examples.jsonl', 'w') as out_file:\n",
    "    last_prompt = \"\"\n",
    "    for row in examples.run_iter(kind=\"row\"):\n",
    "        user = row[\"user\"]\n",
    "        non_user_messages = []\n",
    "        for msg in row[\"conversation\"]:\n",
    "            if msg[\"user\"] != user:\n",
    "                non_user_messages.append(msg[\"text\"])\n",
    "        prompt = format_prompt(non_user_messages)\n",
    "        if prompt and prompt != last_prompt:\n",
    "            example = { \"prompt\": prompt, \"completion\": format_completion(user) }\n",
    "            out_file.write(json.dumps(example) + \"\\n\")\n",
    "            last_prompt = prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a2aee97",
   "metadata": {},
   "source": [
    "## 2. Use few-shot learning with ChatGPT to clean our training examples\n",
    "\n",
    "Before we send our fine-tuning examples to OpenAI to create a custom model, we should clean up our examples to ensure they provide enough signal for determining the interests of our users. \n",
    "\n",
    "Browsing through the generated examples, you may find that some of the messages only contain \"fluff\". These messages can degrade the accuracy of our final fine-tuned model. Furthermore, we probably don't want to alert users about conversations like these, so these become good negative examples for the training set.\n",
    "\n",
    "The cleanup stage will comprise four steps:\n",
    "* 2.1 Human-in-the-loop strong/weak message gathering\n",
    "* 2.2 Strong message summarization, via ChatGPT few-shot learning\n",
    "* 2.3 Example classification, via ChatGPT few-shot learning\n",
    "* 2.4 Final training file creation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3aa1333",
   "metadata": {},
   "source": [
    "### 2.1 Human-in-the-loop strong/weak message gathering\n",
    "\n",
    "First we need to look through the examples we generated to find some examples of messages that provide a strong signal and some examples have a weak signal. To do this, you can manually look through the `examples.jsonl` file or use the included `human.py` script. Ideally we are looking for 50 examples of each type.\n",
    "\n",
    "* A couple examples message with strong signal: \n",
    "\n",
    "    > I'm familiar with Jupyter's employment of CodeMirror, although I'm unsure about the specific tool used for the readme. Currently, I've been utilizing the Rust syntax highlighter for my code blocks. While it's not flawless, it does a reasonably good job of differentiating between functions and literals. Moreover, it highlights instances of 'let' in a unique color. Still, the idea of having a custom highlighter is quite appealing to me.\n",
    "\n",
    "    > The issues stemming from heavy dependence on non-stable components in Kubernetes led to core special interest groups (sigs) being burdened with problems they didn't want to deal with initially. Once a component becomes stable (stable), it can remain in use consistently. However, adding new functionality to stable components involves a rigorous process. Interestingly, in the past couple of years, a new policy has been put in place—new elements can be added at v1beta1 level but require a commitment demonstrated through version updates; otherwise, they will be automatically deprecated.\n",
    "\n",
    "* A few examples of messages with weak signal:\n",
    "\n",
    "    > some very interesting ideas in here, thx for sharing\n",
    "\n",
    "    > were there any issues with this? i'll start verifying a few things in a bit.\n",
    "\n",
    "    > standup?\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "328747a1",
   "metadata": {},
   "source": [
    "### 2.2 Strong message summarization, via ChatGPT few-shot learning\n",
    "\n",
    "In step 2.3, we are going to use few-shot learning to mark our previously generated examples as having strong or weak signal. To do this, we are going to pass the strong & weak examples determined above on every API request. Based on those examples, we will let ChatGPT decide if a message has a strong or weak signal. \n",
    "\n",
    "The examples passed to ChatGPT in this way are super important. It will make decisions based on what it can learn from these few examples. There is also a maximum set of tokes that can be passed to ChatGPT on each request. Therefore we should make sure these examples best represent the type of data that the model may see.\n",
    "\n",
    "From Step 2.1, we see that the messages with strong signal are often quite long. Lets use ChatGPT with few-shot learning to summarize our \"strong\" examples before moving on to the next step.  With summarized text, we should be able to pass more examples to the API and have better overall results.  \n",
    "\n",
    "To do this, we provide instructions to the ChatCompletion API via a set of messages. Each message object contains `role` and `content` properties. The `role` can be either `system`, `user`, or `assistant`. \n",
    "\n",
    "The first message should always be from the `system` role, and provide general instructions to the model of its function. \n",
    "\n",
    "Following this, message pairs of `user` and `assistant` should be added, where the `user` content is our example input and the `assistant` content is our expected response from ChatGPT. These are the \"few-shot\" learnings that ChatGPT uses to help it determine our desired output.\n",
    "\n",
    "Finally, we append a final `user` message that contains the content we to have summarized by the model.  \n",
    "\n",
    "See https://platform.openai.com/docs/guides/gpt for more info."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "271094e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = open(f'examples_strong.jsonl', 'r')\n",
    "out_file = open(f'examples_strong_summarized.jsonl', 'w')\n",
    "\n",
    "prompt_suffix = \"\\n\\n###\\n\\n\"\n",
    "\n",
    "while True:\n",
    "    line = file.readline()\n",
    "\n",
    "    if not line:\n",
    "        break\n",
    "\n",
    "    data = json.loads(line)\n",
    "\n",
    "    prompt = data[\"prompt\"].removesuffix(prompt_suffix)\n",
    "\n",
    "    msgs = [\n",
    "        {\"role\": \"system\", \"content\": \"You are a helpful assistant that provides concise summaries of messages. The response must be significantly shorter than the input. The output should be written as if you were the original author.\"},\n",
    "        {\"role\": \"user\", \"content\": \"I'm familiar with Jupyter's employment of CodeMirror, although I'm unsure about the specific tool used for the readme. Currently, I've been utilizing the Rust syntax highlighter for my code blocks. While it's not flawless, it does a reasonably good job of differentiating between functions and literals. Moreover, it highlights instances of 'let' in a unique color. Still, the idea of having a custom highlighter is quite appealing to me.\"},\n",
    "        {\"role\": \"assistant\", \"content\": \"Familiar with Jupyter's use of CodeMirror, uncertain about readme tool. Using Rust syntax highlighter for code blocks, highlighting 'let' distinctly. Interested in a custom highlighter for better differentiation.\"},\n",
    "        {\"role\": \"user\", \"content\": \"The issues stemming from heavy dependence on non-stable components in Kubernetes led to core special interest groups (sigs) being burdened with problems they didn't want to deal with initially. Once a component becomes stable (stable), it can remain in use consistently. However, adding new functionality to stable components involves a rigorous process. Interestingly, in the past couple of years, a new policy has been put in place—new elements can be added at v1beta1 level but require a commitment demonstrated through version updates; otherwise, they will be automatically deprecated.\"},\n",
    "        {\"role\": \"assistant\", \"content\": \"Heavy reliance on unstable Kubernetes components burdened core SIGs initially. Stable components ensure consistent use, but adding new features is strict. New policy permits v1beta1 additions with commitment, else auto deprecation.\"},\n",
    "        {\"role\": \"user\", \"content\": prompt},\n",
    "    ]\n",
    "\n",
    "    res = openai.ChatCompletion.create(\n",
    "        model = \"gpt-3.5-turbo\",\n",
    "        messages = msgs\n",
    "    )\n",
    "\n",
    "    prompt = res[\"choices\"][0][\"message\"][\"content\"]\n",
    "\n",
    "    data[\"prompt\"] = prompt+prompt_suffix\n",
    "\n",
    "    out_file.write(json.dumps(data) + \"\\n\")\n",
    "\n",
    "file.close()\n",
    "out_file.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c537dd9",
   "metadata": {},
   "source": [
    "### 2.3 Example classification, via ChatGPT few-shot learning\n",
    "\n",
    "Now that we have summarized *strong* examples, we will use few-shot learning again, to classify all our previously generated examples as having *strong* or *weak* signal. \n",
    "\n",
    "This time we will pull our few-shot examples from files instead of including them in the code directly. A few things to note:\n",
    "* If you get an error about too-many tokens used, reduce the `max_count` of included examples, or go back to step 2.2 to further summarize your *strong* examples.\n",
    "* This will cost a fair amount on OpenAI. A rough estimate is $50 per 10,000 examples.\n",
    "* This can take a long time to run to completion. The ChatCompletion API limits the number of tokens used per minute. Running 10,000 examples can take 8 or more hours.\n",
    "\n",
    "First lets build up our \"few-shots\" message array, so we don't have to regenerate that each time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42d0a310",
   "metadata": {},
   "outputs": [],
   "source": [
    "strong = open(f'examples_strong_summarized.jsonl', 'r')\n",
    "weak = open(f'examples_weak.jsonl', 'r')\n",
    "\n",
    "messages = [{\n",
    "    \"role\": \"system\",\n",
    "    \"content\": \"You are a helpful assistant. Your job is to determine if a prompt will be helpful for fine-tuning a model. All prompts start with 'start -->' and end with: '\\\\n\\\\n###\\\\n\\\\n'. You should respond 'yes' if you think the prompt has enough context to be helpful, or 'no' if not. No explanation is needed. You should only respond with 'yes' or 'no'.\"\n",
    "}]\n",
    "\n",
    "count = 0\n",
    "max_count = 50\n",
    "while True:\n",
    "\n",
    "    strong_line = strong.readline()\n",
    "    weak_line = weak.readline()\n",
    "    count += 1\n",
    "\n",
    "    if (not strong_line) or (not weak_line) or (count > max_count):\n",
    "        break\n",
    "\n",
    "    strong_data = json.loads(strong_line)\n",
    "    weak_data = json.loads(weak_line)\n",
    "\n",
    "    messages.append({\"role\": \"user\", \"content\": f'start -->{strong_data[\"prompt\"]}'})\n",
    "    messages.append({\"role\": \"assistant\", \"content\": \"yes\"})\n",
    "    messages.append({\"role\": \"user\", \"content\": f'start -->{weak_data[\"prompt\"]}'})\n",
    "    messages.append({\"role\": \"assistant\",\"content\": \"no\"})\n",
    "\n",
    "strong.close()\n",
    "weak.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1059da2",
   "metadata": {},
   "source": [
    "Then we use those example messages to predict if a prompt will be helpful for training or not. \n",
    "\n",
    "Note that we use the `backoff` library to retry requests that have failed due to a rate-limit error. Even so, sometimes the process stalls and must be manually restarted. The code below appends to the output file instead of replacing it, so that the process can be restarted after an error occurs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6b25924",
   "metadata": {},
   "outputs": [],
   "source": [
    "# only re-run this cell if you want to completely start over after an error occurs\n",
    "starting_count = 0\n",
    "\n",
    "total_count = 0\n",
    "with open(f'examples.jsonl', 'r') as file:\n",
    "    for line in file:\n",
    "        total_count += 1\n",
    "\n",
    "print(f'There are {total_count} lines in the input file.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93acbe47",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, openai, time, logging, backoff\n",
    "\n",
    "# for debugging responses from the API, un-comment this\n",
    "# logging.getLogger('backoff').addHandler(logging.StreamHandler())\n",
    "\n",
    "file = open(f'examples.jsonl', 'r')\n",
    "strong_file = open(f'examples_cleaned_strong.jsonl', 'a')\n",
    "weak_file = open(f'examples_cleaned_weak.jsonl', 'a')\n",
    "\n",
    "@backoff.on_exception(backoff.expo, (openai.error.RateLimitError, openai.error.ServiceUnavailableError))\n",
    "def chat_with_backoff(**kwargs):\n",
    "    time.sleep(1)\n",
    "    try:\n",
    "        return openai.ChatCompletion.create(**kwargs)\n",
    "    except openai.error.InvalidRequestError:\n",
    "        return None\n",
    "\n",
    "count = 0\n",
    "for line in file:\n",
    "    count +=1\n",
    "\n",
    "    # helpful for restarting after issue\n",
    "    if count < starting_count:\n",
    "        continue\n",
    "\n",
    "    data = json.loads(line)\n",
    "\n",
    "    prompt = data[\"prompt\"]\n",
    "\n",
    "    msgs = messages.copy()\n",
    "\n",
    "    msgs.append({\"role\": \"user\", \"content\": f'start -->{prompt}'})\n",
    "\n",
    "    res = chat_with_backoff(\n",
    "        model = \"gpt-3.5-turbo\",\n",
    "        messages = msgs\n",
    "    )\n",
    "    if not res:\n",
    "        continue\n",
    "    response = res[\"choices\"][0][\"message\"][\"content\"]\n",
    "\n",
    "    # for debugging responses from the API, un-comment this\n",
    "    # print(f'Result was `{response}` for prompt: {prompt}')\n",
    "\n",
    "    print(f'Currently processing line {count} of {total_count}')\n",
    "\n",
    "    if response == \"yes\":\n",
    "        strong_file.write(line)\n",
    "        strong_file.flush()\n",
    "    else:\n",
    "        # for weak messages, re-write the completion as ` nil`\n",
    "        data[\"completion\"] = \" nil end\"\n",
    "        weak_file.write(json.dumps(data) + '\\n')\n",
    "        weak_file.flush()\n",
    "\n",
    "    starting_count = count\n",
    "\n",
    "file.close()\n",
    "strong_file.close()\n",
    "weak_file.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3636649d",
   "metadata": {},
   "source": [
    "### 2.4 Final training file creation\n",
    "\n",
    "To create the final training file, we want to ensure that we have an equal number of *strong* and *weak* examples. We will use pandas to grab random examples from each set, combine them in a random order, and the save the final example set to disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d76684cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "strong_df = pd.read_json('examples_cleaned_strong.jsonl', lines=True, orient='records')\n",
    "weak_df = pd.read_json('examples_cleaned_weak.jsonl', lines=True, orient='records')\n",
    "\n",
    "min_length = min([len(strong_df.index), len(weak_df.index)])\n",
    "\n",
    "strong_df = strong_df.sample(min_length)\n",
    "weak_df = weak_df.sample(min_length)\n",
    "\n",
    "combined = pd.concat([strong_df, weak_df])\n",
    "combined = combined.sample(frac=1)\n",
    "combined.to_json('final_examples.jsonl', lines=True, orient='records')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77671e30",
   "metadata": {},
   "source": [
    "Before sending our examples for fine-tuning, we use a tool provided by OpenAI to perform some verification on our input data and then split the dataset into 2 files.  The tool does the following for us:\n",
    "\n",
    "* makes sure all prompts end with same suffix\n",
    "* removes examples that use too many tokens\n",
    "* removes duplicated examples\n",
    "\n",
    "Note: We aren't doing classification, so don't start a fine-tune as suggested by the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5387d01",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "from openai import cli\n",
    "from types import SimpleNamespace\n",
    "\n",
    "args = SimpleNamespace(file='final_examples.jsonl', quiet=True)\n",
    "cli.FineTune.prepare_data(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b171059a",
   "metadata": {},
   "source": [
    "## 3. Send our training data to OpenAI and create a fine-tuned model\n",
    "\n",
    "Finally, we'll send our fine-tuning examples to OpenAI to create a custom model.\n",
    "\n",
    "To do this, we will do the following:\n",
    "* 3.1 Upload training data\n",
    "* 3.2 Create a fine-tuning job\n",
    "* 3.3 Wait for the fine-tuning to start\n",
    "* 3.4 Wait for the fine-tuning to finish\n",
    "* 3.5 Try using the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5802843",
   "metadata": {},
   "source": [
    "### 3.1 Upload training data\n",
    "\n",
    "First we upload the verified final examples file from above to OpenAI. We need to make sure the file has successfully uploaded before moving onto the next step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52c1b41",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "training_file_name = 'final_examples_prepared_train.jsonl'\n",
    "\n",
    "# start the file upload\n",
    "training_file_id = cli.FineTune._get_or_upload(training_file_name, True)\n",
    "\n",
    "# Poll and display the upload status until the file finishes\n",
    "while True:\n",
    "    time.sleep(2)\n",
    "    file_status = openai.File.retrieve(training_file_id)[\"status\"]\n",
    "    print(f'Upload status: {file_status}')\n",
    "    if file_status in [\"succeeded\", \"failed\", \"processed\"]:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa3d6427",
   "metadata": {},
   "source": [
    "### 3.2 Create a fine-tuning job\n",
    "\n",
    "We recommened using either the `curie` or `davinci` models for fine-tuning. We had good success with both of them. Note that the `curie` model is cheaper to use, but takes longer to train a helpful model.\n",
    "\n",
    "There are many parameters to set when training a model. Use our recommendations or try your own. With `curie` used 8 epochs, and with `davinci` we used 4. More help can be found here: https://platform.openai.com/docs/api-reference/completions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83914ada-d108-422b-b4c0-7a0d9576d031",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "create_args = {\n",
    "    \"training_file\": training_file_id,\n",
    "    \"model\": \"curie\",\n",
    "    \"n_epochs\": 8,\n",
    "    \"learning_rate_multiplier\": 0.02,\n",
    "    \"suffix\": \"beep-gpt\"\n",
    "}\n",
    "\n",
    "# Create the fine-tune job and retrieve the job ID\n",
    "resp = openai.FineTune.create(**create_args)\n",
    "job_id = resp[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01215845",
   "metadata": {},
   "source": [
    "### 3.3 Wait for the fine-tuning to start\n",
    "\n",
    "Note, it can take several hours for the fine-tuning to start. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fee6749",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Poll and display the fine-tuning status until the it starts\n",
    "while True:\n",
    "    time.sleep(5)\n",
    "    job_status = openai.FineTune.retrieve(id=job_id)[\"status\"]\n",
    "    print(f'Job status: {job_status}')\n",
    "    if job_status in [\"failed\", \"started\", \"succeeded\"]:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b79218a8",
   "metadata": {},
   "source": [
    "### 3.4 Wait for the fine-tuning to finish\n",
    "\n",
    "Note, it can take a long time for the fine-tuning to finish. A training set of 2000 examples took about 2 hours to train on `curie` with 8 epochs.\n",
    "\n",
    "Run the following code block periodically, until you see a `failed` or `succeeded` status."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ddbc748",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_details = openai.FineTune.retrieve(job_id)\n",
    "\n",
    "print(f'Job status: {job_details[\"status\"]}')\n",
    "print(f'Job events: {job_details[\"events\"]}')\n",
    "\n",
    "if job_details[\"status\"] == \"succeeded\":\n",
    "    model_id = job_details[\"fine_tuned_model\"]\n",
    "    print(f'Successfully fine-tuned model with ID: {model_id}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca17c223",
   "metadata": {},
   "source": [
    "### 3.5 Try using the model\n",
    "\n",
    "Using the validation file, we can try sending a few prompts to our new model and see if it recommends alerting any users."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e306db5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose which row in the validation file to send\n",
    "row = 6\n",
    "\n",
    "valid_df = pd.read_json('final_examples_prepared_valid.jsonl', lines=True, orient='records')\n",
    "\n",
    "prompt = valid_df['prompt'][i]\n",
    "completion = valid_df['completion'][i]\n",
    "\n",
    "# this is the text we send to the model for it to determine if we should alert a user\n",
    "print(f'Prompt: {prompt}')\n",
    "\n",
    "# this is the user (or nil) we would have expected for the response\n",
    "print(f'Completion: {completion}')\n",
    "\n",
    "# this is the response from the model. The `text` feild contains the actual prediction. The `logprobs` arrary contains the log-probability from the 5 highest potential matches.\n",
    "print(f'Prediction:')\n",
    "\n",
    "openai.Completion.create(model=model_id, prompt=prompt, max_tokens=1, n=1, logprobs=5, stop=\" end\", temperature=0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}