gpt_prompt_engineer.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mshumer/gpt-prompt-engineer/blob/main/gpt_prompt_engineer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "WljjH8K3s7kG"
      },
      "source": [
        "# gpt-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 OpenAI key.\n",
        "2. If you don't have GPT-4 access, change `model='gpt-4'` in the second cell to `model='gpt-3.5-turbo'`. If you do have access, skip this step.\n",
        "2. In the last cell, fill in the description of your task, up to 15 test cases, and the number of prompts to generate.\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": {},
      "outputs": [],
      "source": [
        "!pip install openai 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 openai\n",
        "from tqdm import tqdm\n",
        "import itertools\n",
        "import wandb\n",
        "from tenacity import retry, stop_after_attempt, wait_exponential\n",
        "\n",
        "openai.api_key = \"ADD YOUR KEY HERE\" # enter your OpenAI 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": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "system_gen_system_prompt = \"\"\"Your job is to generate system prompts for GPT-4, given a description of the use-case and some test cases.\n",
        "\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",
        "\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. 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",
        "\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",
        "\n",
        "Most importantly, output NOTHING but the prompt. Do not include anything else in your message.\"\"\"\n",
        "\n",
        "\n",
        "ranking_system_prompt = \"\"\"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",
        "\n",
        "You will be provided with the task description, the test prompt, and two generations - one for each system prompt.\n",
        "\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",
        "\n",
        "Remember, to be considered 'better', a generation must not just be good, it must be noticeably superior to the other.\n",
        "\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",
        "\n",
        "Respond with your ranking, and nothing else. Be fair and unbiased in your judgement.\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# K is a constant factor that determines how much ratings change\n",
        "K = 32\n",
        "\n",
        "CANDIDATE_MODEL = 'gpt-4'\n",
        "CANDIDATE_MODEL_TEMPERATURE = 0.9\n",
        "\n",
        "GENERATION_MODEL = 'gpt-3.5-turbo'\n",
        "GENERATION_MODEL_TEMPERATURE = 0.8\n",
        "GENERATION_MODEL_MAX_TOKENS = 60\n",
        "\n",
        "N_RETRIES = 3  # number of times to retry a call to the ranking model if it fails\n",
        "RANKING_MODEL = 'gpt-3.5-turbo'\n",
        "RANKING_MODEL_TEMPERATURE = 0.5\n",
        "\n",
        "NUMBER_OF_PROMPTS = 10 # 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 = \"gpt-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": {},
      "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",
        "      \"system_gen_system_prompt\": system_gen_system_prompt, \n",
        "      \"ranking_system_prompt\": ranking_system_prompt,\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": {},
      "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": 13,
      "metadata": {
        "id": "wXeqMQpzzosx"
      },
      "outputs": [],
      "source": [
        "def generate_candidate_prompts(description, test_cases, number_of_prompts):\n",
        "  outputs = openai.ChatCompletion.create(\n",
        "      model=CANDIDATE_MODEL, # change this to gpt-3.5-turbo if you don't have GPT-4 access\n",
        "      messages=[\n",
        "          {\"role\": \"system\", \"content\": system_gen_system_prompt},\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 prompt, and nothing else. Be creative.\"}\n",
        "          ],\n",
        "      temperature=CANDIDATE_MODEL_TEMPERATURE,\n",
        "      n=number_of_prompts)\n",
        "\n",
        "  prompts = []\n",
        "\n",
        "  for i in outputs.choices:\n",
        "    prompts.append(i.message.content)\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, ranking_model_name, ranking_model_temperature):    \n",
        "    score = openai.ChatCompletion.create(\n",
        "        model=ranking_model_name,\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": ranking_system_prompt},\n",
        "            {\"role\": \"user\", \"content\": f\"\"\"Task: {description.strip()}\n",
        "Prompt: {test_case['prompt']}\n",
        "Generation A: {pos1}\n",
        "Generation B: {pos2}\"\"\"}\n",
        "        ],\n",
        "        logit_bias={\n",
        "              '32': 100,  # 'A' token\n",
        "              '33': 100,  # 'B' token\n",
        "        },\n",
        "        max_tokens=1,\n",
        "        temperature=ranking_model_temperature,\n",
        "    ).choices[0].message.content\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):\n",
        "    generation = openai.ChatCompletion.create(\n",
        "        model=GENERATION_MODEL,\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": prompt},\n",
        "            {\"role\": \"user\", \"content\": f\"{test_case['prompt']}\"}\n",
        "        ],\n",
        "        max_tokens=GENERATION_MODEL_MAX_TOKENS,\n",
        "        temperature=GENERATION_MODEL_TEMPERATURE,\n",
        "    ).choices[0].message.content\n",
        "    return generation\n",
        "\n",
        "def test_candidate_prompts(test_cases, description, 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)\n",
        "          generation2 = get_generation(prompt2, test_case)\n",
        "\n",
        "          # Rank the outputs\n",
        "          score1 = get_score(description, test_case, generation1, generation2, RANKING_MODEL, RANKING_MODEL_TEMPERATURE)\n",
        "          score2 = get_score(description, test_case, generation2, generation1, 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",
        "\n",
        "def generate_optimal_prompt(description, test_cases, number_of_prompts=10, use_wandb=False): \n",
        "  if use_wandb:\n",
        "    wandb_table = wandb.Table(columns=[\"Prompt\", \"Ranking\"])\n",
        "    if wandb.run is None:\n",
        "      start_wandb_run()\n",
        "\n",
        "  prompts = generate_candidate_prompts(description, test_cases, number_of_prompts)\n",
        "  prompt_ratings = test_candidate_prompts(test_cases, description, prompts)\n",
        "\n",
        "  # Print the final ELO ratingsz\n",
        "  table = PrettyTable()\n",
        "  table.field_names = [\"Prompt\", \"Rating\"]\n",
        "  for prompt, rating in sorted(prompt_ratings.items(), key=lambda item: item[1], reverse=True):\n",
        "      table.add_row([prompt, rating])\n",
        "      if use_wandb:\n",
        "         wandb_table.add_data(prompt, rating)\n",
        "\n",
        "  if use_wandb: # log the results to a Weights & Biases table and finsih the run\n",
        "    wandb.log({\"prompt_ratings\": wandb_table})\n",
        "    wandb.finish()\n",
        "  print(table)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "MJSSKFfV_X9F"
      },
      "source": [
        "# In the cell below, fill in your description and test cases"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vCZvLyDepxFP"
      },
      "outputs": [],
      "source": [
        "description = \"Given a prompt, generate a landing page headline.\" # this style of description tends to work well\n",
        "\n",
        "test_cases = [\n",
        "    {\n",
        "        'prompt': 'Promoting an innovative new fitness app, Smartly',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Why a vegan diet is beneficial for your health',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Introducing a new online course on digital marketing',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Launching a new line of eco-friendly clothing',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Promoting a new travel blog focusing on budget travel',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Advertising a new software for efficient project management',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Introducing a new book on mastering Python programming',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Promoting a new online platform for learning languages',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Advertising a new service for personalized meal plans',\n",
        "    },\n",
        "    {\n",
        "        'prompt': 'Launching a new app for mental health and mindfulness',\n",
        "    }\n",
        "]\n",
        "\n",
        "if use_wandb:\n",
        "    wandb.config.update({\"description\": description, \n",
        "                        \"test_cases\": test_cases})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "generate_optimal_prompt(description, test_cases, NUMBER_OF_PROMPTS, use_wandb)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyPMYK+Pn5QaRzPmh3T5a9ca",
      "include_colab_link": true,
      "provenance": []
    },
    "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
}