claude_prompt_engineer.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mshumer/gpt-prompt-engineer/blob/main/claude_prompt_engineer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WljjH8K3s7kG"
      },
      "source": [
        "# claude-prompt-engineer\n",
        "By Matt Shumer (https://twitter.com/mattshumer_)\n",
        "\n",
        "Github repo: https://github.com/mshumer/gpt-prompt-engineer\n",
        "\n",
        "Generate an optimal prompt for a given task.\n",
        "\n",
        "To generate a prompt:\n",
        "1. In the first cell, add in your Anthropic key.\n",
        "2. If you want, adjust your settings in the `Adjust settings here` cell\n",
        "2. In the last cell, fill in the description of your task, the variables the system should account for.\n",
        "3. Run all the cells! The AI will generate a number of candidate prompts, and test them all to find the best one!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JqB2ZUD8lgek",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8d8d2146-5c61-4f8d-f6fa-9701740a7aea"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m195.4/195.4 kB\u001b[0m \u001b[31m15.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m263.5/263.5 kB\u001b[0m \u001b[31m11.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.7/62.7 kB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "!pip install prettytable tqdm tenacity wandb -qq"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dQmMZdkG_RA5"
      },
      "outputs": [],
      "source": [
        "from prettytable import PrettyTable\n",
        "import time\n",
        "import requests\n",
        "from tqdm import tqdm\n",
        "import itertools\n",
        "import wandb\n",
        "from tenacity import retry, stop_after_attempt, wait_exponential\n",
        "\n",
        "ANTHROPIC_API_KEY = \"ADD YOUR KEY HERE\" # enter your Anthropic API key here\n",
        "\n",
        "use_wandb = False # set to True if you want to use wandb to log your config and results"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Adjust settings here"
      ],
      "metadata": {
        "id": "uzDKqEFd5wI5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nIklSQohlgel"
      },
      "outputs": [],
      "source": [
        "# K is a constant factor that determines how much ratings change\n",
        "K = 32\n",
        "\n",
        "CANDIDATE_MODEL = 'claude-3-opus-20240229'\n",
        "CANDIDATE_MODEL_TEMPERATURE = 0.9\n",
        "\n",
        "GENERATION_MODEL = 'claude-3-opus-20240229'\n",
        "GENERATION_MODEL_TEMPERATURE = 0.8\n",
        "GENERATION_MODEL_MAX_TOKENS = 800\n",
        "\n",
        "TEST_CASE_MODEL = 'claude-3-opus-20240229'\n",
        "TEST_CASE_MODEL_TEMPERATURE = .8\n",
        "\n",
        "NUMBER_OF_TEST_CASES = 10 # this determines how many test cases to generate... the higher, the more expensive, but the better the results will be\n",
        "\n",
        "N_RETRIES = 3  # number of times to retry a call to the ranking model if it fails\n",
        "RANKING_MODEL = 'claude-3-opus-20240229'\n",
        "RANKING_MODEL_TEMPERATURE = 0.5\n",
        "\n",
        "NUMBER_OF_PROMPTS = 5 # this determines how many candidate prompts to generate... the higher, the more expensive, but the better the results will be\n",
        "\n",
        "WANDB_PROJECT_NAME = \"claude-prompt-eng\" # used if use_wandb is True, Weights &| Biases project name\n",
        "WANDB_RUN_NAME = None # used if use_wandb is True, optionally set the Weights & Biases run name to identify this run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HyIsyixwlgel"
      },
      "outputs": [],
      "source": [
        "def start_wandb_run():\n",
        "  # start a new wandb run and log the config\n",
        "  wandb.init(\n",
        "    project=WANDB_PROJECT_NAME,\n",
        "    name=WANDB_RUN_NAME,\n",
        "    config={\n",
        "      \"K\": K,\n",
        "      \"candiate_model\": CANDIDATE_MODEL,\n",
        "      \"candidate_model_temperature\": CANDIDATE_MODEL_TEMPERATURE,\n",
        "      \"generation_model\": GENERATION_MODEL,\n",
        "      \"generation_model_temperature\": GENERATION_MODEL_TEMPERATURE,\n",
        "      \"generation_model_max_tokens\": GENERATION_MODEL_MAX_TOKENS,\n",
        "      \"n_retries\": N_RETRIES,\n",
        "      \"ranking_model\": RANKING_MODEL,\n",
        "      \"ranking_model_temperature\": RANKING_MODEL_TEMPERATURE,\n",
        "      \"number_of_prompts\": NUMBER_OF_PROMPTS\n",
        "      })\n",
        "\n",
        "  return"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zk_2Uut-lgel"
      },
      "outputs": [],
      "source": [
        "# Optional logging to Weights & Biases to reocrd the configs, prompts and results\n",
        "if use_wandb:\n",
        "  start_wandb_run()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wXeqMQpzzosx"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import re\n",
        "\n",
        "def remove_first_line(test_string):\n",
        "    if test_string.startswith(\"Here\") and test_string.split(\"\\n\")[0].strip().endswith(\":\"):\n",
        "        return re.sub(r'^.*\\n', '', test_string, count=1)\n",
        "    return test_string\n",
        "\n",
        "def generate_candidate_prompts(description, input_variables, test_cases, number_of_prompts):\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "\n",
        "    variable_descriptions = \"\\n\".join(f\"{var['variable']}: {var['description']}\" for var in input_variables)\n",
        "\n",
        "    data = {\n",
        "        \"model\": CANDIDATE_MODEL,\n",
        "        \"max_tokens\": 1500,\n",
        "        \"temperature\": CANDIDATE_MODEL_TEMPERATURE,\n",
        "        \"system\": f\"\"\"Your job is to generate system prompts for Claude 3, given a description of the use-case, some test cases/input variable examples that will help you understand what the prompt will need to be good at.\n",
        "The prompts you will be generating will be for freeform tasks, such as generating a landing page headline, an intro paragraph, solving a math problem, etc.\n",
        "In your generated prompt, you should describe how the AI should behave in plain English. Include what it will see, and what it's allowed to output.\n",
        "<most_important>Make sure to incorporate the provided input variable placeholders into the prompt, using placeholders like {{{{VARIABLE_NAME}}}} for each variable. Ensure you place placeholders inside four squiggly lines like {{{{VARIABLE_NAME}}}}. At inference time/test time, we will slot the variables into the prompt, like a template.</most_important>\n",
        "Be creative with prompts to get the best possible results. The AI knows it's an AI -- you don't need to tell it this.\n",
        "You will be graded based on the performance of your prompt... but don't cheat! You cannot include specifics about the test cases in your prompt. Any prompts with examples will be disqualified.\n",
        "Here are the input variables and their descriptions:\n",
        "{variable_descriptions}\n",
        "Most importantly, output NOTHING but the prompt (with the variables contained in it like {{{{VARIABLE_NAME}}}}). Do not include anything else in your message.\"\"\",\n",
        "        \"messages\": [\n",
        "            {\"role\": \"user\", \"content\": f\"Here are some test cases:`{test_cases}`\\n\\nHere is the description of the use-case: `{description.strip()}`\\n\\nRespond with your flexible system prompt, and nothing else. Be creative, and remember, the goal is not to complete the task, but write a prompt that will complete the task.\"},\n",
        "        ]\n",
        "    }\n",
        "\n",
        "    prompts = []\n",
        "\n",
        "    for i in range(number_of_prompts):\n",
        "        response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "\n",
        "        message = response.json()\n",
        "\n",
        "        response_text = message['content'][0]['text']\n",
        "\n",
        "        prompts.append(remove_first_line(response_text))\n",
        "\n",
        "    return prompts\n",
        "\n",
        "def expected_score(r1, r2):\n",
        "    return 1 / (1 + 10**((r2 - r1) / 400))\n",
        "\n",
        "def update_elo(r1, r2, score1):\n",
        "    e1 = expected_score(r1, r2)\n",
        "    e2 = expected_score(r2, r1)\n",
        "    return r1 + K * (score1 - e1), r2 + K * ((1 - score1) - e2)\n",
        "\n",
        "# Get Score - retry up to N_RETRIES times, waiting exponentially between retries.\n",
        "@retry(stop=stop_after_attempt(N_RETRIES), wait=wait_exponential(multiplier=1, min=4, max=70))\n",
        "def get_score(description, test_case, pos1, pos2, input_variables, ranking_model_name, ranking_model_temperature):\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "\n",
        "    variable_values = \"\\n\".join(f\"{var['variable']}: {test_case.get(var['variable'], '')}\" for var in input_variables)\n",
        "\n",
        "    data = {\n",
        "        \"model\": RANKING_MODEL,\n",
        "        \"max_tokens\": 1,\n",
        "        \"temperature\": ranking_model_temperature,\n",
        "        \"system\": f\"\"\"Your job is to rank the quality of two outputs generated by different prompts. The prompts are used to generate a response for a given task.\n",
        "You will be provided with the task description, input variable values, and two generations - one for each system prompt.\n",
        "Rank the generations in order of quality. If Generation A is better, respond with 'A'. If Generation B is better, respond with 'B'.\n",
        "Remember, to be considered 'better', a generation must not just be good, it must be noticeably superior to the other.\n",
        "Also, keep in mind that you are a very harsh critic. Only rank a generation as better if it truly impresses you more than the other.\n",
        "Respond with your ranking ('A' or 'B'), and nothing else. Be fair and unbiased in your judgement.\"\"\",\n",
        "        \"messages\": [\n",
        "            {\"role\": \"user\", \"content\": f\"\"\"Task: {description.strip()}\n",
        "Variables: {test_case['variables']}\n",
        "Generation A: {remove_first_line(pos1)}\n",
        "Generation B: {remove_first_line(pos2)}\"\"\"},\n",
        "        ]\n",
        "    }\n",
        "\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "\n",
        "    message = response.json()\n",
        "\n",
        "    score = message['content'][0]['text']\n",
        "\n",
        "    return score\n",
        "\n",
        "@retry(stop=stop_after_attempt(N_RETRIES), wait=wait_exponential(multiplier=1, min=4, max=70))\n",
        "def get_generation(prompt, test_case, input_variables):\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "\n",
        "\n",
        "    # Replace variable placeholders in the prompt with their actual values from the test case\n",
        "    for var_dict in test_case['variables']:\n",
        "        for variable_name, variable_value in var_dict.items():\n",
        "            prompt = prompt.replace(f\"{{{{{variable_name}}}}}\", variable_value)\n",
        "\n",
        "    data = {\n",
        "        \"model\": GENERATION_MODEL,\n",
        "        \"max_tokens\": GENERATION_MODEL_MAX_TOKENS,\n",
        "        \"temperature\": GENERATION_MODEL_TEMPERATURE,\n",
        "        \"system\": 'Complete the task perfectly.',\n",
        "        \"messages\": [\n",
        "            {\"role\": \"user\", \"content\": prompt},\n",
        "        ]\n",
        "    }\n",
        "\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "\n",
        "    message = response.json()\n",
        "\n",
        "    generation = message['content'][0]['text']\n",
        "\n",
        "    return generation\n",
        "\n",
        "def test_candidate_prompts(test_cases, description, input_variables, prompts):\n",
        "    # Initialize each prompt with an ELO rating of 1200\n",
        "    prompt_ratings = {prompt: 1200 for prompt in prompts}\n",
        "\n",
        "    # Calculate total rounds for progress bar\n",
        "    total_rounds = len(test_cases) * len(prompts) * (len(prompts) - 1) // 2\n",
        "\n",
        "    # Initialize progress bar\n",
        "    pbar = tqdm(total=total_rounds, ncols=70)\n",
        "\n",
        "    # For each pair of prompts\n",
        "    for prompt1, prompt2 in itertools.combinations(prompts, 2):\n",
        "        # For each test case\n",
        "        for test_case in test_cases:\n",
        "            # Update progress bar\n",
        "            pbar.update()\n",
        "\n",
        "            # Generate outputs for each prompt\n",
        "            generation1 = get_generation(prompt1, test_case, input_variables)\n",
        "            generation2 = get_generation(prompt2, test_case, input_variables)\n",
        "\n",
        "            # Rank the outputs\n",
        "            score1 = get_score(description, test_case, generation1, generation2, input_variables, RANKING_MODEL, RANKING_MODEL_TEMPERATURE)\n",
        "            score2 = get_score(description, test_case, generation2, generation1, input_variables, RANKING_MODEL, RANKING_MODEL_TEMPERATURE)\n",
        "\n",
        "            # Convert scores to numeric values\n",
        "            score1 = 1 if score1 == 'A' else 0 if score1 == 'B' else 0.5\n",
        "            score2 = 1 if score2 == 'B' else 0 if score2 == 'A' else 0.5\n",
        "\n",
        "            # Average the scores\n",
        "            score = (score1 + score2) / 2\n",
        "\n",
        "            # Update ELO ratings\n",
        "            r1, r2 = prompt_ratings[prompt1], prompt_ratings[prompt2]\n",
        "            r1, r2 = update_elo(r1, r2, score)\n",
        "            prompt_ratings[prompt1], prompt_ratings[prompt2] = r1, r2\n",
        "\n",
        "            # Print the winner of this round\n",
        "            if score > 0.5:\n",
        "                print(f\"Winner: {prompt1}\")\n",
        "            elif score < 0.5:\n",
        "                print(f\"Winner: {prompt2}\")\n",
        "            else:\n",
        "                print(\"Draw\")\n",
        "\n",
        "    # Close progress bar\n",
        "    pbar.close()\n",
        "\n",
        "    return prompt_ratings\n",
        "\n",
        "def generate_optimal_prompt(description, input_variables, num_test_cases=10, number_of_prompts=10, use_wandb=False):\n",
        "    if use_wandb:\n",
        "        wandb_table = wandb.Table(columns=[\"Prompt\", \"Ranking\"] + [var[\"variable\"] for var in input_variables])\n",
        "        if wandb.run is None:\n",
        "            start_wandb_run()\n",
        "\n",
        "    test_cases = generate_test_cases(description, input_variables, num_test_cases)\n",
        "    prompts = generate_candidate_prompts(description, input_variables, test_cases, number_of_prompts)\n",
        "    print('Here are the possible prompts:', prompts)\n",
        "    prompt_ratings = test_candidate_prompts(test_cases, description, input_variables, prompts)\n",
        "\n",
        "    table = PrettyTable()\n",
        "    table.field_names = [\"Prompt\", \"Rating\"] + [var[\"variable\"] for var in input_variables]\n",
        "    for prompt, rating in sorted(prompt_ratings.items(), key=lambda item: item[1], reverse=True):\n",
        "        # Use the first test case as an example for displaying the input variables\n",
        "        example_test_case = test_cases[0]\n",
        "        table.add_row([prompt, rating] + [example_test_case.get(var[\"variable\"], \"\") for var in input_variables])\n",
        "        if use_wandb:\n",
        "            wandb_table.add_data(prompt, rating, *[example_test_case.get(var[\"variable\"], \"\") for var in input_variables])\n",
        "\n",
        "    if use_wandb:\n",
        "        wandb.log({\"prompt_ratings\": wandb_table})\n",
        "        wandb.finish()\n",
        "    print(table)\n",
        "\n",
        "def generate_test_cases(description, input_variables, num_test_cases):\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "\n",
        "    variable_descriptions = \"\\n\".join(f\"{var['variable']}: {var['description']}\" for var in input_variables)\n",
        "\n",
        "    data = {\n",
        "        \"model\": CANDIDATE_MODEL,\n",
        "        \"max_tokens\": 1500,\n",
        "        \"temperature\": CANDIDATE_MODEL_TEMPERATURE,\n",
        "        \"system\": f\"\"\"You are an expert at generating test cases for evaluating AI-generated content.\n",
        "Your task is to generate a list of {num_test_cases} test case prompts based on the given description and input variables.\n",
        "Each test case should be a JSON object with a 'test_design' field containing the overall idea of this test case, and a list of additional JSONs for each input variable, called 'variables'.\n",
        "The test cases should be diverse, covering a range of topics and styles relevant to the description.\n",
        "Here are the input variables and their descriptions:\n",
        "{variable_descriptions}\n",
        "Return the test cases as a JSON list, with no other text or explanation.\"\"\",\n",
        "        \"messages\": [\n",
        "            {\"role\": \"user\", \"content\": f\"Description: {description.strip()}\\n\\nGenerate the test cases. Make sure they are really, really great and diverse:\"},\n",
        "        ]\n",
        "    }\n",
        "\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "    message = response.json()\n",
        "\n",
        "    response_text = message['content'][0]['text']\n",
        "\n",
        "    test_cases = json.loads(response_text)\n",
        "\n",
        "    print('Here are the test cases:', test_cases)\n",
        "\n",
        "    return test_cases"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MJSSKFfV_X9F"
      },
      "source": [
        "# In the cell below, fill in your description and input variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vCZvLyDepxFP"
      },
      "outputs": [],
      "source": [
        "## Example usage\n",
        "description = \"Given a prompt, generate a personalized email response.\" # this style of description tends to work well\n",
        "\n",
        "input_variables = [\n",
        "    {\"variable\": \"SENDER_NAME\", \"description\": \"The name of the person who sent the email.\"},\n",
        "    {\"variable\": \"RECIPIENT_NAME\", \"description\": \"The name of the person receiving the email.\"},\n",
        "    {\"variable\": \"TOPIC\", \"description\": \"The main topic or subject of the email. One to two sentences.\"}\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "if use_wandb:\n",
        "    wandb.config.update({\"description\": description,\n",
        "                         \"input_variables\": input_variables,\n",
        "                         \"num_test_cases\": NUMBER_OF_TEST_CASES,\n",
        "                         \"number_of_prompts\": NUMBER_OF_PROMPTS})"
      ],
      "metadata": {
        "id": "4gxg9uf_vIlr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run this cell to start the prompt engineering process!"
      ],
      "metadata": {
        "id": "Mo9B4x4R_YVl"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7rWjcL2hlgen"
      },
      "outputs": [],
      "source": [
        "generate_optimal_prompt(description, input_variables, NUMBER_OF_TEST_CASES, NUMBER_OF_PROMPTS, use_wandb)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}