Llama_3_1_405B_>_8B_Conversion.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "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/Llama_3_1_405B_%3E_8B_Conversion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Llama 3.1 405B to Llama 3.1 8B - part of the `gpt-prompt-engineer` repo\n",
        "\n",
        "This notebook gives you the ability to go from Llama 3.1 405B to Llama 3.1 8B -- reducing costs massively while keeping quality high.\n",
        "\n",
        "This is powered by OctoAI inference. You'll need to sign up for OctoAI and get an API key to continue.\n",
        "\n",
        "By Matt Shumer (https://twitter.com/mattshumer_) and Ben Hamm (https://www.linkedin.com/in/hammben/)\n",
        "\n",
        "Github repo: https://github.com/mshumer/gpt-prompt-engineer\n",
        "\n"
      ],
      "metadata": {
        "id": "2Jyodln24Bda"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oeDuRwSk3tG6"
      },
      "outputs": [],
      "source": [
        "!pip install openai"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Get your OctoAI API key at: https://octo.ai"
      ],
      "metadata": {
        "id": "ULGwpuIi6Rm0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"OCTOAI_API_KEY\"] = \"PLACE YOUR KEY HERE\""
      ],
      "metadata": {
        "id": "xXqMwX5U6Qac"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import HTML, display\n",
        "import re\n",
        "import json\n",
        "import os\n",
        "from openai import OpenAI\n",
        "\n",
        "def set_css():\n",
        "    display(HTML('''\n",
        "    <style>\n",
        "        pre {\n",
        "            white-space: pre-wrap;\n",
        "        }\n",
        "    </style>\n",
        "    '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)\n",
        "\n",
        "# Initialize the OpenAI client with custom base URL\n",
        "client = OpenAI(\n",
        "    base_url=\"https://text.octoai.run/v1\",\n",
        "    api_key=os.environ['OCTOAI_API_KEY'],\n",
        ")\n",
        "\n",
        "# Define model names\n",
        "small_model = \"meta-llama-3.1-8b-instruct\"\n",
        "big_model = \"meta-llama-3.1-405b-instruct\"\n",
        "\n",
        "def generate_candidate_prompts(task, prompt_example, response_example):\n",
        "    system_prompt = \"\"\"<task>Given an example training sample, create seven additional samples for the same task that are even better. Each example should contain a <prompt> and a <response>.</task>\n",
        "\n",
        "<rules>\n",
        "1. Ensure the new examples are diverse and unique from one another.\n",
        "2. They should all be perfect. If you make a mistake, this system won't work.\n",
        "</rules>\n",
        "\n",
        "Respond in this format:\n",
        "<response_format>\n",
        "<example_one>\n",
        "<prompt>\n",
        "PUT_PROMPT_HERE\n",
        "</prompt>\n",
        "<response>\n",
        "PUT_RESPONSE_HERE\n",
        "</response>\n",
        "</example_one>\n",
        "\n",
        "<example_two>\n",
        "<prompt>\n",
        "PUT_PROMPT_HERE\n",
        "</prompt>\n",
        "<response>\n",
        "PUT_RESPONSE_HERE\n",
        "</response>\n",
        "</example_two>\n",
        "\n",
        "...\n",
        "</response_format>\"\"\"\n",
        "\n",
        "    user_content = f\"\"\"<training_task>{task}</training_task>\n",
        "\n",
        "<prompt_example>\n",
        "{prompt_example}\n",
        "</prompt_example>\n",
        "\n",
        "<response_example>\n",
        "{response_example}\n",
        "</response_example>\"\"\"\n",
        "\n",
        "    response = client.chat.completions.create(\n",
        "        model=big_model,\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": system_prompt},\n",
        "            {\"role\": \"user\", \"content\": user_content}\n",
        "        ],\n",
        "        max_tokens=4000,\n",
        "        temperature=0.5\n",
        "    )\n",
        "\n",
        "    response_text = response.choices[0].message.content\n",
        "\n",
        "    # Parse out the prompts and responses\n",
        "    prompts_and_responses = []\n",
        "    examples = re.findall(r'<example_\\w+>(.*?)</example_\\w+>', response_text, re.DOTALL)\n",
        "    for example in examples:\n",
        "        prompt = re.findall(r'<prompt>(.*?)</prompt>', example, re.DOTALL)[0].strip()\n",
        "        response = re.findall(r'<response>(.*?)</response>', example, re.DOTALL)[0].strip()\n",
        "        prompts_and_responses.append({'prompt': prompt, 'response': response})\n",
        "\n",
        "    return prompts_and_responses\n",
        "\n",
        "def generate_system_prompt(task, prompt_examples):\n",
        "    system_prompt = \"\"\"<your_role>Given a user-description of their <task> a set of prompt / response pairs (it'll be in JSON for easy reading) for the types of outputs we want to generate given inputs, write a fantastic system prompt that describes the task to be done perfectly.</your_role>\n",
        "\n",
        "<rules>\n",
        "1. Do this perfectly.\n",
        "2. Respond only with the system prompt, and nothing else. No other text will be allowed.\n",
        "</rules>\n",
        "\n",
        "Respond in this format:\n",
        "<system_prompt>\n",
        "WRITE_SYSTEM_PROMPT_HERE\n",
        "</system_prompt>\"\"\"\n",
        "\n",
        "    user_content = f\"\"\"<task>{task}</task>\n",
        "\n",
        "<prompt_response_examples>\n",
        "{str(prompt_examples)}\n",
        "</prompt_response_examples>\"\"\"\n",
        "\n",
        "    response = client.chat.completions.create(\n",
        "        model=big_model,\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": system_prompt},\n",
        "            {\"role\": \"user\", \"content\": user_content}\n",
        "        ],\n",
        "        max_tokens=1000,\n",
        "        temperature=0.5\n",
        "    )\n",
        "\n",
        "    response_text = response.choices[0].message.content\n",
        "\n",
        "    # Parse out the prompt\n",
        "    generated_system_prompt = response_text.split('<system_prompt>')[1].split('</system_prompt>')[0].strip()\n",
        "\n",
        "    return generated_system_prompt\n",
        "\n",
        "def test_small_model(generated_examples, prompt_example, system_prompt):\n",
        "    messages = [{\"role\": \"system\", \"content\": system_prompt}]\n",
        "\n",
        "    for example in generated_examples:\n",
        "        messages.append({\"role\": \"user\", \"content\": example['prompt']})\n",
        "        messages.append({\"role\": \"assistant\", \"content\": example['response']})\n",
        "\n",
        "    messages.append({\"role\": \"user\", \"content\": prompt_example.strip()})\n",
        "\n",
        "    response = client.chat.completions.create(\n",
        "        model=small_model,\n",
        "        messages=messages,\n",
        "        max_tokens=2000,\n",
        "        temperature=0.5\n",
        "    )\n",
        "\n",
        "    response_text = response.choices[0].message.content\n",
        "\n",
        "    return response_text\n",
        "\n",
        "def run_conversion_process(task, prompt_example, response_example):\n",
        "    print('Generating the prompts / responses...')\n",
        "    # Generate candidate prompts\n",
        "    generated_examples = generate_candidate_prompts(task, prompt_example, response_example)\n",
        "\n",
        "    print('Prompts / responses generated. Now generating system prompt...')\n",
        "\n",
        "    # Generate the system prompt\n",
        "    system_prompt = generate_system_prompt(task, generated_examples)\n",
        "\n",
        "    print('System prompt generated:', system_prompt)\n",
        "\n",
        "    print(f'\\n\\nTesting the new prompt on {small_model}, using your input example...')\n",
        "    # Test the generated examples and system prompt with the small model\n",
        "    small_model_response = test_small_model(generated_examples, prompt_example, system_prompt)\n",
        "\n",
        "    print(f'{small_model} responded with:')\n",
        "    print(small_model_response)\n",
        "\n",
        "    print('\\n\\n!! CHECK THE FILE DIRECTORY, THE PROMPT IS NOW SAVED THERE !!')\n",
        "\n",
        "    # Create a dictionary with all the relevant information\n",
        "    result = {\n",
        "        \"task\": task,\n",
        "        \"initial_prompt_example\": prompt_example,\n",
        "        \"initial_response_example\": response_example,\n",
        "        \"generated_examples\": generated_examples,\n",
        "        \"system_prompt\": system_prompt,\n",
        "        \"small_model_response\": small_model_response\n",
        "    }\n",
        "\n",
        "    # Save the small model prompt to a Python file\n",
        "    with open(\"small_model_prompt.py\", \"w\") as file:\n",
        "        file.write('system_prompt = \"\"\"' + system_prompt + '\"\"\"\\n\\n')\n",
        "\n",
        "        file.write('messages = [\\n')\n",
        "        for example in generated_examples:\n",
        "            file.write('    {\"role\": \"user\", \"content\": \"\"\"' + example['prompt'] + '\"\"\"},\\n')\n",
        "            file.write('    {\"role\": \"assistant\", \"content\": \"\"\"' + example['response'] + '\"\"\"},\\n')\n",
        "\n",
        "        file.write('    {\"role\": \"user\", \"content\": \"\"\"' + prompt_example.strip() + '\"\"\"}\\n')\n",
        "        file.write(']\\n')\n",
        "\n",
        "    return result"
      ],
      "metadata": {
        "id": "ylWOvJEG4N3I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Fill in your task, prompt_example, and response_example here. Make sure you keep the quality really high here... this is the most important step!"
      ],
      "metadata": {
        "id": "y8PCYOdO4ZYF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "task = \"refactoring complex code\"\n",
        "\n",
        "prompt_example = \"\"\"def calculate_total(prices, tax, discount, shipping_fee, gift_wrap_fee, membership_discount):\n",
        "\n",
        "    total = 0\n",
        "\n",
        "    for i in range(len(prices)):\n",
        "\n",
        "        total += prices[i]\n",
        "\n",
        "    if membership_discount != 0:\n",
        "\n",
        "        total = total - (total * (membership_discount / 100))\n",
        "\n",
        "    if discount != 0:\n",
        "\n",
        "        total = total - (total * (discount / 100))\n",
        "\n",
        "    total = total + (total * (tax / 100))\n",
        "\n",
        "    if total < 50:\n",
        "\n",
        "        total += shipping_fee\n",
        "\n",
        "    else:\n",
        "\n",
        "        total += shipping_fee / 2\n",
        "\n",
        "    if gift_wrap_fee != 0:\n",
        "\n",
        "        total += gift_wrap_fee * len(prices)\n",
        "\n",
        "    if total > 1000:\n",
        "\n",
        "        total -= 50\n",
        "\n",
        "    elif total > 500:\n",
        "\n",
        "        total -= 25\n",
        "\n",
        "    total = round(total, 2)\n",
        "\n",
        "    if total < 0:\n",
        "\n",
        "        total = 0\n",
        "\n",
        "    return total\"\"\"\n",
        "\n",
        "response_example = \"\"\"def calculate_total(prices, tax_rate, discount_rate, shipping_fee, gift_wrap_fee, membership_discount_rate):\n",
        "\n",
        "    def apply_percentage_discount(amount, percentage):\n",
        "\n",
        "        return amount * (1 - percentage / 100)\n",
        "\n",
        "    def calculate_shipping_fee(total):\n",
        "\n",
        "        return shipping_fee if total < 50 else shipping_fee / 2\n",
        "\n",
        "    def apply_tier_discount(total):\n",
        "\n",
        "        if total > 1000:\n",
        "\n",
        "            return total - 50\n",
        "\n",
        "        elif total > 500:\n",
        "\n",
        "            return total - 25\n",
        "\n",
        "        return total\n",
        "\n",
        "    subtotal = sum(prices)\n",
        "\n",
        "    subtotal = apply_percentage_discount(subtotal, membership_discount_rate)\n",
        "\n",
        "    subtotal = apply_percentage_discount(subtotal, discount_rate)\n",
        "\n",
        "\n",
        "\n",
        "    total = subtotal * (1 + tax_rate / 100)\n",
        "\n",
        "    total += calculate_shipping_fee(total)\n",
        "\n",
        "    total += gift_wrap_fee * len(prices)\n",
        "\n",
        "\n",
        "\n",
        "    total = apply_tier_discount(total)\n",
        "\n",
        "    total = max(0, round(total, 2))\n",
        "\n",
        "\n",
        "\n",
        "    return total\"\"\""
      ],
      "metadata": {
        "id": "zdPBIjMc4YIj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Now, let's run this system and get our new prompt! At the end, you'll see a new file pop up in the directory that contains everything you'll need to reduce your costs while keeping quality high w/ Llama 3.1 8B!"
      ],
      "metadata": {
        "id": "xu8WkKy44eRF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "result = run_conversion_process(task, prompt_example, response_example)"
      ],
      "metadata": {
        "id": "VxDzg8944eBB"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}