Adding updated notebooks

notebooks/01_get_data.ipynb

@@ -19,7 +19,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "ea8ae64c-f597-4c94-b93d-1b78060d7953",
    "metadata": {
     "tags": []
@@ -32,15 +32,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "id": "2f9527f9-4756-478b-99ac-a3c8c26ab63e",
    "metadata": {
     "tags": []
    },
    "outputs": [],
    "source": [
-    "proj_dir = str(Path.cwd().parent)\n",
-    "proj_dir\n",
+    "proj_dir_path = Path.cwd().parent\n",
+    "proj_dir = str(proj_dir_path)\n",
     "\n",
     "# So we can import later\n",
     "sys.path.append(proj_dir)"
@@ -56,12 +56,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "8bec29e3-8434-491f-914c-13f303dc68f3",
    "metadata": {
     "tags": []
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
    "source": [
     "%pip install -q -r \"$proj_dir\"/requirements.txt"
    ]
@@ -84,12 +92,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "id": "fe4b357f-88fe-44b5-9fce-354404b1447f",
    "metadata": {
     "tags": []
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Last-Modified: Sun, 01 Oct 2023 23:32:27 GMT\n"
+     ]
+    }
+   ],
    "source": [
     "!curl -I https://dumps.wikimedia.org/simplewiki/latest/simplewiki-latest-pages-articles-multistream.xml.bz2 --silent | grep \"Last-Modified\""
    ]
@@ -104,14 +120,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "id": "0f309c12-12de-4460-a03f-bd5b6fcc942c",
    "metadata": {
     "tags": []
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--2023-10-18 10:55:38--  https://dumps.wikimedia.org/simplewiki/latest/simplewiki-latest-pages-articles-multistream.xml.bz2\n",
+      "Resolving dumps.wikimedia.org (dumps.wikimedia.org)... 208.80.154.142, 2620:0:861:2:208:80:154:142\n",
+      "Connecting to dumps.wikimedia.org (dumps.wikimedia.org)|208.80.154.142|:443... connected.\n",
+      "HTTP request sent, awaiting response... 200 OK\n",
+      "Length: 286759308 (273M) [application/octet-stream]\n",
+      "Saving to: ‘/home/ec2-user/RAGDemo/data/raw/simplewiki-latest-pages-articles-multistream.xml.bz2’\n",
+      "\n",
+      "100%[======================================>] 286,759,308 4.22MB/s   in 66s    \n",
+      "\n",
+      "2023-10-18 10:56:45 (4.13 MB/s) - ‘/home/ec2-user/RAGDemo/data/raw/simplewiki-latest-pages-articles-multistream.xml.bz2’ saved [286759308/286759308]\n",
+      "\n"
+     ]
+    }
+   ],
    "source": [
-    "!wget -P \"$proj_dir\"/data/raw https://dumps.wikimedia.org/simplewiki/latest/simplewiki-latest-pages-articles-multistream.xml.bz2"
+    "!wget -nc -P \"$proj_dir\"/data/raw https://dumps.wikimedia.org/simplewiki/latest/simplewiki-latest-pages-articles-multistream.xml.bz2"
    ]
   },
   {
@@ -119,17 +153,39 @@
    "id": "46af5df6-5785-400a-986c-54a2c98768ea",
    "metadata": {},
    "source": [
-    "## Extract XML into jsonl"
+    "## Extract from XML\n",
+    "The download format from wikipedia is in XML. `wikiextractor` will convert this into a jsonl format split into many folders and files."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 9,
    "id": "c22dedcd-73b3-4aad-8eb7-1063954967ed",
-   "metadata": {},
-   "outputs": [],
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO: Preprocessing '/home/ec2-user/RAGDemo/data/raw/simplewiki-latest-pages-articles-multistream.xml.bz2' to collect template definitions: this may take some time.\n",
+      "INFO: Preprocessed 100000 pages\n",
+      "INFO: Preprocessed 200000 pages\n",
+      "INFO: Preprocessed 300000 pages\n",
+      "INFO: Preprocessed 400000 pages\n",
+      "INFO: Loaded 36594 templates in 54.1s\n",
+      "INFO: Starting page extraction from /home/ec2-user/RAGDemo/data/raw/simplewiki-latest-pages-articles-multistream.xml.bz2.\n",
+      "INFO: Using 3 extract processes.\n",
+      "INFO: Extracted 100000 articles (3481.4 art/s)\n",
+      "INFO: Extracted 200000 articles (3764.9 art/s)\n",
+      "INFO: Extracted 300000 articles (4175.8 art/s)\n",
+      "INFO: Finished 3-process extraction of 332024 articles in 86.9s (3822.7 art/s)\n"
+     ]
+    }
+   ],
    "source": [
-    "!wikiextractor -o \"$proj_dir\"/data/raw/output  --json simplewiki-latest-pages-articles-multistream.xml.bz2 "
+    "!wikiextractor -o \"$proj_dir\"/data/raw/output  --json \"$proj_dir\"/data/raw/simplewiki-latest-pages-articles-multistream.xml.bz2 "
    ]
   },
   {
@@ -139,12 +195,14 @@
    "source": [
     "## Consolidate into json\n",
     "\n",
-    "Some of this is boring, and most people dont care how you format it, just that its correct. Feel free to check out the consolidate file for more details."
+    "The split format is tedious to deal with, so now we we will consolidate this into 1 json file. This is fine since our data fits easily in RAM. But if it didnt, there are better options.\n",
+    "\n",
+    "Feel free to check out the [consolidate file](../src/preprocessing/consolidate.py) for more details."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 14,
    "id": "0a4ce3aa-9c1e-45e4-8219-a1714f482371",
    "metadata": {
     "tags": []
@@ -156,13 +214,39 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 17,
    "id": "3e93da6a-e304-450c-a81e-ffecaf0d8a9a",
-   "metadata": {},
-   "outputs": [],
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3f045c61ef544f34a1d6f7c4236b206c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Processing:   0%|          | 0/206 [00:00<?, ?file/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Wiki processed in 2.92 seconds!\n",
+      "Writing file!\n",
+      "File written in 3.08 seconds!\n"
+     ]
+    }
+   ],
    "source": [
-    "folder = proj_dir / 'data/raw/output'\n",
-    "file_out = proj_dir / 'data/consolidated/simple_wiki.json'\n",
+    "folder = proj_dir_path / 'data/raw/output'\n",
+    "file_out = proj_dir_path / 'data/consolidated/simple_wiki.json'\n",
     "folder_to_json(folder, file_out)"
    ]
   }
@@ -183,7 +267,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.10.9"
   }
  },
  "nbformat": 4,

notebooks/02_preprocessing.ipynb

@@ -37,7 +37,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "/Users/derekthomas/projects/spaces/RAGDemo\n"
+      "/home/ec2-user/RAGDemo\n"
      ]
     }
    ],
@@ -56,7 +56,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 13,
    "id": "f6f74545-54a7-4f41-9f02-96964e1417f0",
    "metadata": {
     "tags": []
@@ -64,7 +64,7 @@
    "outputs": [],
    "source": [
     "file_in = proj_dir / 'data/consolidated/simple_wiki.json'\n",
-    "file_out = proj_dir / 'data/processed/simple_wiki_processed.plk'"
+    "file_out = proj_dir / 'data/processed/simple_wiki_processed.pkl'"
    ]
   },
   {
@@ -96,7 +96,16 @@
    "metadata": {
     "tags": []
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[nltk_data] Downloading package punkt to /home/ec2-user/nltk_data...\n",
+      "[nltk_data]   Unzipping tokenizers/punkt.zip.\n"
+     ]
+    }
+   ],
    "source": [
     "pp = PreProcessor(clean_whitespace = True,\n",
     "             clean_header_footer = False,\n",
@@ -139,10 +148,10 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Preprocessing:   0%|█▏                                                                                                                                                                                                                                                        | 1510/332023 [00:01<04:12, 1308.89docs/s]We found one or more sentences whose word count is higher than the split length.\n",
-      "Preprocessing:  83%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍                                         | 276351/332023 [01:06<00:10, 5510.66docs/s]Document 81972e5bc1997b1ed4fb86d17f061a41 is 21206 characters long after preprocessing, where the maximum length should be 10000. Something might be wrong with the splitting, check the document affected to prevent issues at query time. This document will be now hard-split at 10000 chars recursively.\n",
+      "Preprocessing:   0%|▌                                                                                                                      | 1551/332023 [00:02<09:44, 565.82docs/s]We found one or more sentences whose word count is higher than the split length.\n",
+      "Preprocessing:  83%|████████████████████████████████████████████████████████████████████████████████████████████████▌                   | 276427/332023 [02:12<00:20, 2652.57docs/s]Document 81972e5bc1997b1ed4fb86d17f061a41 is 21206 characters long after preprocessing, where the maximum length should be 10000. Something might be wrong with the splitting, check the document affected to prevent issues at query time. This document will be now hard-split at 10000 chars recursively.\n",
       "Document 5e63e848e42966ddc747257fb7cf4092 is 11206 characters long after preprocessing, where the maximum length should be 10000. Something might be wrong with the splitting, check the document affected to prevent issues at query time. This document will be now hard-split at 10000 chars recursively.\n",
-      "Preprocessing: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 332023/332023 [01:15<00:00, 4403.36docs/s]\n"
+      "Preprocessing: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████��██████| 332023/332023 [02:29<00:00, 2219.16docs/s]\n"
      ]
     }
    ],
@@ -175,7 +184,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dbdb8eadba804b4485cc3e7f11a8b863",
+       "model_id": "ba764e7bf29f4202a74e08576a29f4e4",
        "version_major": 2,
        "version_minor": 0
       },
@@ -206,7 +215,9 @@
    "cell_type": "code",
    "execution_count": 8,
    "id": "02c1c6c8-6283-49a8-9d29-c355f1b08540",
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
    "outputs": [
     {
      "data": {
@@ -271,7 +282,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 11,
    "id": "5485cc27-3d3f-4b96-8884-accf5324da2d",
    "metadata": {
     "tags": []
@@ -340,7 +351,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.10.9"
   }
  },
  "nbformat": 4,

notebooks/03_get_embeddings.ipynb

@@ -0,0 +1,422 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a0f21cb1-fbc8-4282-b902-f47d92974df8",
+   "metadata": {},
+   "source": [
+    "# Pre-requisites"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3102abce-ea42-4da6-8c98-c6dd4edf7f0b",
+   "metadata": {},
+   "source": [
+    "## Start TEI\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. \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 the smaller [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) instead of the large. Its just as good on [mteb/leaderboard](https://huggingface.co/spaces/mteb/leaderboard) but its faster and smaller. TEI is fast, but this will make our life easier for storage and retrieval.\n",
+    "\n",
+    "I use the `revision=refs/pr/1` 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/BAAI/bge-base-en-v1.5/discussions/1) 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/data\n",
+    "# model=BAAI/bge-base-en-v1.5\n",
+    "# revision=refs/pr/1\n",
+    "# docker run \\\n",
+    "#     --gpus all \\\n",
+    "#     -p 8080:80 \\\n",
+    "#     -v $volume:/data \\\n",
+    "#     --pull always \\\n",
+    "#     ghcr.io/huggingface/text-embeddings-inference:latest \\\n",
+    "#     --model-id $model \\\n",
+    "#     --revision $revision \\\n",
+    "#     --pooling cls \\\n",
+    "#     --max-batch-tokens 65536"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "86a5ff83-1038-4880-8c90-dc3cab75cb49",
+   "metadata": {},
+   "source": [
+    "## Test Endpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "52edfc97-5b6f-44f9-8d89-8578cf79fae9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "passed\n"
+     ]
+    }
+   ],
+   "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": "b1b28232-b65d-41ce-88de-fd70b93a528d",
+   "metadata": {},
+   "source": [
+    "# Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "88408486-566a-4791-8ef2-5ee3e6941156",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from IPython.core.interactiveshell import InteractiveShell\n",
+    "InteractiveShell.ast_node_interactivity = 'all'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "abb5186b-ee67-4e1e-882d-3d8d5b4575d4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import asyncio\n",
+    "from pathlib import Path\n",
+    "import pickle\n",
+    "\n",
+    "import aiohttp\n",
+    "from tqdm.notebook import tqdm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "c4b82ea2-8b30-4c2e-99f0-9a30f2f1bfb7",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/home/ec2-user/RAGDemo\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": "markdown",
+   "id": "0d2bcda7-b245-45e3-a347-34166f217e1e",
+   "metadata": {},
+   "source": [
+    "I'm putting the documents in pickle files. The compression is nice, though its important to note pickles are known to be a security risk."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "f6f74545-54a7-4f41-9f02-96964e1417f0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "file_in = proj_dir / 'data/processed/simple_wiki_processed.pkl'\n",
+    "file_out = proj_dir / 'data/processed/simple_wiki_embeddings.pkl'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d2dd0df0-4274-45b3-9ee5-0205494e4d75",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Setup\n",
+    "Read in our list of dictionaries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "3c08e039-3686-4eca-9f87-7c469e3f19bc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 6.24 s, sys: 928 ms, total: 7.17 s\n",
+      "Wall time: 6.61 s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "with open(file_in, 'rb') as handle:\n",
+    "    documents = pickle.load(handle)\n",
+    "\n",
+    "documents = [document.to_dict() for document in documents]"
+   ]
+  },
+  {
+   "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": "code",
+   "execution_count": 8,
+   "id": "949d6bf8-804f-496b-a59a-834483cc7073",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Constants\n",
+    "ENDPOINT = \"http://127.0.0.1:8080/embed\"\n",
+    "HEADERS = {'Content-Type': 'application/json'}\n",
+    "MAX_WORKERS = 512"
+   ]
+  },
+  {
+   "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": "3353c849-a36c-4047-bb81-93dac6c49b68",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "async def fetch(session, url, document):\n",
+    "    payload = {\"inputs\": [document[\"content\"]], 'truncate':True}\n",
+    "    async with session.post(url, json=payload) as response:\n",
+    "        if response.status == 200:\n",
+    "            resp_json = await response.json()\n",
+    "            # Assuming the server's response contains an 'embedding' field\n",
+    "            document[\"embedding\"] = resp_json[0]\n",
+    "        else:\n",
+    "            print(f\"Error {response.status}: {await response.text()}\")\n",
+    "            # Handle error appropriately if needed\n",
+    "\n",
+    "async def main(documents):\n",
+    "    async with aiohttp.ClientSession(headers=HEADERS) as session:\n",
+    "        tasks = [fetch(session, ENDPOINT, doc) for doc in documents]\n",
+    "        await asyncio.gather(*tasks)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "f0d17264-72dc-40be-aa46-17cde38c8189",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f0ff772e915f4432971317e2150b60f2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Processing documents:   0%|          | 0/526 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# Create a list of async tasks\n",
+    "tasks = [main(documents[i:i+MAX_WORKERS]) for i in range(0, len(documents), MAX_WORKERS)]\n",
+    "\n",
+    "# Add a progress bar for visual feedback and run tasks\n",
+    "for task in tqdm(tasks, desc=\"Processing documents\"):\n",
+    "    await task"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f90a0ed7-b5e9-4ae4-9e87-4c04875ebcc9",
+   "metadata": {},
+   "source": [
+    "Lets double check that we got all the embeddings we expected!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "3950fa88-9961-4b33-9719-d5804509d4cf",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "268980"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "268980"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "count = 0\n",
+    "for document in documents:\n",
+    "    if len(document['embedding']) == 768:\n",
+    "        count += 1\n",
+    "count\n",
+    "len(documents)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b78bfa4-d365-4906-a71c-f444eabf6bf8",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "Great, we can see that they match.\n",
+    "\n",
+    "Let's write our embeddings to file"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "58d437a5-473f-4eae-9dbf-e8e6992754f6",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 5.68 s, sys: 640 ms, total: 6.32 s\n",
+      "Wall time: 14.1 s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "with open(file_out, 'wb') as handle:\n",
+    "    pickle.dump(documents, handle, protocol=pickle.HIGHEST_PROTOCOL)"
+   ]
+  },
+  {
+   "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.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}