notebooks/04_get_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0f21cb1-fbc8-4282-b902-f47d92974df8",
   "metadata": {},
   "source": [
    "# Pre-requisites"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f625807-0707-4e2f-a0e0-8fbcdf08c865",
   "metadata": {},
   "source": [
    "## Why TEI\n",
    "There are 2 **unsung** challenges with RAG at scale:\n",
    "1. Getting the embeddings efficiently\n",
    "1. Efficient ingestion into the vector DB\n",
    "\n",
    "The issue with `1.` is that there are techniques but they are not widely *applied*. TEI solves a number of aspects:\n",
    "- Token Based Dynamic Batching\n",
    "- Using latest optimizations (Flash Attention, Candle and cuBLASLt)\n",
    "- Fast loading with safetensors\n",
    "\n",
    "The issue with `2.` is that it takes a bit of planning. We wont go much into that side of things here though."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3102abce-ea42-4da6-8c98-c6dd4edf7f0b",
   "metadata": {},
   "source": [
    "## Start TEI Locally\n",
    "Run [TEI](https://github.com/huggingface/text-embeddings-inference#docker), I have this running in a nvidia-docker container, but you can install as you like. Note that I ran this in a different terminal for monitoring and seperation. \n",
    "\n",
    "Note that as its running, its always going to pull the latest. Its at a very early stage at the time of writing. \n",
    "\n",
    "I chose [sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2](https://huggingface.co/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2) based on the STS ar-ar performance on [mteb/leaderboard](https://huggingface.co/spaces/mteb/leaderboard), its the top performer and half the size of second place! TEI is fast, but this will make our life easier for storage and retrieval.\n",
    "\n",
    "I use the `revision=refs/pr/8` because this has the pull request with [safetensors](https://github.com/huggingface/safetensors) which is required by TEI. Check out the [pull request](https://huggingface.co/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2/discussions/8) if you want to use a different embedding model and it doesnt have safetensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7e873652-8257-4aae-92bc-94e1bac54b73",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "# volume=$pwd/tei\n",
    "# model=sentence-transformers/paraphrase-multilingual-minilm-l12-v2\n",
    "# revision=refs/pr/8\n",
    "# docker run \\\n",
    "#     --gpus all \\\n",
    "#     -p 8080:80 \\\n",
    "#     -v $volume:/data \\\n",
    "#     -v /home/ec2-user/.cache/huggingface/token:/root/.cache/huggingface/token \\\n",
    "#     --pull always \\\n",
    "#     ghcr.io/huggingface/text-embeddings-inference:latest \\\n",
    "#     --model-id $model \\\n",
    "#     --revision $revision \\\n",
    "#     --pooling mean \\\n",
    "#     --max-batch-tokens 65536"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51959ef4-186e-4a32-826a-731813eaf593",
   "metadata": {},
   "source": [
    "### Test Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "52edfc97-5b6f-44f9-8d89-8578cf79fae9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "# response_code=$(curl -s -o /dev/null -w \"%{http_code}\" 127.0.0.1:8080/embed \\\n",
    "#     -X POST \\\n",
    "#     -d '{\"inputs\":\"What is Deep Learning?\"}' \\\n",
    "#     -H 'Content-Type: application/json')\n",
    "\n",
    "# if [ \"$response_code\" -eq 200 ]; then\n",
    "#     echo \"passed\"\n",
    "# else\n",
    "#     echo \"failed\"\n",
    "# fi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9d6b54a-02bd-49aa-b180-27a7ab90154e",
   "metadata": {},
   "source": [
    "## Start TEI with Inference Endpoints\n",
    "Another option is to run TEI on [Inference Endpoints](https://huggingface.co/inference-endpoints). Its cheap and fast. It took me less than 5 minutes to get it up and running!\n",
    "\n",
    "Check here for a [comprehensive guide](https://huggingface.co/blog/inference-endpoints-embeddings#3-deploy-embedding-model-as-inference-endpoint). Make sure to set these options **IN ORDER**:\n",
    "1. Model Repository = `transformers/paraphrase-multilingual-minilm-l12-v2`\n",
    "1. Name your endpoint\n",
    "1. Choose a GPU, I chose `Nvidia A10G` which is **$1.3/hr**.\n",
    "1. Advanced Configuration\n",
    "    1. Task = `Sentence Embeddings`\n",
    "    1. Revision (based on [this pull request for safetensors](https://huggingface.co/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2/discussions/8) = `a21e6630`\n",
    "    1. Container Type = `Text Embeddings Inference`\n",
    "    \n",
    "Set the other options as you prefer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec78c98a-6b7b-4689-8ef8-582c3fcdf66e",
   "metadata": {},
   "source": [
    "### Test Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a69e2ee1-67f2-4f0a-b496-02f5415a52ca",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "What is your API_URL? ········\n",
      "What is your BEARER TOKEN? Check your endpoint. ········\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "API_URL = getpass.getpass(prompt='What is your API_URL?')\n",
    "bearer_token = getpass.getpass(prompt='What is your BEARER TOKEN? Check your endpoint.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "949d6bf8-804f-496b-a59a-834483cc7073",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Constants\n",
    "HEADERS = {\n",
    "\t\"Authorization\": f\"Bearer {bearer_token}\",\n",
    "\t\"Content-Type\": \"application/json\"\n",
    "}\n",
    "MAX_WORKERS = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d00b4af1-8fbc-4f7a-8a78-e1c52dd77a66",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0047912598, -0.03164673, -0.018051147, -0.057739258, -0.04498291]...\n"
     ]
    }
   ],
   "source": [
    "import requests\n",
    "\n",
    "\n",
    "def query(payload):\n",
    "\tresponse = requests.post(API_URL, headers=HEADERS, json=payload)\n",
    "\treturn response.json()\n",
    "\t\n",
    "output = query({\n",
    "\t\"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",
    "print(f'{output[0][:5]}...')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1b28232-b65d-41ce-88de-fd70b93a528d",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "abb5186b-ee67-4e1e-882d-3d8d5b4575d4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import asyncio\n",
    "from pathlib import Path\n",
    "import json\n",
    "import time\n",
    "\n",
    "\n",
    "from aiohttp import ClientSession, ClientTimeout\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c4b82ea2-8b30-4c2e-99f0-9a30f2f1bfb7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/ec2-user/arabic-wiki\n"
     ]
    }
   ],
   "source": [
    "proj_dir = Path.cwd().parent\n",
    "print(proj_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76119e74-f601-436d-a253-63c5a19d1c83",
   "metadata": {},
   "source": [
    "# Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f6f74545-54a7-4f41-9f02-96964e1417f0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "files_in = list((proj_dir / 'data/processed/').glob('*.ndjson'))\n",
    "folder_out = proj_dir / 'data/embedded/'\n",
    "folder_out_str = str(folder_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e73235d-6274-4958-9e57-977afeeb5f1b",
   "metadata": {},
   "source": [
    "# Embed\n",
    "## Strategy\n",
    "TEI allows multiple concurrent requests, so its important that we dont waste the potential we have. I used the default `max-concurrent-requests` value of `512`, so I want to use that many `MAX_WORKERS`.\n",
    "\n",
    "Im using an `async` way of making requests that uses `aiohttp` as well as a nice progress bar. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf3da8cc-1651-4704-9091-39c2a1b835be",
   "metadata": {},
   "source": [
    "Note that Im using `'truncate':True` as even with our `350` word split earlier, there are always exceptions. Its important that as this scales we have as few issues as possible when embedding. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e455dd52-aad3-4313-8738-03141ee5152a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "async def request(document, semaphore):\n",
    "    # Semaphore guard\n",
    "    async with semaphore:\n",
    "        payload = {\n",
    "            \"inputs\": document['content'],\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(API_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(512)\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\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f0d17264-72dc-40be-aa46-17cde38c8189",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c4b7384336ad4c39a417a54a5a00a4ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b034dc636df440594550f56dc152c8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/243068 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 1: Embeddings = 243068 documents = 243068\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0203531009644b75abb22725a38b3ace",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/104065 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 2: Embeddings = 104065 documents = 104065\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a4c781089c42466ba380b0b598b2f9e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/123154 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 3: Embeddings = 123154 documents = 123154\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66a0feea106145a0aadeb64fab48b6f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/135965 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 4: Embeddings = 135965 documents = 135965\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6976832a78e48c5be335c5fef14bb5d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/99138 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 5: Embeddings = 99138 documents = 99138\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "764ceac837b040a39c2541074386e1f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/83678 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 6: Embeddings = 83678 documents = 83678\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d268907dd9844cd8f81f48f8568f576",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30573 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 7: Embeddings = 30573 documents = 30573\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "53bb70b332774c3d867cdb1cb3c48958",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/78957 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 8: Embeddings = 78957 documents = 78957\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "525ae3cf63af47b2acad508cb3c3efb7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/86327 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 9: Embeddings = 86327 documents = 86327\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a4c5103a1184ca1999b452c716131be",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/83111 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 10: Embeddings = 83111 documents = 83111\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9a06e2e21c6d4d12a04a55ca746594a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/92664 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 11: Embeddings = 92664 documents = 92664\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "933d457b2f4f4f1fa3d20b469dc22d75",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/66404 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 12: Embeddings = 66404 documents = 66404\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e4ce5ea591f431ca1ba6497ccf82b84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/62844 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 13: Embeddings = 62844 documents = 62844\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1d267ec29d864694b9f89fbf15e3e34a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/59349 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 14: Embeddings = 59349 documents = 59349\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfb1ecea3b2143c1916beb446201fe7f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/52554 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 15: Embeddings = 52554 documents = 52554\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdb83413a46e4ba984eb261994d05cd3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/34240 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 16: Embeddings = 34240 documents = 34240\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6cb8095952f4db3ab6d31219c21087e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/35933 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 17: Embeddings = 35933 documents = 35933\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0699bd10530c4a34aaaf9e88523ad5e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/64575 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 18: Embeddings = 64575 documents = 64575\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "736895c24eb84a8a9f514d99c628bdc7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/94244 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 19: Embeddings = 94244 documents = 94244\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "632c1bfc4370488ab977bedc8c31d404",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/124472 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 20: Embeddings = 124472 documents = 124472\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d2c550797b5f4444b91b954c3f3958b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/121849 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 21: Embeddings = 121849 documents = 121849\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8918701c109f4ecdbb1d73e5fe97d6b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/147110 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 22: Embeddings = 147110 documents = 147110\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c69c13c3a1354f5c900d268500ffcb00",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70322 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 23: Embeddings = 70322 documents = 70322\n",
      "104 min 32.33 sec\n"
     ]
    }
   ],
   "source": [
    "start = time.perf_counter()\n",
    "for i, file_in in tqdm(enumerate(files_in)):\n",
    "\n",
    "    with open(file_in, 'r') as f:\n",
    "        documents = [json.loads(line) for line in f]\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']) == 384:\n",
    "            count += 1\n",
    "    print(f'Batch {i+1}: Embeddings = {count} documents = {len(documents)}')\n",
    "\n",
    "    # Write to file\n",
    "    with open(folder_out/file_in.name, 'w', encoding='utf-8') as f:\n",
    "        for document in documents:\n",
    "            json_str = json.dumps(document, ensure_ascii=False)\n",
    "            f.write(json_str + '\\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": "3f0d9e6d-68f2-4086-9bcc-ffb27971fd63",
   "metadata": {},
   "source": [
    "Lets make sure that we still have all our documents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "abc6dccc-0e5c-45e2-a269-b9f02cff2d05",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/ec2-user/arabic-wiki/data/embedded\n",
      "2094596\n"
     ]
    }
   ],
   "source": [
    "!echo \"$folder_out_str\" && cat \"$folder_out_str\"/*.ndjson | wc -l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cdee2b1c-0493-4b3e-8ecb-9d79109c756e",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'content': 'عشاء كارين هو من سلسلة مطاعم استرالية يهدف عمداً عن تجربة تناول طعام غير سارَة ويتم توجيه الموظفين لإهانة العملاء طوال وجباتهم.\\nاقتبس اسم المطعم من المصطلح العامي على الإنترنت (كارين) والذي يستخدم لوصف امرأة بيضاء مسنة وقحة بشكل نمطي.\\nتاريخ المطعم.\\nتم إنشاء السلسلة في أستراليا (سيدني) في عام 2021 من قبل إيدين ليفن وجيمس فاريل. المطعم ذو طابع خاص يعتمد على خدمة تجربة طعام غير سارة حيث يدفع العملاء للموظفين لإهانتهم وكان من المفترض ان يكون المطعم مطعماً منبثقاً لمدة ستة أشهر في وورلد سكوير.\\nاثارت فكرة المطعم في البداية ردات فعل متغايرة مما أثار الخوف بشأن ما إذا كانت الإهانات المتبادلة من الممكن ان تعرض الموظفين لسوء المعاملة من قبل العملاء.\\nاسم (كارين) هو إشارة إلى الإسم المستخدم في الميمات (النكت التي تنشهر بسرعة في مواقع التواصل) لوصف امرأة بيضاء في منتصف العمر ووقحة بشكل نمطي.\\nيطلب من الموظفين ارتداء شخصية وقحة والسخرية من العملاء بشكل هزلي اثناء تناول وجباتهم ومن المتوقع ان يعيد العملاء هذا السلوك من خلال التصرف بوقاحة تجاه الموظفين ومع ذلك يُحظر على العملاء والموظفين استخدام الإهانات العنصرية أو التحيز الجنسي أو رهاب المثلية الجنسية.\\nتتضمن العديد من هذه التبادلات لغة نابية ويجب ان يكون برفقة الاشخاص اللذين يقلون عن 16 عاماََ بالغين.\\nكما يمكن لمالكي بطاقة هوية تظهر ان اسمهم كارين الحصول على مشروب مجاني.\\n',\n",
       " 'content_type': 'text',\n",
       " 'score': None,\n",
       " 'meta': {'id': '8974231',\n",
       "  'revid': '593870',\n",
       "  'url': 'https://ar.wikipedia.org/wiki?curid=8974231',\n",
       "  'title': 'مطعم عشاء كارين',\n",
       "  '_split_id': 0,\n",
       "  '_split_overlap': [{'doc_id': '288196225044b53e6ff86f2485257a0a',\n",
       "    'range': [790, 1225]}]},\n",
       " 'id_hash_keys': ['content'],\n",
       " 'embedding': [0.053985596,\n",
       "  -0.06933594,\n",
       "  -0.046417236,\n",
       "  0.07788086,\n",
       "  -0.06768799,\n",
       "  -0.01789856,\n",
       "  0.03265381,\n",
       "  -0.1060791,\n",
       "  0.046325684,\n",
       "  0.00022745132,\n",
       "  0.01524353,\n",
       "  0.021408081,\n",
       "  -0.039398193,\n",
       "  -0.076049805,\n",
       "  -0.058380127,\n",
       "  -0.027786255,\n",
       "  0.067993164,\n",
       "  -0.00894928,\n",
       "  0.11645508,\n",
       "  0.09039307,\n",
       "  0.032684326,\n",
       "  -0.023635864,\n",
       "  0.00970459,\n",
       "  0.08312988,\n",
       "  0.01638794,\n",
       "  -0.058380127,\n",
       "  0.0501709,\n",
       "  0.035583496,\n",
       "  -0.042297363,\n",
       "  -0.011688232,\n",
       "  0.010314941,\n",
       "  0.026626587,\n",
       "  0.039978027,\n",
       "  0.04095459,\n",
       "  0.039398193,\n",
       "  -0.029754639,\n",
       "  0.01360321,\n",
       "  -0.0015621185,\n",
       "  -0.009384155,\n",
       "  0.043151855,\n",
       "  0.09710693,\n",
       "  -0.0044937134,\n",
       "  0.0059928894,\n",
       "  -0.07293701,\n",
       "  -0.012710571,\n",
       "  0.013046265,\n",
       "  -0.08099365,\n",
       "  0.015457153,\n",
       "  -0.037963867,\n",
       "  -0.016601562,\n",
       "  -0.09564209,\n",
       "  0.038238525,\n",
       "  -0.020751953,\n",
       "  -0.016494751,\n",
       "  0.05529785,\n",
       "  -0.024490356,\n",
       "  0.011192322,\n",
       "  0.021087646,\n",
       "  -0.1274414,\n",
       "  0.04840088,\n",
       "  0.068603516,\n",
       "  -0.026535034,\n",
       "  0.010772705,\n",
       "  -0.03778076,\n",
       "  -0.046447754,\n",
       "  0.012870789,\n",
       "  -0.008674622,\n",
       "  0.062561035,\n",
       "  -0.06188965,\n",
       "  -0.05831909,\n",
       "  -0.05557251,\n",
       "  0.018417358,\n",
       "  -0.028656006,\n",
       "  0.015487671,\n",
       "  0.0146865845,\n",
       "  -0.010131836,\n",
       "  0.06652832,\n",
       "  -0.09710693,\n",
       "  -0.012542725,\n",
       "  0.04815674,\n",
       "  0.019470215,\n",
       "  0.0017337799,\n",
       "  -0.0069770813,\n",
       "  -0.027648926,\n",
       "  -0.103149414,\n",
       "  -0.019058228,\n",
       "  -0.049926758,\n",
       "  -0.05307007,\n",
       "  0.13562012,\n",
       "  0.016311646,\n",
       "  -0.10913086,\n",
       "  -0.04837036,\n",
       "  0.033996582,\n",
       "  -0.042144775,\n",
       "  0.043060303,\n",
       "  -0.10797119,\n",
       "  -0.087402344,\n",
       "  -0.051452637,\n",
       "  0.018203735,\n",
       "  0.15771484,\n",
       "  -0.025131226,\n",
       "  0.074401855,\n",
       "  -0.0033817291,\n",
       "  -0.026138306,\n",
       "  0.032165527,\n",
       "  -0.010025024,\n",
       "  -0.016189575,\n",
       "  -0.001121521,\n",
       "  -0.049560547,\n",
       "  -0.034057617,\n",
       "  -0.04336548,\n",
       "  0.11694336,\n",
       "  -0.035949707,\n",
       "  -0.061767578,\n",
       "  -0.101623535,\n",
       "  0.0051116943,\n",
       "  0.039245605,\n",
       "  0.0053977966,\n",
       "  0.016708374,\n",
       "  -0.012908936,\n",
       "  -0.022369385,\n",
       "  0.03475952,\n",
       "  0.035583496,\n",
       "  0.046417236,\n",
       "  -0.06097412,\n",
       "  0.07409668,\n",
       "  -0.07312012,\n",
       "  -0.019683838,\n",
       "  -0.07086182,\n",
       "  -0.03967285,\n",
       "  -0.016937256,\n",
       "  -0.02658081,\n",
       "  0.105651855,\n",
       "  0.05831909,\n",
       "  0.03778076,\n",
       "  -0.07543945,\n",
       "  0.055267334,\n",
       "  0.012268066,\n",
       "  0.07720947,\n",
       "  -0.068359375,\n",
       "  -0.011100769,\n",
       "  -0.0072898865,\n",
       "  0.051635742,\n",
       "  0.027954102,\n",
       "  0.043121338,\n",
       "  0.032440186,\n",
       "  0.081604004,\n",
       "  -0.058807373,\n",
       "  -0.048706055,\n",
       "  0.07867432,\n",
       "  0.014404297,\n",
       "  -0.0024490356,\n",
       "  -0.008575439,\n",
       "  0.010345459,\n",
       "  -0.013870239,\n",
       "  -0.019424438,\n",
       "  0.020889282,\n",
       "  0.013290405,\n",
       "  0.012817383,\n",
       "  0.015930176,\n",
       "  0.03768921,\n",
       "  -0.0012264252,\n",
       "  0.0010223389,\n",
       "  -0.023483276,\n",
       "  -0.005252838,\n",
       "  -0.051574707,\n",
       "  0.034729004,\n",
       "  -0.004081726,\n",
       "  0.0317688,\n",
       "  0.0087890625,\n",
       "  0.11047363,\n",
       "  0.05291748,\n",
       "  -0.12841797,\n",
       "  0.031799316,\n",
       "  -0.09881592,\n",
       "  0.07299805,\n",
       "  0.02859497,\n",
       "  0.024780273,\n",
       "  0.049316406,\n",
       "  -0.07122803,\n",
       "  0.03930664,\n",
       "  0.012702942,\n",
       "  0.06915283,\n",
       "  -0.03967285,\n",
       "  0.035949707,\n",
       "  -0.045166016,\n",
       "  0.09625244,\n",
       "  0.001203537,\n",
       "  0.022750854,\n",
       "  -0.03665161,\n",
       "  -0.013633728,\n",
       "  -0.018112183,\n",
       "  0.054107666,\n",
       "  -0.007106781,\n",
       "  0.004924774,\n",
       "  -0.014953613,\n",
       "  0.07147217,\n",
       "  -0.013092041,\n",
       "  0.016845703,\n",
       "  -0.002910614,\n",
       "  -0.05593872,\n",
       "  0.027404785,\n",
       "  0.002696991,\n",
       "  0.05822754,\n",
       "  0.0066566467,\n",
       "  -0.09729004,\n",
       "  0.040100098,\n",
       "  -0.00868988,\n",
       "  0.10290527,\n",
       "  0.04144287,\n",
       "  -0.012680054,\n",
       "  0.039215088,\n",
       "  -0.14074707,\n",
       "  0.08215332,\n",
       "  -0.05078125,\n",
       "  -0.028549194,\n",
       "  0.011962891,\n",
       "  0.028900146,\n",
       "  -0.02444458,\n",
       "  0.004207611,\n",
       "  -0.00995636,\n",
       "  0.028717041,\n",
       "  -0.08325195,\n",
       "  -0.047424316,\n",
       "  0.032043457,\n",
       "  -0.04675293,\n",
       "  -0.064575195,\n",
       "  -0.03857422,\n",
       "  0.0070266724,\n",
       "  -0.12634277,\n",
       "  -0.0803833,\n",
       "  -0.05419922,\n",
       "  0.064331055,\n",
       "  -0.004421234,\n",
       "  -0.00844574,\n",
       "  -0.05923462,\n",
       "  0.052490234,\n",
       "  0.032592773,\n",
       "  0.024230957,\n",
       "  0.075683594,\n",
       "  0.011390686,\n",
       "  0.013252258,\n",
       "  -0.029403687,\n",
       "  -0.03338623,\n",
       "  -0.045928955,\n",
       "  0.015022278,\n",
       "  -0.08343506,\n",
       "  0.060180664,\n",
       "  0.076171875,\n",
       "  0.058898926,\n",
       "  0.026184082,\n",
       "  -0.04031372,\n",
       "  -0.0847168,\n",
       "  -0.06628418,\n",
       "  -0.017974854,\n",
       "  -0.09967041,\n",
       "  0.07952881,\n",
       "  0.012413025,\n",
       "  0.04006958,\n",
       "  0.07788086,\n",
       "  0.014640808,\n",
       "  -0.029281616,\n",
       "  -0.04949951,\n",
       "  0.012565613,\n",
       "  0.025466919,\n",
       "  -0.018478394,\n",
       "  -0.072753906,\n",
       "  0.08905029,\n",
       "  0.023849487,\n",
       "  0.012306213,\n",
       "  -0.089538574,\n",
       "  -0.05657959,\n",
       "  0.11804199,\n",
       "  -0.035827637,\n",
       "  0.049194336,\n",
       "  -0.008911133,\n",
       "  0.09680176,\n",
       "  -0.03616333,\n",
       "  -0.057525635,\n",
       "  0.03375244,\n",
       "  -0.027435303,\n",
       "  0.0035476685,\n",
       "  0.010269165,\n",
       "  0.031921387,\n",
       "  0.0011024475,\n",
       "  0.045715332,\n",
       "  -0.026885986,\n",
       "  0.033935547,\n",
       "  0.06341553,\n",
       "  0.019958496,\n",
       "  0.008239746,\n",
       "  0.015174866,\n",
       "  -0.023071289,\n",
       "  0.0056762695,\n",
       "  0.064575195,\n",
       "  0.0042533875,\n",
       "  -0.05718994,\n",
       "  0.04486084,\n",
       "  0.020614624,\n",
       "  0.01461792,\n",
       "  -0.09283447,\n",
       "  0.019592285,\n",
       "  -0.022644043,\n",
       "  0.011512756,\n",
       "  -0.005874634,\n",
       "  -0.018569946,\n",
       "  0.006614685,\n",
       "  0.009269714,\n",
       "  -0.04296875,\n",
       "  -0.052856445,\n",
       "  0.084106445,\n",
       "  0.0043563843,\n",
       "  -0.020721436,\n",
       "  0.029022217,\n",
       "  0.03982544,\n",
       "  -0.109436035,\n",
       "  -0.036071777,\n",
       "  0.03253174,\n",
       "  0.011558533,\n",
       "  -0.10650635,\n",
       "  0.034454346,\n",
       "  -0.06951904,\n",
       "  -0.025817871,\n",
       "  0.10668945,\n",
       "  0.010101318,\n",
       "  -0.070739746,\n",
       "  0.049621582,\n",
       "  -0.09057617,\n",
       "  0.037231445,\n",
       "  -0.03152466,\n",
       "  -0.043914795,\n",
       "  0.07507324,\n",
       "  0.061645508,\n",
       "  0.0085372925,\n",
       "  0.004142761,\n",
       "  -0.051971436,\n",
       "  -0.05480957,\n",
       "  0.0030975342,\n",
       "  -0.046875,\n",
       "  -0.039398193,\n",
       "  0.08782959,\n",
       "  -0.012550354,\n",
       "  -0.003955841,\n",
       "  -0.07775879,\n",
       "  -0.021133423,\n",
       "  0.0062713623,\n",
       "  -0.02255249,\n",
       "  0.017868042,\n",
       "  0.049560547,\n",
       "  0.028121948,\n",
       "  0.031707764,\n",
       "  0.041168213,\n",
       "  0.009559631,\n",
       "  0.036956787,\n",
       "  0.008987427,\n",
       "  0.0024776459,\n",
       "  -0.003440857,\n",
       "  -0.0067749023,\n",
       "  -0.06439209,\n",
       "  -0.010902405,\n",
       "  -0.07104492,\n",
       "  0.006214142,\n",
       "  -0.06359863,\n",
       "  0.062316895,\n",
       "  0.005367279,\n",
       "  0.015197754,\n",
       "  -0.043182373,\n",
       "  0.050933838,\n",
       "  0.0035800934,\n",
       "  0.0032138824,\n",
       "  -0.017974854,\n",
       "  0.08972168,\n",
       "  0.011268616,\n",
       "  0.020477295,\n",
       "  -0.05050659,\n",
       "  -0.07232666,\n",
       "  0.07055664,\n",
       "  -0.010002136,\n",
       "  0.11480713,\n",
       "  0.02130127,\n",
       "  0.039093018,\n",
       "  0.009597778,\n",
       "  -0.0619812,\n",
       "  -0.016952515],\n",
       " 'id': '1af84f3b4cc6a9f1018f2f80b4fd3ba7'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "documents[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93d6ab01-bd3b-479d-918d-2bdb30b00fac",
   "metadata": {},
   "source": [
    "# Performance and Cost Analysis\n",
    "You can see that we are quite cost effective!\n",
    "![Cost](../media/arabic-rag-embeddings-cost.png)\n",
    "Note that the performance is over just the last 30 min window.\n",
    "Observations:\n",
    "- We have a througput of `~333/s`\n",
    "- Our median latency per request is `~50ms`\n",
    "![Metrics](../media/arabic-rag-embeddings-metrics.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc1e7cc5-b878-42bb-9fb4-e810f3f5006a",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Next Steps\n",
    "We need to import this into a vector db. "
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}