openai-cookbook/examples/Chat_finetuning_data_prep.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a06ec76c",
   "metadata": {},
   "source": [
    "# Data preparation and analysis for chat model fine-tuning\n",
    "\n",
    "This notebook serves as a tool to preprocess and analyze the chat dataset used for fine-tuning a chat model. \n",
    "It checks for format errors, provides basic statistics, and estimates token counts for fine-tuning costs.\n",
    "The method shown here corresponds to the [current fine-tuning method](https://platform.openai.com/docs/guides/fine-tuning) for gpt-3.5-turbo.\n",
    "See [legacy fine-tuning](https://platform.openai.com/docs/guides/legacy-fine-tuning) for models like babbage-002 and davinci-002."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4e63973b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import tiktoken # for token counting\n",
    "import numpy as np\n",
    "from collections import defaultdict"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "013bdbc4",
   "metadata": {},
   "source": [
    "## Data loading\n",
    "\n",
    "We first load the chat dataset from an [example JSONL file](https://github.com/openai/openai-cookbook/blob/main/examples/data/toy_chat_fine_tuning.jsonl)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c248ccd1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num examples: 5\n",
      "First example:\n",
      "{'role': 'system', 'content': 'You are a happy assistant that puts a positive spin on everything.'}\n",
      "{'role': 'user', 'content': 'I fell off my bike today.'}\n",
      "{'role': 'assistant', 'content': \"It's great that you're getting exercise outdoors!\"}\n"
     ]
    }
   ],
   "source": [
    "data_path = \"data/toy_chat_fine_tuning.jsonl\"\n",
    "\n",
    "# Load the dataset\n",
    "with open(data_path, 'r', encoding='utf-8') as f:\n",
    "    dataset = [json.loads(line) for line in f]\n",
    "\n",
    "# Initial dataset stats\n",
    "print(\"Num examples:\", len(dataset))\n",
    "print(\"First example:\")\n",
    "for message in dataset[0][\"messages\"]:\n",
    "    print(message)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "17903d61",
   "metadata": {},
   "source": [
    "## Format validation\n",
    "\n",
    "We can perform a variety of error checks to validate that each conversation in the dataset adheres to the format expected by the fine-tuning API. Errors are categorized based on their nature for easier debugging.\n",
    "\n",
    "1. **Data Type Check**: Checks whether each entry in the dataset is a dictionary (`dict`). Error type: `data_type`.\n",
    "2. **Presence of Message List**: Checks if a `messages` list is present in each entry. Error type: `missing_messages_list`.\n",
    "3. **Message Keys Check**: Validates that each message in the `messages` list contains the keys `role` and `content`. Error type: `message_missing_key`.\n",
    "4. **Unrecognized Keys in Messages**: Logs if a message has keys other than `role`, `content`, and `name`. Error type: `message_unrecognized_key`.\n",
    "5. **Role Validation**: Ensures the `role` is one of \"system\", \"user\", or \"assistant\". Error type: `unrecognized_role`.\n",
    "6. **Content Validation**: Verifies that `content` has textual data and is a string. Error type: `missing_content`.\n",
    "7. **Assistant Message Presence**: Checks that each conversation has at least one message from the assistant. Error type: `example_missing_assistant_message`.\n",
    "\n",
    "The code below performs these checks, and outputs counts for each type of error found are printed. This is useful for debugging and ensuring the dataset is ready for the next steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d9f3ccbf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No errors found\n"
     ]
    }
   ],
   "source": [
    "# Format error checks\n",
    "format_errors = defaultdict(int)\n",
    "\n",
    "for ex in dataset:\n",
    "    if not isinstance(ex, dict):\n",
    "        format_errors[\"data_type\"] += 1\n",
    "        continue\n",
    "        \n",
    "    messages = ex.get(\"messages\", None)\n",
    "    if not messages:\n",
    "        format_errors[\"missing_messages_list\"] += 1\n",
    "        continue\n",
    "        \n",
    "    for message in messages:\n",
    "        if \"role\" not in message or \"content\" not in message:\n",
    "            format_errors[\"message_missing_key\"] += 1\n",
    "        \n",
    "        if any(k not in (\"role\", \"content\", \"name\", \"function_call\") for k in message):\n",
    "            format_errors[\"message_unrecognized_key\"] += 1\n",
    "        \n",
    "        if message.get(\"role\", None) not in (\"system\", \"user\", \"assistant\", \"function\"):\n",
    "            format_errors[\"unrecognized_role\"] += 1\n",
    "            \n",
    "        content = message.get(\"content\", None)\n",
    "        function_call = message.get(\"function_call\", None)\n",
    "        \n",
    "        if (not content and not function_call) or not isinstance(content, str):\n",
    "            format_errors[\"missing_content\"] += 1\n",
    "    \n",
    "    if not any(message.get(\"role\", None) == \"assistant\" for message in messages):\n",
    "        format_errors[\"example_missing_assistant_message\"] += 1\n",
    "\n",
    "if format_errors:\n",
    "    print(\"Found errors:\")\n",
    "    for k, v in format_errors.items():\n",
    "        print(f\"{k}: {v}\")\n",
    "else:\n",
    "    print(\"No errors found\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "981e77da",
   "metadata": {},
   "source": [
    "## Token Counting Utilities\n",
    "\n",
    "Lets define a few helpful utilities to be used in the rest of the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8f4b47b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoding = tiktoken.get_encoding(\"cl100k_base\")\n",
    "\n",
    "# not exact!\n",
    "# simplified from https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb\n",
    "def num_tokens_from_messages(messages, tokens_per_message=3, tokens_per_name=1):\n",
    "    num_tokens = 0\n",
    "    for message in messages:\n",
    "        num_tokens += tokens_per_message\n",
    "        for key, value in message.items():\n",
    "            num_tokens += len(encoding.encode(value))\n",
    "            if key == \"name\":\n",
    "                num_tokens += tokens_per_name\n",
    "    num_tokens += 3\n",
    "    return num_tokens\n",
    "\n",
    "def num_assistant_tokens_from_messages(messages):\n",
    "    num_tokens = 0\n",
    "    for message in messages:\n",
    "        if message[\"role\"] == \"assistant\":\n",
    "            num_tokens += len(encoding.encode(message[\"content\"]))\n",
    "    return num_tokens\n",
    "\n",
    "def print_distribution(values, name):\n",
    "    print(f\"\\n#### Distribution of {name}:\")\n",
    "    print(f\"min / max: {min(values)}, {max(values)}\")\n",
    "    print(f\"mean / median: {np.mean(values)}, {np.median(values)}\")\n",
    "    print(f\"p5 / p95: {np.quantile(values, 0.1)}, {np.quantile(values, 0.9)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0fdff67d",
   "metadata": {},
   "source": [
    "## Data Warnings and Token Counts \n",
    "\n",
    "With some lightweight analysis we can identify potential issues in the dataset, like missing messages, and provide statistical insights into message and token counts.\n",
    "\n",
    "1. **Missing System/User Messages**: Counts the number of conversations missing a \"system\" or \"user\" message. Such messages are critical for defining the assistant's behavior and initiating the conversation.\n",
    "2. **Number of Messages Per Example**: Summarizes the distribution of the number of messages in each conversation, providing insight into dialogue complexity.\n",
    "3. **Total Tokens Per Example**: Calculates and summarizes the distribution of the total number of tokens in each conversation. Important for understanding fine-tuning costs.\n",
    "4. **Tokens in Assistant's Messages**: Calculates the number of tokens in the assistant's messages per conversation and summarizes this distribution. Useful for understanding the assistant's verbosity.\n",
    "5. **Token Limit Warnings**: Checks if any examples exceed the maximum token limit (4096 tokens), as such examples will be truncated during fine-tuning, potentially resulting in data loss.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "52e58ee4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num examples missing system message: 1\n",
      "Num examples missing user message: 1\n",
      "\n",
      "#### Distribution of num_messages_per_example:\n",
      "min / max: 2, 9\n",
      "mean / median: 3.8, 3.0\n",
      "p5 / p95: 2.0, 6.6000000000000005\n",
      "\n",
      "#### Distribution of num_total_tokens_per_example:\n",
      "min / max: 26, 8032\n",
      "mean / median: 1648.4, 45.0\n",
      "p5 / p95: 26.8, 4863.6\n",
      "\n",
      "#### Distribution of num_assistant_tokens_per_example:\n",
      "min / max: 4, 8000\n",
      "mean / median: 1610.2, 10.0\n",
      "p5 / p95: 6.0, 4811.200000000001\n",
      "\n",
      "1 examples may be over the 4096 token limit, they will be truncated during fine-tuning\n"
     ]
    }
   ],
   "source": [
    "# Warnings and tokens counts\n",
    "n_missing_system = 0\n",
    "n_missing_user = 0\n",
    "n_messages = []\n",
    "convo_lens = []\n",
    "assistant_message_lens = []\n",
    "\n",
    "for ex in dataset:\n",
    "    messages = ex[\"messages\"]\n",
    "    if not any(message[\"role\"] == \"system\" for message in messages):\n",
    "        n_missing_system += 1\n",
    "    if not any(message[\"role\"] == \"user\" for message in messages):\n",
    "        n_missing_user += 1\n",
    "    n_messages.append(len(messages))\n",
    "    convo_lens.append(num_tokens_from_messages(messages))\n",
    "    assistant_message_lens.append(num_assistant_tokens_from_messages(messages))\n",
    "    \n",
    "print(\"Num examples missing system message:\", n_missing_system)\n",
    "print(\"Num examples missing user message:\", n_missing_user)\n",
    "print_distribution(n_messages, \"num_messages_per_example\")\n",
    "print_distribution(convo_lens, \"num_total_tokens_per_example\")\n",
    "print_distribution(assistant_message_lens, \"num_assistant_tokens_per_example\")\n",
    "n_too_long = sum(l > 4096 for l in convo_lens)\n",
    "print(f\"\\n{n_too_long} examples may be over the 4096 token limit, they will be truncated during fine-tuning\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2afb04df",
   "metadata": {},
   "source": [
    "## Cost Estimation\n",
    "\n",
    "In this final section, we estimate the total number of tokens that will be used for fine-tuning, which allows us to approximate the cost. It is worth noting that the duration of the fine-tuning jobs will also increase with the token count. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fb95a7ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset has ~4306 tokens that will be charged for during training\n",
      "By default, you'll train for 20 epochs on this dataset\n",
      "By default, you'll be charged for ~86120 tokens\n"
     ]
    }
   ],
   "source": [
    "# Pricing and default n_epochs estimate\n",
    "MAX_TOKENS_PER_EXAMPLE = 4096\n",
    "\n",
    "TARGET_EPOCHS = 3\n",
    "MIN_TARGET_EXAMPLES = 100\n",
    "MAX_TARGET_EXAMPLES = 25000\n",
    "MIN_DEFAULT_EPOCHS = 1\n",
    "MAX_DEFAULT_EPOCHS = 25\n",
    "\n",
    "n_epochs = TARGET_EPOCHS\n",
    "n_train_examples = len(dataset)\n",
    "if n_train_examples * TARGET_EPOCHS < MIN_TARGET_EXAMPLES:\n",
    "    n_epochs = min(MAX_DEFAULT_EPOCHS, MIN_TARGET_EXAMPLES // n_train_examples)\n",
    "elif n_train_examples * TARGET_EPOCHS > MAX_TARGET_EXAMPLES:\n",
    "    n_epochs = max(MIN_DEFAULT_EPOCHS, MAX_TARGET_EXAMPLES // n_train_examples)\n",
    "\n",
    "n_billing_tokens_in_dataset = sum(min(MAX_TOKENS_PER_EXAMPLE, length) for length in convo_lens)\n",
    "print(f\"Dataset has ~{n_billing_tokens_in_dataset} tokens that will be charged for during training\")\n",
    "print(f\"By default, you'll train for {n_epochs} epochs on this dataset\")\n",
    "print(f\"By default, you'll be charged for ~{n_epochs * n_billing_tokens_in_dataset} tokens\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a0ad0369",
   "metadata": {},
   "source": [
    "See https://openai.com/pricing to estimate total costs."
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}