Adding automated_embeddings

notebooks/automated_embeddings.ipynb

@@ -0,0 +1,721 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "5d9aca72-957a-4ee2-862f-e011b9cd3a62",
+   "metadata": {},
+   "source": [
+    "# Introduction\n",
+    "## Goal\n",
+    "I have a dataset I want to embed for semantic search (or QA, or RAG), I want the easiest way to do embed this and put it in a new dataset.\n",
+    "\n",
+    "## Approach\n",
+    "Im using a dataset from my favorite subreddit [r/bestofredditorupdates](). Since it has such long entries, I will use the new [jinaai/jina-embeddings-v2-base-en](https://huggingface.co/jinaai/jina-embeddings-v2-base-en) since it has an 8k context length. Since Im GPU-poor I will deploy this using [Inference Endpoint](https://huggingface.co/inference-endpoints) to save money and time. To follow this you will need to add a payment method. To make it even easier, I'll make this fully API based."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d2534669-003d-490c-9d7a-32607fa5f404",
+   "metadata": {},
+   "source": [
+    "# Setup"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b6f72042-173d-4a72-ade1-9304b43b528d",
+   "metadata": {},
+   "source": [
+    "## Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "e2beecdd-d033-4736-bd45-6754ec53b4ac",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import asyncio\n",
+    "from getpass import getpass\n",
+    "import json\n",
+    "from pathlib import Path\n",
+    "import time\n",
+    "\n",
+    "from aiohttp import ClientSession, ClientTimeout\n",
+    "from datasets import load_dataset, Dataset, DatasetDict\n",
+    "from huggingface_hub import notebook_login\n",
+    "import pandas as pd\n",
+    "import requests\n",
+    "from tqdm.auto import tqdm"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5eece903-64ce-435d-a2fd-096c0ff650bf",
+   "metadata": {},
+   "source": [
+    "## Config\n",
+    "You need to fill this in with your desired repos. Note I used 5 for the `MAX_WORKERS` since `jina-embeddings-v2` are quite memory hungry. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "dcd7daed-6aca-4fe7-85ce-534bdcd8bc87",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "dataset_in = 'derek-thomas/dataset-creator-reddit-bestofredditorupdates'\n",
+    "dataset_out = \"processed-bestofredditorupdates\"\n",
+    "endpoint_name = \"boru-jina-embeddings-demo\"\n",
+    "\n",
+    "MAX_WORKERS = 5  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "88cdbd73-5923-4ae9-9940-b6be935f70fa",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "What is your Hugging Face 🤗 username? (with a credit card) ········\n",
+      "What is your Hugging Face 🤗 token? ········\n"
+     ]
+    }
+   ],
+   "source": [
+    "username = getpass(prompt=\"What is your Hugging Face 🤗 username? (with an added payment method)\")\n",
+    "hf_token = getpass(prompt='What is your Hugging Face 🤗 token?')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b972a719-2aed-4d2e-a24f-fae7776d5fa4",
+   "metadata": {},
+   "source": [
+    "## Get Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "27835fa4-3a4f-44b1-a02a-5e31584a1bba",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Dataset({\n",
+       "    features: ['date_utc', 'title', 'flair', 'content', 'poster', 'permalink', 'id', 'content_length', 'score'],\n",
+       "    num_rows: 9991\n",
+       "})"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dataset = load_dataset(dataset_in, token=hf_token)\n",
+    "dataset['train']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "8846087e-4d0d-4c0e-8aeb-ea95d9e97126",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(9991,\n",
+       " {'date_utc': Timestamp('2022-12-31 18:16:22'),\n",
+       "  'title': 'To All BORU contributors, Thank you :)',\n",
+       "  'flair': 'CONCLUDED',\n",
+       "  'content': '[removed]',\n",
+       "  'poster': 'IsItAcOnSeQuEnCe',\n",
+       "  'permalink': '/r/BestofRedditorUpdates/comments/10004zw/to_all_boru_contributors_thank_you/',\n",
+       "  'id': '10004zw',\n",
+       "  'content_length': 9,\n",
+       "  'score': 1})"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "documents = dataset['train'].to_pandas().to_dict('records')\n",
+    "len(documents), documents[0]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93096cbc-81c6-4137-a283-6afb0f48fbb9",
+   "metadata": {},
+   "source": [
+    "# Inference Endpoints\n",
+    "## Create Inference Endpoint\n",
+    "We are going to use the [API](https://huggingface.co/docs/inference-endpoints/api_reference) to create an [Inference Endpoint](https://huggingface.co/inference-endpoints). This should provide a few main benefits:\n",
+    "- It's convenient (No clicking)\n",
+    "- It's repeatable (We have the code to run it easily)\n",
+    "- It's cheaper (No time spent waiting for it to load, and automatically shut it down)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "3a8f67b9-6ac6-4b5e-91ee-e48463191e1b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "headers = {\n",
+    "\t\"Authorization\": f\"Bearer {hf_token}\",\n",
+    "\t\"Content-Type\": \"application/json\"\n",
+    "}\n",
+    "base_url = f\"https://api.endpoints.huggingface.cloud/v2/endpoint/{username}\"\n",
+    "endpoint_url = f\"https://api.endpoints.huggingface.cloud/v2/endpoint/{username}/{endpoint_name}\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0f2c97dc-34e8-49e9-b60e-f5b7366294c0",
+   "metadata": {},
+   "source": [
+    "There are a few design choices here:\n",
+    "- I'm using the `g5.2xlarge` since it is big and `jina-embeddings-v2` are memory hungry (remember the 8k context length). \n",
+    "- I didnt alter the default `MAX_BATCH_TOKENS` or `MAX_CONCURRENT_REQUESTS`\n",
+    "    - You should consider this if you are making this production ready\n",
+    "    - You will need to restrict these to match the HW you are running on\n",
+    "- As mentioned before, I chose the repo and the corresponding revision\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "f1ea29cb-b69d-4340-859f-3646d650c68e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "202\n"
+     ]
+    }
+   ],
+   "source": [
+    "data = {\n",
+    "    \"accountId\": None,\n",
+    "    \"compute\": {\n",
+    "        \"accelerator\": \"gpu\",\n",
+    "        \"instanceType\": \"g5.2xlarge\",\n",
+    "        \"instanceSize\": \"medium\",\n",
+    "        \"scaling\": {\n",
+    "            \"maxReplica\": 1,\n",
+    "            \"minReplica\": 1\n",
+    "        }\n",
+    "    },\n",
+    "    \"model\": {\n",
+    "        \"framework\": \"pytorch\",\n",
+    "        \"image\": {\n",
+    "          \"custom\": {\n",
+    "            \"url\": \"ghcr.io/huggingface/text-embeddings-inference:0.3.0\",\n",
+    "            \"health_route\": \"/health\",\n",
+    "            \"env\": {\n",
+    "              \"MAX_BATCH_TOKENS\": \"16384\",\n",
+    "              \"MAX_CONCURRENT_REQUESTS\": \"512\",\n",
+    "              \"MODEL_ID\": \"/repository\"\n",
+    "            }\n",
+    "          }\n",
+    "        },\n",
+    "        \"repository\": \"jinaai/jina-embeddings-v2-base-en\",\n",
+    "        \"revision\": \"8705ed9657208b2d5220fffad1c3a30980d279d0\",\n",
+    "        \"task\": \"sentence-embeddings\",\n",
+    "    },\n",
+    "    \"name\": endpoint_name,\n",
+    "    \"provider\": {\n",
+    "        \"region\": \"us-east-1\",\n",
+    "        \"vendor\": \"aws\"\n",
+    "    },\n",
+    "    \"type\": \"protected\"\n",
+    "}\n",
+    "\n",
+    "response = requests.post(base_url, headers={**headers, 'accept': 'application/json'}, json=data)\n",
+    "\n",
+    "\n",
+    "print(response.status_code)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "96d173b2-8980-4554-9039-c62843d3fc7d",
+   "metadata": {},
+   "source": [
+    "## Wait until its running\n",
+    "Here we use `tqdm` as a pretty way of displaying our status. It took about ~30s for this model to get the Inference Endpoint running."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "b8aa66a9-3c8a-4040-9465-382c744f36cf",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a6f27d86f68b4000aa40e09ae079c6b0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Waiting for status to change: 0s [00:00, ?s/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Status is 'running'.\n"
+     ]
+    }
+   ],
+   "source": [
+    "with tqdm(desc=\"Waiting for status to change\", unit=\"s\") as pbar:\n",
+    "    while True:\n",
+    "        response_json = requests.get(endpoint_url, headers=headers).json()\n",
+    "        current_status = response_json['status']['state']\n",
+    "\n",
+    "        if current_status == 'running':\n",
+    "            print(\"Status is 'running'.\")\n",
+    "            break\n",
+    "\n",
+    "        pbar.set_description(f\"Status: {current_status}\")\n",
+    "        time.sleep(2)\n",
+    "        pbar.update(1)\n",
+    "\n",
+    "embedding_url = response_json['status']['url']"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "063fa066-e4d0-4a65-a82d-cf17db4af8d8",
+   "metadata": {},
+   "source": [
+    "I found that even though the status is running, I want to get a test message to run first before running our batch in parallel."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "66e00960-1d3d-490d-bedc-3eaf1924db76",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4e03e5a3d07a498ca6b3631605724b62",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Waiting for endpoint to accept requests: 0s [00:00, ?s/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Endpoint is accepting requests\n"
+     ]
+    }
+   ],
+   "source": [
+    "payload = {\"inputs\": \"This sound track was beautiful! It paints the senery in your mind so well I would recomend it even to people who hate vid. game music!\"}\n",
+    "\n",
+    "with tqdm(desc=\"Waiting for endpoint to accept requests\", unit=\"s\") as pbar:\n",
+    "    while True:\n",
+    "        try:\n",
+    "            response_json = requests.post(embedding_url, headers=headers, json=payload).json()\n",
+    "\n",
+    "            # Assuming the successful response has a specific structure\n",
+    "            if len(response_json[0]) == 768:\n",
+    "                print(\"Endpoint is accepting requests\")\n",
+    "                break\n",
+    "\n",
+    "        except requests.ConnectionError as e:\n",
+    "            pass\n",
+    "\n",
+    "        # Delay between retries\n",
+    "        time.sleep(5)\n",
+    "        pbar.update(1)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f7186126-ef6a-47d0-b158-112810649cd9",
+   "metadata": {},
+   "source": [
+    "# Get Embeddings"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1dadfd68-6d46-4ce8-a165-bfeb43b1f114",
+   "metadata": {},
+   "source": [
+    "Here I send a document, update it with the embedding, and return it. This happens in parallel with `MAX_WORKERS`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "ad3193fb-3def-42a8-968e-c63f2b864ca8",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "async def request(document, semaphore):\n",
+    "    # Semaphore guard\n",
+    "    async with semaphore:\n",
+    "        payload = {\n",
+    "            \"inputs\": document['content'] or document['title'] or '[deleted]',\n",
+    "            \"truncate\": True\n",
+    "        }\n",
+    "        \n",
+    "        timeout = ClientTimeout(total=10)  # Set a timeout for requests (10 seconds here)\n",
+    "\n",
+    "        async with ClientSession(timeout=timeout, headers=headers) as session:\n",
+    "            async with session.post(embedding_url, json=payload) as resp:\n",
+    "                if resp.status != 200:\n",
+    "                    raise RuntimeError(await resp.text())\n",
+    "                result = await resp.json()\n",
+    "                \n",
+    "        document['embedding'] = result[0]  # Assuming the API's output can be directly assigned\n",
+    "        return document\n",
+    "\n",
+    "async def main(documents):\n",
+    "    # Semaphore to limit concurrent requests. Adjust the number as needed.\n",
+    "    semaphore = asyncio.BoundedSemaphore(MAX_WORKERS)\n",
+    "\n",
+    "    # Creating a list of tasks\n",
+    "    tasks = [request(document, semaphore) for document in documents]\n",
+    "    \n",
+    "    # Using tqdm to show progress. It's been integrated into the async loop.\n",
+    "    for f in tqdm(asyncio.as_completed(tasks), total=len(documents)):\n",
+    "        await f"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "ec4983af-65eb-4841-808a-3738fb4d682d",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "cb73af52244e40d2aab8bdac3a55d443",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/9991 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Embeddings = 9991 documents = 9991\n",
+      "32 min 14.53 sec\n"
+     ]
+    }
+   ],
+   "source": [
+    "start = time.perf_counter()\n",
+    "\n",
+    "# Get embeddings\n",
+    "await main(documents)\n",
+    "\n",
+    "# Make sure we got it all\n",
+    "count = 0\n",
+    "for document in documents:\n",
+    "    if document['embedding'] and len(document['embedding']) == 768:\n",
+    "        count += 1\n",
+    "print(f'Embeddings = {count} documents = {len(documents)}')\n",
+    "\n",
+    "            \n",
+    "# Print elapsed time\n",
+    "elapsed_time = time.perf_counter() - start\n",
+    "minutes, seconds = divmod(elapsed_time, 60)\n",
+    "print(f\"{int(minutes)} min {seconds:.2f} sec\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bab97c7b-7bac-4bf5-9752-b528294dadc7",
+   "metadata": {},
+   "source": [
+    "## Pause Inference Endpoint\n",
+    "Now that we have finished, lets pause the endpoint so we don't incur any extra charges, this will also allow us to analyze the cost."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "540a0978-7670-4ce3-95c1-3823cc113b85",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "200\n",
+      "paused\n"
+     ]
+    }
+   ],
+   "source": [
+    "response = requests.post(endpoint_url + '/pause', headers=headers)\n",
+    "\n",
+    "print(response.status_code)\n",
+    "print(response.json()['status']['state'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "45ad65b7-3da2-4113-9b95-8fb4e21ae793",
+   "metadata": {},
+   "source": [
+    "# Push updated dataset to Hub\n",
+    "We now have our documents updated with the embeddings we wanted. First we need to convert it back to a `Dataset` format. I find its easiest to go from list of dicts -> `pd.DataFrame` -> `Dataset`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "9bb993f8-d624-4192-9626-8e9ed9888a1b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "df = pd.DataFrame(documents)\n",
+    "dd = DatasetDict({'train': Dataset.from_pandas(df)})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "f48e7c55-d5b7-4ed6-8516-272ae38716b1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "84a481e0cf74494cb2eb9d9857701212",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Pushing dataset shards to the dataset hub:   0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b8f128dfe7c546bcbc8f04817e3ca48c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Creating parquet from Arrow format:   0%|          | 0/10 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2dcc1d54036a49f1a1346a6be64e765a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Upload 1 LFS files:   0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "dd.push_to_hub(dataset_out, token=hf_token)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "41abea64-379d-49de-8d9a-355c2f4ce1ac",
+   "metadata": {},
+   "source": [
+    "# Analyze Usage\n",
+    "1. Go to your `dashboard_url` printed below\n",
+    "1. Click on the Usage & Cost tab\n",
+    "1. See how much you have spent"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "16815445-3079-43da-b14e-b54176a07a62",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "https://ui.endpoints.huggingface.co/HF-test-lab/endpoints/boru-jina-embeddings-demo\n"
+     ]
+    }
+   ],
+   "source": [
+    "dashboard_url = f'https://ui.endpoints.huggingface.co/{username}/endpoints/{endpoint_name}'\n",
+    "print(dashboard_url)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "81096c6f-d12f-4781-84ec-9066cfa465b3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "Hit enter to continue with the notebook \n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "''"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "input(\"Hit enter to continue with the notebook\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b953d5be-2494-4ff8-be42-9daf00c99c41",
+   "metadata": {},
+   "source": [
+    "# Delete Endpoint\n",
+    "We should see a `200` if everything went correctly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "c310c0f3-6f12-4d5c-838b-3a4c1f2e54ad",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "200\n"
+     ]
+    }
+   ],
+   "source": [
+    "response = requests.delete(endpoint_url, headers=headers)\n",
+    "\n",
+    "print(response.status_code)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5db1b1c3-16c3-403a-9472-a97e730826d5",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}