AEWF_L&L_June_26th_LLM_Quality_Optimization_Bootcamp.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/sms-astanley/octoai/blob/main/AEWF_L%26L_June_26th_LLM_Quality_Optimization_Bootcamp.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# LLM Quality Optimization Bootcamp\n",
        "\n",
        "### Author: Thierry Moreau - Co-founder, Head of DevRel @ OctoAI\n",
        "\n",
        "\n",
        "In this notebook you'll learn how to fine-tune an open source LLM (Llama3-8B) from scratch to perform a specialized task - PII redaction via function calling.\n",
        "\n",
        "**We'll show that taking a small and efficient LLM like Llama3-8B and fine-tuning it, you can achieve significant quality improvements over a state of the art model like GPT-4-Turbo, while also achieving significant cost savings (in the order of 100x).**\n",
        "\n",
        "This notebook is divided into 4 parts:\n",
        "\n",
        "1. Fine-tuning dataset collection\n",
        "2. Kick off fine-tuning on OpenPipe\n",
        "3. Deploy your fine-tune on OctoAI\n",
        "4. Evaluate your fine-tune against GPT-4-Turbo\n",
        "\n",
        "![llm deployment cycle](https://raw.githubusercontent.com/tmoreau89/image-assets/main/fine_tune/llm_deployment_cycle.png)\n",
        "\n",
        "\n",
        "## Pre-requisites\n",
        "\n",
        "### OctoAI\n",
        "\n",
        "We'll use OctoAI to deploy our fine-tune. OctoAI offers efficient, customizable and reliable GenAI inference endpoints. You can sign up for an account on http://octoai.cloud/, and get access to the latest and greatest open source LLMs.\n",
        "\n",
        "Create an API token by following [this guide](https://octo.ai/docs/getting-started/how-to-create-an-octoai-access-token) and save it somewhere safe: we'll need it later in this notebook.\n",
        "\n",
        "By signing up you will automatically get $10 in credits, enough to generate 66.7M Llama3-8B tokens.\n",
        "\n",
        "### OpenPipe\n",
        "\n",
        "We'll use OpenPipe to fine tune our Llama3-8B model. OpenPipe lets developers build fine tuning datasets, fire off fine-tune jobs across a variety of base models, and run comprehensive quality evaluations.\n",
        "\n",
        "Get started by signing up for an account: https://openpipe.ai/.\n",
        "\n",
        "### OpenAI (optional)\n",
        "\n",
        "We'll use OpenAI for quality comparisons against our fine-tuned LLM. More specifically we'll compare our fine-tune against GPT-4-Turbo, and GPT-3.5-Turbo.\n",
        "\n",
        "You can create an OpenAI account at the following URL: https://platform.openai.com. Create a new API key on [this link](https://platform.openai.com/api-keys) once you've created an account.\n",
        "\n",
        "### Python Packages\n",
        "\n",
        "Last but not least, run the cell below to install the necessary pip packages."
      ],
      "metadata": {
        "id": "kAP3nI3kUUb0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Ignore the dependency resolver error\n",
        "! pip install -q openai datasets"
      ],
      "metadata": {
        "id": "YLjnW6jWtjyZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from getpass import getpass\n",
        "\n",
        "# Enter your OctoAI Token\n",
        "OCTOAI_TOKEN = getpass()\n",
        "os.environ[\"OCTOAI_TOKEN\"] = OCTOAI_TOKEN"
      ],
      "metadata": {
        "id": "-k9AEzEZzMcD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Enter your OpenAI Token\n",
        "OPENAI_API_KEY = getpass()\n",
        "os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY"
      ],
      "metadata": {
        "id": "KUt2H7CqzjWS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1. Building a Fine-Tuning Dataset\n",
        "\n",
        "A fine-tuned model is only as good as the dataset it's been trained on. Therefore the dataset collection step is critical in order to get a high quality fine tune.\n",
        "\n",
        "There are two approaches to building a fine-tuning dataset using OpenPipe's SDK:\n",
        "* Record a sufficient number of requests and responses from an LLM (e.g. GPT-4)\n",
        "* Upload a pre-built dataset directly in JSONL format\n",
        "\n",
        "In this notebook we'll opt for the latter approach by taking an already labeled dataset from HuggingFace and turning into a synthetic LLM request log.\n",
        "\n",
        "## PII Masking Dataset\n",
        "\n",
        "The dataset we're using in this notebook is the Personally Identifiable Information (PII) masking 200k dataset from [AI4Privacy](https://www.ai4privacy.com/), available via this [link on HuggingFace](https://huggingface.co/datasets/ai4privacy/pii-masking-200k).\n",
        "\n",
        "This dataset has about 200k synthetic text samples that each contain one or more PII entries across 54 PII classes.\n",
        "\n",
        "An example of PII redaction looks as follow.\n",
        "\n",
        "Input text:\n",
        "```\n",
        "Dear Omer, as per our records, your license 78B5R2MVFAHJ48500 is still registered in our records for access to the educational tools.\n",
        "```\n",
        "\n",
        "Redacted text:\n",
        "```\n",
        "Dear [FIRSTNAME], as per our records, your license [VEHICLEVIN] is still registered in our records for access to the educational tools.\n",
        "```\n",
        "\n",
        "Privacy mask:\n",
        "```\n",
        "[ { \"value\": \"Omer\", \"start\": 5, \"end\": 9, \"label\": \"FIRSTNAME\" }, { \"value\": \"78B5R2MVFAHJ48500\", \"start\": 44, \"end\": 61, \"label\": \"VEHICLEVIN\" } ]\n",
        "```\n",
        "\n",
        "## Function calling for PII redaction\n",
        "\n",
        "Instead of training an LLM to do the PII redaction directly on the input text, we'll use the LLM's function calling ability to call a function that will perform the redaction on the original text.\n",
        "\n",
        "Using a function to perform the redaction gives us flexibility to implement different redaction approaches after the LLM has been fine tuned.\n",
        "* We can redact by replacing the PII with the PII class it belogs to, e.g. `Omer` becomes `[FIRSTNAME`]`.\n",
        "* We can redact by replacing the PII with masked information, e.g. `Omer` becomes `XXXXXX`.\n",
        "* We can redact by replacing the PII with a fake PII by mapping each unique original PII to a corresponding fake PII substitue from a database, e.g. `Omer` becomes `Kendall`.\n",
        "\n",
        "First we specify a system prompt that indicates to the LLM what categories it needs to redact:\n",
        "```python\n",
        "system_prompt =\n",
        "\"\"\"\n",
        "You are an expert model trained to redact potentially sensitive information from documents. You have been given a document to redact. Your goal is to accurately redact the sensitive information from the document. Sensitive information can be in one of the following categories:\n",
        "\n",
        "- ACCOUNTNAME: name of an account\n",
        "...\n",
        "- ZIPCODE: zipcode indicating location or address\n",
        "            \n",
        "You are a function calling AI and should return the specific string that needs to be redacted, along with the category of sensitive information that it belongs to. If there is no sensitive information in the document, return no strings.\n",
        "\"\"\"\n",
        "```\n",
        "\n",
        "Second, we'll need to provide the function call specification of the function that performs the actual redaction:\n",
        "```python\n",
        "tool_choice = {\n",
        "  \"type\": \"function\",\n",
        "  \"function\": {\"name\": \"redact\"}\n",
        "}\n",
        "\n",
        "tools = [\n",
        "  {\n",
        "    \"function\": {\n",
        "      \"name\": \"redact\",\n",
        "      \"parameters\": {\n",
        "        \"type\": \"object\",\n",
        "        \"properties\": {\n",
        "          \"fields_to_redact\": {\n",
        "            \"type\": \"array\",\n",
        "            \"items\": {\n",
        "              \"type\": \"object\",\n",
        "              \"required\": [\n",
        "                \"string\",\n",
        "                \"pii_type\"\n",
        "              ],\n",
        "              \"properties\": {\n",
        "                \"string\": {\n",
        "                  \"type\": \"string\",\n",
        "                  \"description\": \"The exact matching string to redact. Include any whitespace or punctuation. Must be an exact string match!\"\n",
        "                },\n",
        "                \"pii_type\": {\n",
        "                  \"enum\": [\n",
        "                    \"ACCOUNTNAME\",\n",
        "                    ...\n",
        "                    \"ZIPCODE\"\n",
        "                  ],\n",
        "                  \"type\": \"string\"\n",
        "                }\n",
        "              }\n",
        "            }\n",
        "          }\n",
        "        }\n",
        "      }\n",
        "    },\n",
        "    \"type\": \"function\"\n",
        "  }\n",
        "]\n",
        "```\n",
        "\n",
        "With the above system prompt and the function call definition, we'll invoke the LLM by passing in the text that needs to be redacted:\n",
        "\n",
        "```python\n",
        "\n",
        "import requests\n",
        "\n",
        "user_prompt = \"Dear Omer, as per our records, your license 78B5R2MVFAHJ48500 is still registered in our records for access to the educational tools.\"\n",
        "\n",
        "req = requests.post(\"https://text.octoai.run/v1/chat/completions\",\n",
        "    headers={\n",
        "        \"Content-Type\": \"application/json\",\n",
        "        \"Authorization\": f\"Bearer {OCTOAI_TOKEN}\"\n",
        "    },\n",
        "    json={\n",
        "        \"messages\": [\n",
        "            {\"role\": \"system\", \"content\": system_prompt},\n",
        "            {\"role\": \"user\", \"content\": user_prompt}]\n",
        "        \"model\": \"openpipe-llama-3-8b-32k\",\n",
        "        \"max_tokens\": 512,\n",
        "        \"presence_penalty\": 0,\n",
        "        \"temperature\": 0,\n",
        "        \"top_p\": 0.9,\n",
        "        \"peft\": lora_asset_name, # you'll know what to assign this to below\n",
        "        \"tool_choice\": tool_choice,\n",
        "        \"tools\": tools\n",
        "    }\n",
        ")\n",
        "\n",
        "print(response.json()[\"choices\"][0][\"message\"][\"tool_calls\"][0][\"function\"][\"arguments\"])\n",
        "```\n",
        "\n",
        "Which when fine-tuned correcty, should return the following chat completions response containing the  arguments to pass into the `redact()` function call:\n",
        "```json\n",
        "{\n",
        "  'fields_to_redact':\n",
        "  [\n",
        "    {\n",
        "      'string': 'Omer',\n",
        "      'pii_type': 'FIRSTNAME'\n",
        "    },\n",
        "    {\n",
        "      'string': '78B5R2MVFAHJ48500',\n",
        "      'pii_type': 'VEHICLEVIN'\n",
        "    }\n",
        "  ]\n",
        "}\n",
        "```"
      ],
      "metadata": {
        "id": "HoBISfX3or4r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the system prompt with all of the PII categories\n",
        "# The nice thing about this system prompt is that it can be very easily extended\n",
        "# to fit your unique use case.\n",
        "system_prompt = \"\"\"\n",
        "You are an expert model trained to redact potentially sensitive information from documents. You have been given a document to redact. Your goal is to accurately redact the sensitive information from the document. Sensitive information can be in one of the following categories:\n",
        "\n",
        "- ACCOUNTNAME: name of an account\n",
        "- ACCOUNTNUMBER: number of an account\n",
        "- AGE: a person's age\n",
        "- AMOUNT: information indicating a certain monetary amount\n",
        "- BIC: a business identifier code\n",
        "- BITCOINADDRESS: bitcoint address, generally stored in a cryptocurrency wallet\n",
        "- BUILDINGNUMBER: number of a building in a physical address\n",
        "- CITY: name of a city indicating location or address\n",
        "- COMPANYNAME: name of a company\n",
        "- COUNTRY: name of a country indicating location or address\n",
        "- CREDITCARDCVV: credit card CVV\n",
        "- CREDITCARDISSUER: credit card issuer\n",
        "- CREDITCARDNUMBER: credit card number\n",
        "- CURRENCY: currency of a balance or transaction\n",
        "- CURRENCYCODE: the code a currency (e.g. USD)\n",
        "- CURRENCYNAME: name of a currency (e.g. US dollar)\n",
        "- CURRENCYSYMBOL: symbol of a currency (e.g. $)\n",
        "- DATE: a specific calendar date\n",
        "- DOB: a specific calendar date representing birth\n",
        "- EMAIL: an email ID\n",
        "- ETHEREUMADDRESS: ethereum address, generally stored in a cryptocurrency wallet\n",
        "- EYECOLOR: eye color, used to identify a person\n",
        "- FIRSTNAME: first name of a person\n",
        "- GENDER: a gender identifier\n",
        "- HEIGHT: height of a person\n",
        "- IBAN: international banking account number\n",
        "- IP: IP address\n",
        "- IPV4: IP v4 address\n",
        "- IPV6: IP v6 address\n",
        "- JOBAREA: job area, specialization or category\n",
        "- JOBTITLE: job title\n",
        "- LASTNAME: last name of a person\n",
        "- LITECOINADDRESS: litecoin address, generally stored in a cryptocurrency wallet\n",
        "- MAC: MAC address\n",
        "- MASKEDNUMBER: masked number\n",
        "- MIDDLENAME: middle name of a person\n",
        "- NEARBYGPSCOORDINATE: nearby GPS coordinates\n",
        "- ORDINALDIRECTION: ordinal direction (north, south, northeast, etc.)\n",
        "- PASSWORD: a secure string used for authentication\n",
        "- PHONEIMEI: the IMEI of a phone\n",
        "- PHONENUMBER: a telephone number\n",
        "- PIN: a personal identificaiton number (PIN)\n",
        "- PREFIX: prefix used to identify a person (Mr., Mrs., Dr. etc.)\n",
        "- SECONDARY ADDRESS: a secondary physical address address\n",
        "- SEX: a sex identifier (male/female)\n",
        "- SSN: a social security number\n",
        "- STATE: name of a state indicating location or address\n",
        "- STREET: name of a street indicating location or address\n",
        "- TIME: time of the day\n",
        "- URL: URL of a website\n",
        "- USERAGENT: user agent to identify the application, operating system, vendor etc.\n",
        "- USERNAME: user name to identify user\n",
        "- VERHICLEVIN: vehicle identification number or license number\n",
        "- VEHICLEVRM: vehicle registration mark\n",
        "- ZIPCODE: zipcode indicating location or address\n",
        "\n",
        "You should return the specific string that needs to be redacted, along with the category of sensitive information that it belongs to. If there is no sensitive information in the document, return no strings.\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "E_-q4Yw-l0jJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the tools for the LLM to invoke\n",
        "\n",
        "tool_choice = {\n",
        "  \"type\": \"function\",\n",
        "  \"function\": {\"name\": \"redact\"}\n",
        "}\n",
        "\n",
        "tools = [\n",
        "  {\n",
        "    \"function\": {\n",
        "      \"name\": \"redact\",\n",
        "      \"parameters\": {\n",
        "        \"type\": \"object\",\n",
        "        \"properties\": {\n",
        "          \"fields_to_redact\": {\n",
        "            \"type\": \"array\",\n",
        "            \"items\": {\n",
        "              \"type\": \"object\",\n",
        "              \"required\": [\n",
        "                \"string\",\n",
        "                \"pii_type\"\n",
        "              ],\n",
        "              \"properties\": {\n",
        "                \"string\": {\n",
        "                  \"type\": \"string\",\n",
        "                  \"description\": \"The exact matching string to redact. Include any whitespace or punctuation. Must be an exact string match!\"\n",
        "                },\n",
        "                \"pii_type\": {\n",
        "                  \"enum\": [\n",
        "                    \"ACCOUNTNAME\",\n",
        "                    \"ACCOUNTNUMBER\",\n",
        "                    \"AGE\",\n",
        "                    \"AMOUNT\",\n",
        "                    \"BIC\",\n",
        "                    \"BITCOINADDRESS\",\n",
        "                    \"BUILDINGNUMBER\",\n",
        "                    \"CITY\",\n",
        "                    \"COMPANYNAME\",\n",
        "                    \"COUNTY\",\n",
        "                    \"CREDITCARDCVV\",\n",
        "                    \"CREDITCARDISSUER\",\n",
        "                    \"CREDITCARDNUMBER\",\n",
        "                    \"CURRENCY\",\n",
        "                    \"CURRENCYCODE\",\n",
        "                    \"CURRENCYNAME\",\n",
        "                    \"CURRENCYSYMBOL\",\n",
        "                    \"DATE\",\n",
        "                    \"DOB\",\n",
        "                    \"EMAIL\",\n",
        "                    \"ETHEREUMADDRESS\",\n",
        "                    \"EYECOLOR\",\n",
        "                    \"FIRSTNAME\",\n",
        "                    \"GENDER\",\n",
        "                    \"HEIGHT\",\n",
        "                    \"IBAN\",\n",
        "                    \"IP\",\n",
        "                    \"IPV4\",\n",
        "                    \"IPV6\",\n",
        "                    \"JOBAREA\",\n",
        "                    \"JOBTITLE\",\n",
        "                    \"JOBTYPE\",\n",
        "                    \"LASTNAME\",\n",
        "                    \"LITECOINADDRESS\",\n",
        "                    \"MAC\",\n",
        "                    \"MASKEDNUMBER\",\n",
        "                    \"MIDDLENAME\",\n",
        "                    \"NEARBYGPSCOORDINATE\",\n",
        "                    \"ORDINALDIRECTION\",\n",
        "                    \"PASSWORD\",\n",
        "                    \"PHONEIMEI\",\n",
        "                    \"PHONENUMBER\",\n",
        "                    \"PIN\",\n",
        "                    \"PREFIX\",\n",
        "                    \"SECONDARYADDRESS\",\n",
        "                    \"SEX\",\n",
        "                    \"SSN\",\n",
        "                    \"STATE\",\n",
        "                    \"STREET\",\n",
        "                    \"TIME\",\n",
        "                    \"URL\",\n",
        "                    \"USERAGENT\",\n",
        "                    \"USERNAME\",\n",
        "                    \"VEHICLEVIN\",\n",
        "                    \"VEHICLEVRM\",\n",
        "                    \"ZIPCODE\"\n",
        "                  ],\n",
        "                  \"type\": \"string\"\n",
        "                }\n",
        "              }\n",
        "            }\n",
        "          }\n",
        "        }\n",
        "      }\n",
        "    },\n",
        "    \"type\": \"function\"\n",
        "  }\n",
        "]"
      ],
      "metadata": {
        "id": "Gpo7VMozl-3V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Preparing the dataset\n",
        "\n",
        "In the next cell we'll load the dataset from Huggingface and generate and build a 10k sample large fine-tuning dataset using a combination of the input text as LLM prompt, and a re-worked privacy mask as the expected tools call response from the LLM."
      ],
      "metadata": {
        "id": "32n1tctTnBg4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "# Load the dataset from huggingface\n",
        "dataset = load_dataset(\"ai4privacy/pii-masking-200k\")"
      ],
      "metadata": {
        "id": "Wsr4U_7LtFdc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "from google.colab import files\n",
        "\n",
        "# OpenPipe dataset size\n",
        "# To improve on accuracy results you can bump the size to a larger value, e.g. 10,000\n",
        "TRAINING_SIZE = 10000\n",
        "\n",
        "# Create a dataset for OpenPipe\n",
        "openpipe_dataset = []\n",
        "for idx, item in enumerate(dataset['train'].select(range(0, TRAINING_SIZE))):\n",
        "    function_arguments = {\n",
        "        \"fields_to_redact\": []\n",
        "    }\n",
        "    for i in item[\"privacy_mask\"]:\n",
        "        function_arguments[\"fields_to_redact\"].append({\n",
        "            \"string\": i[\"value\"],\n",
        "            \"pii_type\": i[\"label\"]\n",
        "        })\n",
        "    dataitem = {\n",
        "        \"messages\": [\n",
        "            {\n",
        "                \"role\": \"system\",\n",
        "                \"content\": system_prompt\n",
        "            },\n",
        "            {\n",
        "                \"role\": \"user\",\n",
        "                \"content\": item[\"source_text\"]\n",
        "            },\n",
        "            {\n",
        "                \"role\": \"assistant\",\n",
        "                \"content\": None,\n",
        "                \"tool_calls\":\n",
        "                    [\n",
        "                        {\n",
        "                            \"id\":\"\",\n",
        "                            \"type\":\"function\",\n",
        "                            \"function\":\n",
        "                            {\n",
        "                                \"name\": \"redact\",\n",
        "                                \"arguments\": json.dumps(function_arguments, indent=2)\n",
        "                            }\n",
        "                        }\n",
        "                    ]\n",
        "            },\n",
        "        ],\n",
        "        \"tools\": tools,\n",
        "        \"tool_choice\": tool_choice\n",
        "    }\n",
        "    openpipe_dataset.append(dataitem)\n",
        "\n",
        "with open('openpipe_dataset.jsonl', 'w') as outfile:\n",
        "    for entry in openpipe_dataset:\n",
        "        json.dump(entry, outfile)\n",
        "        outfile.write('\\n')\n",
        "\n",
        "# This will let you download the file on your browser if you're using Google CoLab\n",
        "files.download('openpipe_dataset.jsonl')"
      ],
      "metadata": {
        "id": "c1Mrqq9ss5Rt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2. Fine tune the LLM\n",
        "\n",
        "We'll use OpenPipe for this step.\n",
        "\n",
        "## Upload the dataset to OpenPipe\n",
        "\n",
        "Check your downloads folder, you should find an `openpipe_dataset.jsonl` in there.\n",
        "\n",
        "Now follow along the instructions on [this page](https://docs.openpipe.ai/features/exporting-data#dataset-export) to upload your dataset on OpenPipe.\n",
        "\n",
        "1. On your OpenPipe console, click on \"Datasets\" listed in the bar on the left.\n",
        "2. Click on \"+ New Dataset\" button at the top right of the window.\n",
        "3. Click on \"Upload Data\" button at the top left of the window.\n",
        "4. Drop the jsonl file that was just downloaded in the \"Upload File\" window.\n",
        "5. Click on the Upload button and wait for the dataset to get uploaded.\n",
        "\n",
        "You'll get the confirmation window below if the dataset gets successfully uploaded.\n",
        "\n",
        "![upload confirmation window](https://raw.githubusercontent.com/tmoreau89/image-assets/main/fine_tune/openpipe_dataset_uploaded.png)\n",
        "\n",
        "\n",
        "You'll see your dataset entries under the \"Dataset view\" - 10,000 of them which should have gotten split into a 9,000 training and 1,000 test set.\n",
        "\n",
        "![dataset view](https://raw.githubusercontent.com/tmoreau89/image-assets/main/fine_tune/dataset_view.png)\n",
        "\n",
        "Hit \"Settings\", and rename your Dataset if need be.\n",
        "\n",
        "## Launch a fine-tune\n",
        "\n",
        "Under the Dataset view on the OpenPipe console, you can launch a fine tune by clicking on the \"Fine Tune\" button at the top right of the window.\n",
        "\n",
        "You can define a model ID - this lets us uniquely identify the resulting fine tuned model.\n",
        "\n",
        "Next you can chose your base model:\n",
        "* You can choose between open source models (Llama, Mistral) or closed source models (GPT). Selecting an open source model gives you ownership of the weights, and lets you deploy the model on the platform of your choice. For this notebook we'll select the \"Llama 3 8B 32K\" model.\n",
        "* You'll see that we have a good working set size to work with with 10k training samples.\n",
        "* Finally you can choose to tweak the advanced options but we'll leave them as-is.\n",
        "\n",
        "![fine tuning settings](https://raw.githubusercontent.com/tmoreau89/image-assets/main/fine_tune/finetuning_launch.png)\n",
        "\n",
        "Let's go ahead and hit \"Start Training\" to kick off the fine-tuning job."
      ],
      "metadata": {
        "id": "ICN_Pa784NWz"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3. Deploy the fine-tuned LLM\n",
        "\n",
        "In this section we'll export the model weights in LoRA FP16 form to be hosted on OctoAI.\n",
        "\n",
        "You'll be informed that the fine tune job has completed by email. Once you've been notified you can proceed with the steps below.\n",
        "\n",
        "## Export your model weights\n",
        "\n",
        "Access your fine-tuned model by clicking on \"Fine Tune\" on the OpenPipe console.\n",
        "Click on the model that was just fine tuned.\n",
        "\n",
        "![model fine tune](https://raw.githubusercontent.com/tmoreau89/image-assets/main/fine_tune/model_finetune.png)\n",
        "\n",
        "At the bottom of the page, you can select a format that the model weights gets exported in. Select \"LoRA:FP32\" under the \"Format\" drop-down and hit Export Weights.\n",
        "\n",
        "It will take a couple of minutes until your model weights are ready for download. When the weights are ready, right click on \"Download Weights\" to copy the link to the weights and set the URL aside, which we'll need in the next step to set `lora_url`.\n",
        "\n",
        "![fine tune download link](https://raw.githubusercontent.com/tmoreau89/image-assets/main/fine_tune/model_finetune_link.png)\n",
        "\n",
        "## Upload the LoRA to OctoAI\n",
        "\n",
        "\n",
        "First, we'll install the `octoai` CLI (and the jq library for later use)\n",
        "\n"
      ],
      "metadata": {
        "id": "4-BcaElTf7jh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!curl https://s3.amazonaws.com/downloads.octoai.cloud/octoai/install_octoai_cli_and_sdk.sh -sSfL | sh\n",
        "!apt-get install jq"
      ],
      "metadata": {
        "id": "biM-NdlpkU7c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "Next, login to the CLI with your token. If you don't have one, you can follow the instructions under the pre-requisites section.\n",
        "\n",
        "To enable you to upload LoRA assets, please add your credit card to the account. Note that you have $10 in free credits, which should more then cover test usage."
      ],
      "metadata": {
        "id": "IDTCkIVwkjzB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!octoai login"
      ],
      "metadata": {
        "id": "oV6q_m54kxut"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below, you'll need to set the `lora_url` to the URL you just copied from the \"Download Weights\" link in the previous step."
      ],
      "metadata": {
        "id": "uiU7l7Kalv3v"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "\n",
        "# Llama-3-8B-Instruct (32K token context)\n",
        "checkpoint_name = \"octoai:openpipe-llama-3-8b-32k\"\n",
        "lora_url = \"SET ME\"\n",
        "assert(lora_url != \"SET ME\") # Please update lora_url\n",
        "\n",
        "# Define an asset name on OctoAI to uniquely identify the LoRA\n",
        "lora_asset_name = \"pii-redaction-finetune-{}\".format(str(random.randint(1000, 999999)))\n",
        "\n",
        "# Set checkpoint name and LoRA URL as env vars\n",
        "%env checkpoint_name=$checkpoint_name\n",
        "%env lora_url=$lora_url\n",
        "%env lora_asset_name=$lora_asset_name"
      ],
      "metadata": {
        "id": "dyegeqoblvcA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below, we upload the LoRA to OctoAI.\n",
        "\n",
        "We need to specify what base checkpoint and architecture (\"engine\") the model corresponds to.\n",
        "\n",
        "The command below uses \"--upload-from-url\" which lets you upload these files from the OpenPipe download URL. Note also that there is an \"upload-from-dir\" that lets you specify a local directory if you've downloaded the LoRA zip file on your local drive.\n",
        "\n",
        "The \"--wait\" flag allows to block until the upload has completed, making scripting possible."
      ],
      "metadata": {
        "id": "Eywtn0P-mxqs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%bash\n",
        "\n",
        "octoai asset create \\\n",
        "  --checkpoint $checkpoint_name \\\n",
        "  --format safetensors \\\n",
        "  --type lora \\\n",
        "  --engine text/llama-3-8b \\\n",
        "  --name $lora_asset_name \\\n",
        "  --data-type fp16 \\\n",
        "  --upload-from-url $lora_url \\\n",
        "  --wait"
      ],
      "metadata": {
        "id": "OixoUUS1mvJT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can double check to make sure the LoRA got added using \"octoai asset get\" command:"
      ],
      "metadata": {
        "id": "HRJcLN76n4rF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!octoai asset get -n $lora_asset_name"
      ],
      "metadata": {
        "id": "i3zNKgWgn-Hl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sanity checking that the fine-tune is running on OctoAI\n",
        "\n",
        "Let's go ahead and use OctoAI to run a test inference with the code below. It's doable by supplying the LoRA name as `peft` parameter (a LoRA is a type of Parameter-Efficient Fine Tune) when making a call to the model."
      ],
      "metadata": {
        "id": "QTsL8swioFKI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "\n",
        "print(\"Using assset \", lora_asset_name)\n",
        "\n",
        "test_prompt = \"Dear Omer, as per our records, your license 78B5R2MVFAHJ48500 is still registered in our records for access to the educational tools. Please feedback on it's operability.\"\n",
        "\n",
        "messages=[\n",
        "    {\n",
        "        \"role\": \"system\",\n",
        "        \"content\": system_prompt\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": test_prompt\n",
        "    },\n",
        "]\n",
        "\n",
        "req = requests.post(\"https://text.octoai.run/v1/chat/completions\",\n",
        "    headers={\n",
        "        \"Content-Type\": \"application/json\",\n",
        "        \"Authorization\": f\"Bearer {OCTOAI_TOKEN}\"\n",
        "    },\n",
        "    json={\n",
        "        \"messages\": messages,\n",
        "        \"model\": \"openpipe-llama-3-8b-32k\",\n",
        "        \"max_tokens\": 512,\n",
        "        \"presence_penalty\": 0,\n",
        "        \"temperature\": 0,\n",
        "        \"top_p\": 0.9,\n",
        "        \"peft\": lora_asset_name,\n",
        "        \"tool_choice\": tool_choice,\n",
        "        \"tools\": tools\n",
        "    }\n",
        ")\n",
        "\n",
        "print(json.dumps(req.json(),indent=4))"
      ],
      "metadata": {
        "id": "fUT1waisoXFs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The output looks good! Let's now run a more exhaustive set of quality evaluations to see where this model stands next to very capable LLMs like GPT-4."
      ],
      "metadata": {
        "id": "UiMkaUJLqGzp"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 4. Quality Evaluations\n",
        "\n",
        "In this section we'll run more extensive quality tests between our fine-tuned model and OpenAI's ChatGPT-4-Turbo.\n",
        "\n",
        "While the other parts of this tutorial were covered by the \\$100 fine-tuning credits on OpenPipe and \\$10 inference credits on OctoAI, you'll have to spend a bit on OpenAI inference to obtain the full set of comparative results.\n",
        "\n",
        "We estimate that running this evaluation on 1000 test samples should cost you about \\$10 in GPT-4-Turbo usage. If you wish to spend less, you can simply reduce the size of the test set (e.g. 100 samples should run you \\$1).\n",
        "\n",
        "## LLM inferences\n",
        "\n",
        "In the code below we'll evaluate our fine tune against GPT-4-Turbo against 1000 test samples. You can of course feel free to scale down the evaluation to a smaller number, e.g. 100 since it'll take some time to generate the responses.\n",
        "\n",
        "Most importanly we need to make sure that the test samples are not taken from the data used to fine-tune our LoRA, this is why we start at offset 10,000 (our fine-tuning dataset was 10,000 large)."
      ],
      "metadata": {
        "id": "UrsJSNB4qPkG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import concurrent.futures\n",
        "\n",
        "# Max concurrent threads\n",
        "MAX_THREADS = 10\n",
        "\n",
        "# Test size\n",
        "TEST_SIZE = 1000\n",
        "\n",
        "# Run the LLM prediction\n",
        "def predict_redaction(text, model):\n",
        "    model_args = {\n",
        "        \"model\": model,\n",
        "        \"messages\": [\n",
        "            {\"role\": \"system\", \"content\": system_prompt},\n",
        "            {\"role\": \"user\", \"content\": text}],\n",
        "        \"max_tokens\": 512,\n",
        "        \"temperature\": 0,\n",
        "        \"tool_choice\": tool_choice,\n",
        "        \"tools\": tools,\n",
        "    }\n",
        "\n",
        "    if model.startswith(\"gpt\"):\n",
        "        # Run on OpenAI\n",
        "        from openai import OpenAI\n",
        "        client = OpenAI(api_key=os.environ.get(\"OPENAI_API_KEY\", \"<your OpenAI API key if not set as env var>\"))\n",
        "        response = client.chat.completions.create(**model_args)\n",
        "        try:\n",
        "            return json.loads(response.choices[0].message.tool_calls[0].function.arguments)\n",
        "        except:\n",
        "            return None\n",
        "    else:\n",
        "        # Run on OctoAI\n",
        "        model_args[\"model\"] = \"openpipe-llama-3-8b-32k\"\n",
        "        model_args[\"peft\"] = model\n",
        "        response = requests.post(\n",
        "            \"https://text.octoai.run/v1/chat/completions\",\n",
        "            headers={\n",
        "                \"Content-Type\": \"application/json\",\n",
        "                \"Authorization\": \"Bearer {}\".format(OCTOAI_TOKEN)\n",
        "            },\n",
        "            json=model_args\n",
        "        )\n",
        "        try:\n",
        "            tool_call = response.json()[\"choices\"][0][\"message\"][\"tool_calls\"][0][\"function\"][\"arguments\"]\n",
        "            return json.loads(tool_call)\n",
        "        except:\n",
        "            return None\n",
        "\n",
        "# Task to run in parallel\n",
        "def task(idx, item):\n",
        "    print(\"Evaluating test sample {}\".format(idx))\n",
        "    entry = {\n",
        "      \"input\": item[\"source_text\"],\n",
        "      \"output\": {\n",
        "         \"fields_to_redact\": []\n",
        "      }\n",
        "    }\n",
        "    for i in item[\"privacy_mask\"]:\n",
        "        entry[\"output\"][\"fields_to_redact\"].append({\n",
        "            \"string\": i[\"value\"],\n",
        "            \"pii_type\": i[\"label\"]\n",
        "        })\n",
        "    entry[\"gpt4_output\"] = predict_redaction(item[\"source_text\"], \"gpt-4-0125-preview\")\n",
        "    entry[\"finetune_output\"] = predict_redaction(item[\"source_text\"], lora_asset_name)\n",
        "    return entry\n",
        "\n",
        "# We store the results in test_data\n",
        "test_data = []\n",
        "\n",
        "with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_THREADS) as executor:\n",
        "    # Submit tasks to the executor\n",
        "    futures = [\n",
        "        executor.submit(task, idx, item) for idx, item in enumerate(dataset['train'].select(range(TRAINING_SIZE, TRAINING_SIZE+TEST_SIZE)))\n",
        "    ]\n",
        "    # Collect the results\n",
        "    test_data = [future.result() for future in concurrent.futures.as_completed(futures)]\n",
        "\n",
        "# Show test results\n",
        "print(test_data)\n"
      ],
      "metadata": {
        "id": "zwmN2K2uxRFT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "## Quality Metric\n",
        "\n",
        "All quality evaluations start by defining a quality metric. In our case, we already have a labeled dataset from AI4Privacy, which we can use as our ground evaluation ground truth.\n",
        "\n",
        "We introduce a scoring system that works fairly simply. Each PII that needs to be redacted is represented as a pair containing:\n",
        "* The PII string itself, e.g. `5943919109159496`\n",
        "* The PII class, e.g. `CREDITCARDNUMBER`\n",
        "\n",
        "We use the SequenceMatcher library to obtain a similarity score between the ground truth PII and the one that's been inferred by the LLMs.\n",
        "\n",
        "If the PII string and class match perfectly, we get a score of 1.0. If any information starts to divert (e.g. LLM classifies PII as `MIDDLENAME` instead of `FIRSTNAME`, the score becomes lower, but is not 0."
      ],
      "metadata": {
        "id": "9Kozwls5wDuO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from difflib import SequenceMatcher\n",
        "\n",
        "def similar(a, b):\n",
        "    a_string = \"{}, {}\".format(a['string'], a['pii_type'])\n",
        "    b_string = \"{}, {}\".format(b['string'], b['pii_type'])\n",
        "    return SequenceMatcher(None, a_string, b_string).ratio()\n",
        "\n",
        "def derive_score(ref, test):\n",
        "    # Assess OpenPipe accuracy\n",
        "    final_score = 0\n",
        "    if test and \"fields_to_redact\" in test:\n",
        "        # sum of all of the best similarity scores across the test PII\n",
        "        score = 0\n",
        "        for t in test[\"fields_to_redact\"]:\n",
        "            # we retain the best similarity score across all pairwise PII comparisons\n",
        "            best_score = 0\n",
        "            for r in ref[\"fields_to_redact\"]:\n",
        "                sim_score = similar(r, t)\n",
        "                if sim_score > best_score:\n",
        "                    best_score = sim_score\n",
        "            score += best_score\n",
        "        # divide the sum by the max of PII classes in reference data, and test data\n",
        "        # this is a simple formula to introduce a penalty in case we have a false positive or false negative\n",
        "        final_score = score/max(len(ref[\"fields_to_redact\"]), len(test[\"fields_to_redact\"]))\n",
        "    return final_score\n",
        "\n",
        "gpt4_total_score = 0\n",
        "finetune_total_score = 0\n",
        "\n",
        "for elem in test_data:\n",
        "    # Retrieve pii masks\n",
        "    ref_output = elem[\"output\"]\n",
        "    gpt4_output = elem[\"gpt4_output\"]\n",
        "    ft_output = elem[\"finetune_output\"]\n",
        "    print(\"Ground truth\")\n",
        "    print(\"\\t{}\".format(ref_output))\n",
        "    # Assess GPT4 accuracy\n",
        "    print(\"Eval GPT4\")\n",
        "    print(\"\\t{}\".format(gpt4_output))\n",
        "    score = derive_score(ref_output, gpt4_output)\n",
        "    print(\"\\tScore = {}\".format(score))\n",
        "    gpt4_total_score += score\n",
        "    # Assess fine-tune accuracy\n",
        "    print(\"Eval fine-tune\")\n",
        "    print(\"\\t{}\".format(ft_output))\n",
        "    score = derive_score(ref_output, ft_output)\n",
        "    print(\"\\tScore = {}\".format(score))\n",
        "    finetune_total_score += score\n",
        "    print(\"\\n\\n\")\n"
      ],
      "metadata": {
        "id": "WWMs_QDB8NtN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Take a look at the scores now for GPT-4-Turbo vs. your fine-tune!\n",
        "\n",
        "GPT-4 costs \\$30.00 / 1M output tokens, while Llama-3-8B on OctoAI costs \\$0.15 / 1M output tokens.\n",
        "\n",
        "Essentially with fine tuning you're getting a 200x cost reduction, while improving overall quality of processing significantly!"
      ],
      "metadata": {
        "id": "nVlTj8oY-YNe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"GPT-4-Turbo Total Score: {:.3f}\".format(gpt4_total_score/TEST_SIZE))\n",
        "print(\"Fine-tune Total Score: {:.3f}\".format(finetune_total_score/TEST_SIZE))"
      ],
      "metadata": {
        "id": "wSbfVIkR-ek5"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}