classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8c73ef23",
   "metadata": {},
   "source": [
    "# Automatic Prompt Engineering for classification\n",
    "\n",
    "Given only (text -> label), this notebook generates and optimizes system and user prompts.\n",
    "\n",
    "This is how classification is intended to be done.\n",
    "- (system prompt, user prompt prefix + text + user prompt suffix) -Haiku-> bot response -extract-> label\n",
    "\n",
    "The notebook will produce\n",
    "- the system prompt\n",
    "- the user prompt prefix\n",
    "- the user prompt suffix\n",
    "\n",
    "You can simply run this notebook with just\n",
    "- an Anthropic API key and an OpenAI API key\n",
    "\n",
    "If you want to change the classification task, you will need to\n",
    "- provide a dataset (text -> label)\n",
    "\n",
    "This is how prompt tuning is done\n",
    "- Sample from the full dataset.\n",
    "- Haiku takes in (system prompt, user prompt prefix + text + user prompt suffix) and produces model_response\n",
    "- Extract the label from the model_response.\n",
    "- Sample from the mistakes and the correct results.\n",
    "- o1-mini summarizes the mistakes and update the prompts (model parameters).\n",
    "- Repeat.\n",
    "\n",
    "You will need to have these Python modules installed\n",
    "- pandas\n",
    "- scikit-learn\n",
    "- anthropic\n",
    "- openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3c074e16",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/Caskroom/miniforge/base/lib/python3.9/site-packages/scipy/__init__.py:155: UserWarning: A NumPy version >=1.18.5 and <1.25.0 is required for this version of SciPy (detected version 1.26.4\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import re\n",
    "import random\n",
    "import inspect\n",
    "import textwrap\n",
    "import collections\n",
    "import itertools\n",
    "import concurrent.futures\n",
    "import pandas as pd\n",
    "import html\n",
    "from collections import defaultdict\n",
    "\n",
    "from IPython.display import display, HTML\n",
    "from sklearn.metrics import precision_score, recall_score\n",
    "\n",
    "import anthropic\n",
    "from openai import OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4c8e450",
   "metadata": {},
   "source": [
    "# Use your Anthropic API key and OpenAI API key here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c127a818",
   "metadata": {},
   "outputs": [],
   "source": [
    "anthropic_api_key = os.environ.get(\"ANTHROPIC_API_KEY\")\n",
    "anthropic_client = anthropic.Anthropic(api_key=anthropic_api_key)\n",
    "\n",
    "openai_api_key = os.environ.get(\"OPENAI_API_KEY\")\n",
    "openai_client = OpenAI(api_key=openai_api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e7c9b939",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sk-ant-api03-Mbs\n"
     ]
    }
   ],
   "source": [
    "print(anthropic_client.api_key[:16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "39725119",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sk-T0uvUYsHoQOag\n"
     ]
    }
   ],
   "source": [
    "print(openai_client.api_key[:16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6130e282",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_PARALLEL_FORWARD_PASS_API_CALLS = 100  # see https://docs.anthropic.com/claude/reference/rate-limits\n",
    "NUM_SAMPLES_FORWARD_PASS_FOR_EACH_LABEL = 100\n",
    "NUM_SAMPLES_MISTAKE_GRADIENT_CALCULATION_FOR_EACH_LABEL = 20\n",
    "NUM_SAMPLES_CORRECT_GRADIENT_CALCULATION_FOR_EACH_LABEL = 10\n",
    "NUM_ITERATIONS = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a33cf13",
   "metadata": {},
   "source": [
    "# Define the dataset here\n",
    "\n",
    "The outcome of this section should be the variable `dataset: list[tuple[str, str]]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "11a5d8c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "tsv = \"\"\"\n",
    "What are some examples of sorting algorithms that require more conditional statements?\tTRUE\n",
    "Can you elaborate on how radix sort and counting sort functions work?\tTRUE\n",
    "What are some specialized sorting algorithms?\tTRUE\n",
    "Could you explain why sorting is equivalent to discovering the permutation of items?\tTRUE\n",
    "What are some real-world applications of sorting algorithms?\tTRUE\n",
    "Do you think some billionaires intentionally underreport their wealth?\tTRUE\n",
    "Could the media's focus on billionaires be detrimental in any way?\tTRUE\n",
    "Do you think that billionaire wealth lists provide any meaningful insights into the global economy?\tTRUE\n",
    "What are some activities that Japanese students could do to be more proactive in their learning?\tTRUE\n",
    "Is the trend of students becoming more \"obedient and good kids\" a positive or negative development for Japan?\tFALSE\n",
    "What are the implications of Japanese students becoming less likely to read for pleasure?\tTRUE\n",
    "What are some ways that teachers can help students transition from being students to graduates?\tTRUE\n",
    "How might this \"good kid\" syndrome affect Japanese society in the long term?\tFALSE\n",
    "Do you think DeSantis would have been a good president?\tTRUE\n",
    "Do you think the sympathy wave would have been enough to secure a win for DeSantis?\tFALSE\n",
    "What are some other factors that you think would have impacted the election outcome?\tFALSE\n",
    "Why do you think it's useful to think of dying as a process rather than an event?\tTRUE\n",
    "What other human cell types have a low metabolic rate?\tFALSE\n",
    "How long does it take for muscle and skin cells to die after circulation stops?\tTRUE\n",
    "What are some other factors that forensic scientists use to estimate time of death?\tFALSE\n",
    "What is the most surprising thing about how the human body decomposes?\tTRUE\n",
    "What are some examples of how apathy can affect a project?\tTRUE\n",
    "What is the trade-off between shipping updates quickly and releasing quality software?\tTRUE\n",
    "Is this apathy issue specific to software development, or is it a broader problem?\tFALSE\n",
    "What are the average costs associated with burying someone in a standard casket compared to an extra-large casket?\tTRUE\n",
    "Do cultural or religious beliefs influence the choice of casket size or burial practices?\tTRUE\n",
    "Are extra charges associated with cremation of obese individuals justified?\tTRUE\n",
    "Do you think the price of clothing for larger people is unfairly inflated?\tTRUE\n",
    "How does this perspective hold up when looking at specific Fighting-type Pokémon that are not typically heroic?\tFALSE\n",
    "Does the Fairy type fill the role of a \"Light\" type?\tFALSE\n",
    "What are some other positive qualities besides fighting and righteousness that the Fighting type embodies?\tFALSE\n",
    "Why do you think it's important to condemn violence in politics?\tTRUE\n",
    "What do you mean by \"Trumpism would turn from a political movement to a religion\"?\tFALSE\n",
    "Why do you hope the perpetrator was not a Democrat?\tFALSE\n",
    "How does obtaining guns easily make it easier to take a shot at a president?\tTRUE\n",
    "What are the most important factors when judging the success of a prime minister?\tTRUE\n",
    "What was it like living through the time of Harold Wilson?\tTRUE\n",
    "Did your dad ever explain why he acted that way?\tFALSE\n",
    "Do you think this experience changed your relationship with your father?\tFALSE\n",
    "How do you feel about your father now?\tFALSE\n",
    "Is it common for fathers to act this way?\tFALSE\n",
    "What's the most thoughtful thing your dad ever did for you?\tTRUE\n",
    "What are some benefits you received during your deployment?\tTRUE\n",
    "What were some of the challenges of living on a remote combat outpost in Afghanistan?\tTRUE\n",
    "How did you manage your money during your deployment?\tTRUE\n",
    "Do other cultures use stars and constellations to represent concepts like the afterlife?\tFALSE\n",
    "What are some names for Canis Minor in Sanskrit?\tTRUE\n",
    "Are there other examples of using celestial figures as a way to navigate or remember things?\tFALSE\n",
    "Is there any evidence that ancient Hindus actually used these constellations for navigating the path of the departed souls?\tFALSE\n",
    "Does the concept of Pitriloka and the path of the departed souls play a significant role in modern Hinduism?\tTRUE\n",
    "What are the implications of the current situation between the US and China?\tFALSE\n",
    "Why does China not need nuclear-fueled aircraft carriers?\tTRUE\n",
    "What is the strategy behind China's approach of building smaller ships?\tTRUE\n",
    "Did you think the employees at the first Lowe's thought you were going to spend a lot of money and just didn’t want to deal with it?\tFALSE\n",
    "Did you use financing for all $12,000?\tFALSE\n",
    "Why would someone choose to invest in real estate?\tTRUE\n",
    "Should your colleague factor in the interest on the loan into his calculations?\tFALSE\n",
    "What are some downsides to investing in real estate?\tTRUE\n",
    "What homework should people do before investing in real estate?\tTRUE\n",
    "Is your son's reaction common in boys his age?\tFALSE\n",
    "What could you do to help your son feel more comfortable expressing his emotions?\tTRUE\n",
    "What are the societal pressures that contribute to men being discouraged from expressing their emotions?\tTRUE\n",
    "What are your thoughts on the idea that men should \"toughen up\" and not cry?\tTRUE\n",
    "What does this story say about the nature of desire?\tFALSE\n",
    "What are some other stories of divine play (leela) in Hinduism?\tFALSE\n",
    "What are some other examples of the \"divine play\" in Hindu mythology?\tFALSE\n",
    "Are there other similar stories in other religions or mythologies that involve deities taking on human forms?\tFALSE\n",
    "Do you think Omaha Steaks could be considered a luxury brand?\tTRUE\n",
    "What are your thoughts on companies that incentivize large purchases with discounts and bundling?\tTRUE\n",
    "Why don't Omaha Steaks show photos of their raw steaks?\tTRUE\n",
    "What other information would you need to feel comfortable purchasing from Omaha Steaks?\tFALSE\n",
    "What is the difference between welfare and food stamps?\tTRUE\n",
    "What are some ways companies can create a more flexible work environment?\tTRUE\n",
    "What are some ways employees can communicate effectively with their teams about their schedule?\tTRUE\n",
    "How has the importance of being present in the office changed in recent years?\tTRUE\n",
    "Are there situations where being physically present in the office is still important?\tTRUE\n",
    "What are some companies that have a culture of flexible work hours?\tTRUE\n",
    "What are some examples of children who exhibit similar behavior?\tFALSE\n",
    "Is there any connection between E.'s behavior and his early childhood experiences?\tFALSE\n",
    "What are the long-term effects of \"Defiance and Anger Management Disorder\"?\tFALSE\n",
    "What are some of the challenges of dealing with children with \"Defiance and Anger Management Disorder\"?\tFALSE\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7e1a3f31",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = [row.split('\\t') for row in tsv.split(\"\\n\")]\n",
    "dataset = [(text, \"0\" if label == \"TRUE\" else \"1\") for text, label in dataset]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0c7c1744",
   "metadata": {},
   "outputs": [],
   "source": [
    "for text, target in dataset:\n",
    "    # note that you need to cast the target into a string\n",
    "    assert type(text) == str\n",
    "    assert type(target) == str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "53e4e290",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Counter({'0': 48, '1': 33})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# make sure the number of types of labels is small\n",
    "collections.Counter(label for _, label in dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b9380f08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('What are some examples of sorting algorithms that require more conditional statements?',\n",
       " '0')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dc4cd7a",
   "metadata": {},
   "source": [
    "# Model definition\n",
    "If you want to use a different model, you can configure here this section.\n",
    "\n",
    "Otherwise you can just run the rest of the code and see what prompts is o1 optimizing for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "89fdc763",
   "metadata": {},
   "outputs": [],
   "source": [
    "# usually o1 is good enough to produce working prompts from nothing\n",
    "model_parameters = {\n",
    "    \"system_prompt\": f\"Classify into one of the following labels {set(label for _, label in dataset)}\",\n",
    "    \"user_prompt_prefix\": \"Classify the following text <text>\",\n",
    "    \"user_prompt_suffix\": \"</text>\\nYour response should produce a brief reasoning and end with Label: <label>{{label}}</label>\",\n",
    "}\n",
    "\n",
    "token_counts = defaultdict(int)\n",
    "token_costs = defaultdict(float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c037ad70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the code for this function is being used as o1 input\n",
    "def compute_model_response(text, model_parameters):    \n",
    "    user_message = model_parameters[\"user_prompt_prefix\"] + text + model_parameters[\"user_prompt_suffix\"]\n",
    "    \n",
    "    message = anthropic_client.messages.create(\n",
    "        model=\"claude-3-haiku-20240307\",\n",
    "        max_tokens=2000,\n",
    "        temperature=0,\n",
    "        messages=[\n",
    "            {\"role\": \"user\", \"content\": user_message},\n",
    "            {\"role\": \"assistant\", \"content\": \"Reasoning:\"},\n",
    "        ],\n",
    "        system=model_parameters[\"system_prompt\"],\n",
    "        timeout=10\n",
    "    )\n",
    "    token_counts[\"haiku_input\"] += message.usage.input_tokens\n",
    "    token_counts[\"haiku_output\"] += message.usage.output_tokens\n",
    "    token_costs[\"haiku_input\"] += message.usage.input_tokens * 0.25 * 1e-6\n",
    "    token_costs[\"haiku_output\"] += message.usage.output_tokens * 1.25 * 1e-6\n",
    "\n",
    "    return message.content[0].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fdccd181",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "def extract_from_model_response(text):\n",
    "    pattern = r'<label>(.*?)</label>'\n",
    "    match = re.search(pattern, text, re.DOTALL)\n",
    "    if match:\n",
    "        return match.group(1).strip()\n",
    "    else:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d1b3a15",
   "metadata": {},
   "source": [
    "# Optimization functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "94a927c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_texts_and_labels(dataset):\n",
    "    dataset = [data for data in dataset]\n",
    "    random.shuffle(dataset)\n",
    "    label_set = set(label for _,label in dataset)\n",
    "\n",
    "    sampled_dataset = []\n",
    "    for target_label in label_set:\n",
    "        dataset_with_label = [(data, label) for data, label in dataset if label == target_label]\n",
    "        sampled_dataset += dataset_with_label[:NUM_SAMPLES_FORWARD_PASS_FOR_EACH_LABEL]\n",
    "    random.shuffle(sampled_dataset)\n",
    "\n",
    "    return [data for data, _ in sampled_dataset], [label for _, label in sampled_dataset]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "576c5225",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_pass(texts, model_parameters):\n",
    "\n",
    "    with concurrent.futures.ThreadPoolExecutor(max_workers=NUM_PARALLEL_FORWARD_PASS_API_CALLS) as executor:\n",
    "        model_response = executor.map(compute_model_response, texts, [model_parameters]*len(texts))\n",
    "\n",
    "    model_response = list(model_response)\n",
    "\n",
    "    with concurrent.futures.ThreadPoolExecutor(max_workers=NUM_PARALLEL_FORWARD_PASS_API_CALLS) as executor:\n",
    "        predicted_labels = executor.map(extract_from_model_response, model_response)\n",
    "\n",
    "    predicted_labels = list(predicted_labels)\n",
    "\n",
    "    return model_response, predicted_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "07a50c74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_model_parameters(\n",
    "    texts, model_responses, predicted_labels, correct_labels,\n",
    "    model_parameters, metrics,\n",
    "    parameters_update_response_history, metrics_history,\n",
    "):\n",
    "    conversation_history = []\n",
    "    assert len(parameters_update_response_history) == len(metrics_history[:-1])\n",
    "    for parameters_update_response_past, metrics_past in zip(\n",
    "        parameters_update_response_history, metrics_history[:-1]\n",
    "    ):\n",
    "        # the latest metrics will be given in the last message\n",
    "        conversation_history.append(\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": str(metrics_past),  \n",
    "            }\n",
    "        )\n",
    "        conversation_history.append(\n",
    "            {\n",
    "                \"role\": \"assistant\",\n",
    "                \"content\": parameters_update_response_past,\n",
    "            },\n",
    "        )\n",
    "    \n",
    "    mistake_counts = collections.defaultdict(int)\n",
    "    correct_counts = collections.defaultdict(int)\n",
    "\n",
    "    user_message = textwrap.dedent(\n",
    "        f\"\"\"\n",
    "        You are given\n",
    "        - a set of (text, model response, predicted label, correct label)\n",
    "            - predicted label may be None if it is not found in the model response\n",
    "        - the current set of prompts (which may be empty) for an LLM\n",
    "\n",
    "        You will improve the prompts so that the LLM will predict the correct label.\n",
    "        Please spend some time to think what each label means, based on the examples.\n",
    "\n",
    "        The LLM input has the following parameters.\n",
    "        - {list(model_parameters.keys())}\n",
    "        \n",
    "        This is how the parameters are used by the LLM\n",
    "        {inspect.getsource(compute_model_response)}\n",
    "\n",
    "        Please ensure that the prompts contains these following instructions\n",
    "        - The label should have the same exact text as the actual labels\n",
    "        - The information on what the set of labels are\n",
    "            - Some information on what is considered invalid reasons for a label\n",
    "        - A small number of example inputs and outputs\n",
    "        - The LLM should only classify the text. The LLM should not respond to the text or decline classifying the text.\n",
    "        - The LLM should provide a concise reasoning. The reasoning should happen before the label.\n",
    "        - The response should always end with Label: <label>{{label}}</label>\n",
    "            - Note that the label needs exactly the text in the correct label\n",
    "            - Note that you need the html tags\n",
    "        \n",
    "        The current metrics is {str(metrics)}.\n",
    "        If the metrics is especially bad (i.e. correctness is less than random), you likely need to rethink the interpretation of the labels.\n",
    "        \"\"\"\n",
    "    ) + \"\\n\\n\\n\"\n",
    "    \n",
    "    for text, model_response, predicted_label, correct_label in zip(\n",
    "        texts, model_responses, predicted_labels, correct_labels\n",
    "    ):\n",
    "        correctness_verdict = \"\"\n",
    "        correct_labels_set = set(correct_labels)\n",
    "        if predicted_label == correct_label:\n",
    "            correctness_verdict = \"This prediction is correct.\"\n",
    "            if correct_counts[correct_label] > NUM_SAMPLES_CORRECT_GRADIENT_CALCULATION_FOR_EACH_LABEL:\n",
    "                continue\n",
    "            correct_counts[correct_label] += 1\n",
    "        elif predicted_label not in correct_labels_set:\n",
    "            correctness_verdict = \"This predicted label could not be extracted or does not belong to one of the actual labels.\"\n",
    "            if mistake_counts[None] > NUM_SAMPLES_MISTAKE_GRADIENT_CALCULATION_FOR_EACH_LABEL:\n",
    "                continue\n",
    "            mistake_counts[None] += 1\n",
    "        else:\n",
    "            correctness_verdict = \"This prediction is incorrect.\"\n",
    "            if mistake_counts[correct_label] > NUM_SAMPLES_MISTAKE_GRADIENT_CALCULATION_FOR_EACH_LABEL:\n",
    "                continue\n",
    "            mistake_counts[correct_label] += 1\n",
    "            \n",
    "        user_message += textwrap.dedent(\n",
    "            f\"\"\"\n",
    "            <datapoint>\n",
    "                {correctness_verdict}\n",
    "            \n",
    "                <text>{text}</text>\n",
    "\n",
    "                <model_response>{model_response}</model_response>\n",
    "\n",
    "                <predicted_label>{predicted_label}</predicted_label>\n",
    "\n",
    "                <correct_label>{correct_label}</correct_label>\n",
    "            \n",
    "            </datapoint>\n",
    "            \"\"\"\n",
    "        ) + \"\\n\\n\"\n",
    "    \n",
    "    user_message += \"\\n\\n\\nThis the current set of prompts\\n\"\n",
    "    \n",
    "    for model_parameter_key, current_model_parameter_value in model_parameters.items():\n",
    "        user_message += textwrap.dedent(f\"\"\"\n",
    "        <{model_parameter_key}>\n",
    "        {current_model_parameter_value}\n",
    "        </{model_parameter_key}>\n",
    "        \"\"\") + \"\\n\\n\"\n",
    "\n",
    "    user_message += textwrap.dedent(f\"\"\"\n",
    "    Your reply (not the LLM you are tuning prompts for) should include the following\n",
    "    \n",
    "    Your informed interpretation of the labels\n",
    "        \n",
    "    The prompt parameters in the following format within the xml tags\n",
    "    (please make sure each prompt parameter has some meaningful content)    \n",
    "    \"\"\") + \"\\n\\n\"\n",
    "\n",
    "    for model_parameter_key in model_parameters.keys():\n",
    "        user_message += textwrap.dedent(f\"\"\"\n",
    "        <{model_parameter_key}>\n",
    "        the new {model_parameter_key} here\n",
    "        </{model_parameter_key}>\n",
    "        \"\"\") + \"\\n\\n\"        \n",
    "        \n",
    "    user_message += \"\\n\\nPlease spend time to check that your improved prompts will actually fix the mistakes.\"\n",
    "        \n",
    "    conversation_history.append({\"role\": \"user\", \"content\": user_message})\n",
    "    \n",
    "    response = openai_client.chat.completions.create(\n",
    "        model=\"o1-preview\",\n",
    "        messages=conversation_history\n",
    "    )\n",
    "    \n",
    "    token_counts[\"o1_input\"] += response.usage.prompt_tokens\n",
    "    token_counts[\"o1_output\"] += response.usage.completion_tokens\n",
    "    token_costs[\"o1_input\"] += response.usage.prompt_tokens * 3 * 1e-6\n",
    "    token_costs[\"o1_output\"] += response.usage.completion_tokens * 15 * 1e-6\n",
    "\n",
    "    parameters_update_response = response.choices[0].message.content\n",
    "\n",
    "    for model_parameter_key in model_parameters.keys():\n",
    "        groups = re.search(\n",
    "            fr'<{model_parameter_key}>(.*?)</{model_parameter_key}>',\n",
    "            parameters_update_response, re.DOTALL\n",
    "        )\n",
    "        model_parameters[model_parameter_key] = groups.group(1) if groups else \"\"\n",
    "    \n",
    "    return model_parameters, parameters_update_response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1d4f661",
   "metadata": {},
   "source": [
    "# Display functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "040e056a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_metrics(predicted_labels, correct_labels):\n",
    "    metrics = {}\n",
    "    correct_labels_set = set(correct_labels)\n",
    "    for label in sorted(correct_labels_set):\n",
    "        metrics[f\"{label}_recall\"] = recall_score(\n",
    "            [correct_label == label for correct_label in correct_labels],\n",
    "            [predicted_label == label for predicted_label in predicted_labels],\n",
    "        )\n",
    "    metrics[\"accuracy\"] = sum(\n",
    "        predicted_label == correct_label for predicted_label, correct_label in zip(predicted_labels, correct_labels)\n",
    "    )\n",
    "    metrics[\"missing\"] = sum(\n",
    "        [predicted_label not in correct_labels_set for predicted_label in predicted_labels]\n",
    "    ) / len(predicted_labels)\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "04eaf8a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_and_display_prompt_history(model_parameters_history, parameters_update_response_history, metrics_history):\n",
    "\n",
    "    iteration_data_all = []\n",
    "\n",
    "    for model_parameters, parameters_update_response, metrics in itertools.zip_longest(\n",
    "        model_parameters_history, parameters_update_response_history, metrics_history, fillvalue={}\n",
    "    ):\n",
    "        iteration_data = {}\n",
    "        for k,v in model_parameters.items():\n",
    "            iteration_data[k] = v\n",
    "        for k,v in metrics.items():\n",
    "            iteration_data[k] = f\"{v:.3f}\"\n",
    "        if parameters_update_response:\n",
    "            iteration_data[\"parameters_update_response\"] = parameters_update_response\n",
    "        iteration_data_all.append(iteration_data)\n",
    "\n",
    "    df = pd.DataFrame(iteration_data_all).fillna(\"\")\n",
    "\n",
    "    html_prefix = '''\n",
    "    <meta charset=\"UTF-8\">\n",
    "    <style>\n",
    "    table {\n",
    "        border-collapse: collapse;\n",
    "    }\n",
    "    td, th {\n",
    "        border: 1px solid black;\n",
    "        padding: 5px;\n",
    "        vertical-align: top;\n",
    "    }\n",
    "    td {\n",
    "        white-space: pre-wrap;\n",
    "        font-family: monospace;\n",
    "    }\n",
    "    </style>\n",
    "    '''\n",
    "    \n",
    "    os.makedirs(\"html_output\", exist_ok=True)\n",
    "    prompt_info_file_name = \"html_output/prompt-history-classification.html\"\n",
    "    with open(prompt_info_file_name, 'w') as f:\n",
    "        f.write(\n",
    "            html_prefix + df.replace(\n",
    "                {r'\\n': '__NEWLINE__'}, regex=True\n",
    "            ).applymap(str).applymap(html.escape).replace(\n",
    "                {'__NEWLINE__': '<br>'}, regex=True\n",
    "            ).style.set_table_styles(\n",
    "                [\n",
    "                    dict(selector=\"tr:nth-child(even)\", props=[(\"background-color\", \"#f2f2f2\")]),\n",
    "                    dict(selector=\"tr:nth-child(odd)\", props=[(\"background-color\", \"white\")]),\n",
    "                ]\n",
    "            ).render(\n",
    "                index=False, escape=False\n",
    "            )\n",
    "        )\n",
    "\n",
    "    link = f'<a href=\"{prompt_info_file_name}\" target=\"_blank\">{prompt_info_file_name}</a>'\n",
    "    display(HTML(link))\n",
    "        \n",
    "    \n",
    "def save_and_display_current_iteration(iteration_idx, texts, model_response, predicted_labels, correct_labels):\n",
    "    \n",
    "    df = pd.DataFrame({\n",
    "        \"text\": texts,\n",
    "        \"model_response\": model_response,\n",
    "        \"predicted_label\": predicted_labels,\n",
    "        \"correct_label\": correct_labels,\n",
    "    })\n",
    "    \n",
    "    def highlight_diff(row):\n",
    "        if row['predicted_label'] == row['correct_label']:\n",
    "            return ['background-color: #90EE90'] * len(row)  # green\n",
    "        return ['background-color: #FFB6C1'] * len(row)  # red\n",
    "\n",
    "    html_prefix = '''\n",
    "    <meta charset=\"UTF-8\">\n",
    "    <style>\n",
    "    table {\n",
    "        border-collapse: collapse;\n",
    "    }\n",
    "    td, th {\n",
    "        border: 1px solid black;\n",
    "        padding: 5px;\n",
    "        vertical-align: top;\n",
    "    }\n",
    "    td {\n",
    "        white-space: pre-wrap;\n",
    "        font-family: monospace;\n",
    "    }\n",
    "    </style>\n",
    "    '''\n",
    "    \n",
    "    os.makedirs(\"html_output\", exist_ok=True)\n",
    "    iteration_info_file_name = f\"html_output/iteration-classification-{iteration_idx:03}.html\"\n",
    "    with open(iteration_info_file_name, 'w') as f:\n",
    "        f.write(\n",
    "            html_prefix + df.replace(\n",
    "                {r'\\n': '__NEWLINE__'}, regex=True\n",
    "            ).applymap(str).applymap(html.escape).replace(\n",
    "                {'__NEWLINE__': '<br>'}, regex=True\n",
    "            ).style.apply(highlight_diff, axis=1).render(\n",
    "                index=False, escape=False\n",
    "            )\n",
    "        )\n",
    "    \n",
    "    link = f'<a href=\"{iteration_info_file_name}\" target=\"_blank\">{iteration_info_file_name}</a>'\n",
    "    display(HTML(link))   \n",
    "    \n",
    "    os.makedirs(\"html_output\", exist_ok=True)\n",
    "    iteration_info_file_name = f\"html_output/iteration-classification-{iteration_idx:03}-diff.html\"\n",
    "    with open(iteration_info_file_name, 'w') as f:\n",
    "        f.write(\n",
    "            html_prefix + df[df[\"predicted_label\"] != df[\"correct_label\"]].sort_values(\"correct_label\").replace(\n",
    "                {r'\\n': '__NEWLINE__'}, regex=True\n",
    "            ).applymap(str).applymap(html.escape).replace(\n",
    "                {r'__NEWLINE__': '<br>'}, regex=True\n",
    "            ).style.set_table_styles(\n",
    "                [\n",
    "                    dict(selector=\"tr:nth-child(even)\", props=[(\"background-color\", \"#f2f2f2\")]),\n",
    "                    dict(selector=\"tr:nth-child(odd)\", props=[(\"background-color\", \"white\")]),                    \n",
    "                ]\n",
    "            ).render(\n",
    "                index=False, escape=False\n",
    "            )\n",
    "        )\n",
    "\n",
    "    link = f'<a href=\"{iteration_info_file_name}\" target=\"_blank\">{iteration_info_file_name}</a>'\n",
    "    display(HTML(link))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "138725e1",
   "metadata": {},
   "source": [
    "# Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b1647bb3",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/prompt-history-classification.html\" target=\"_blank\">html_output/prompt-history-classification.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-001.html\" target=\"_blank\">html_output/iteration-classification-001.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-001-diff.html\" target=\"_blank\">html_output/iteration-classification-001-diff.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/prompt-history-classification.html\" target=\"_blank\">html_output/prompt-history-classification.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-002.html\" target=\"_blank\">html_output/iteration-classification-002.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-002-diff.html\" target=\"_blank\">html_output/iteration-classification-002-diff.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/prompt-history-classification.html\" target=\"_blank\">html_output/prompt-history-classification.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-003.html\" target=\"_blank\">html_output/iteration-classification-003.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-003-diff.html\" target=\"_blank\">html_output/iteration-classification-003-diff.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/prompt-history-classification.html\" target=\"_blank\">html_output/prompt-history-classification.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-004.html\" target=\"_blank\">html_output/iteration-classification-004.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-004-diff.html\" target=\"_blank\">html_output/iteration-classification-004-diff.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/prompt-history-classification.html\" target=\"_blank\">html_output/prompt-history-classification.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-005.html\" target=\"_blank\">html_output/iteration-classification-005.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"html_output/iteration-classification-005-diff.html\" target=\"_blank\">html_output/iteration-classification-005-diff.html</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model_parameters_history = [{k:v for k,v in model_parameters.items()}]\n",
    "parameters_update_response_history = []\n",
    "metrics_history = []\n",
    "\n",
    "for iteration_idx in range(1, NUM_ITERATIONS+1):\n",
    "    samples, correct_labels = get_texts_and_labels(dataset)\n",
    "\n",
    "    model_responses, predicted_labels = forward_pass(samples, model_parameters)\n",
    "    metrics = calculate_metrics(predicted_labels, correct_labels)\n",
    "    metrics_history.append(metrics)\n",
    "\n",
    "    if iteration_idx != NUM_ITERATIONS:  # don't update parameters for the last iteration\n",
    "        model_parameters, parameters_update_response = update_model_parameters(\n",
    "            samples, model_responses,\n",
    "            predicted_labels, correct_labels,\n",
    "            model_parameters, metrics,\n",
    "            parameters_update_response_history, metrics_history,\n",
    "        )\n",
    "        parameters_update_response_history.append(parameters_update_response)\n",
    "        model_parameters_history.append({k:v for k,v in model_parameters.items()})\n",
    "\n",
    "    save_and_display_prompt_history(model_parameters_history, parameters_update_response_history, metrics_history)\n",
    "    save_and_display_current_iteration(iteration_idx, samples, model_responses, predicted_labels, correct_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a188d56b",
   "metadata": {},
   "source": [
    "# Cost tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9154d3ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "defaultdict(int,\n",
       "            {'haiku_input': 240596,\n",
       "             'haiku_output': 20808,\n",
       "             'o1_input': 37424,\n",
       "             'o1_output': 31341})"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "token_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e509034e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cost: $0.67\n"
     ]
    }
   ],
   "source": [
    "print(f\"Cost: ${sum(token_costs.values()):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0434c634",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}