Adding updated notebooks

notebooks/03_preprocessing.ipynb

@@ -0,0 +1,825 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b1b28232-b65d-41ce-88de-fd70b93a528d",
+   "metadata": {},
+   "source": [
+    "# Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "abb5186b-ee67-4e1e-882d-3d8d5b4575d4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from pathlib import Path\n",
+    "import pickle\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "from haystack.nodes.preprocessor import PreProcessor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "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": 3,
+   "id": "f6f74545-54a7-4f41-9f02-96964e1417f0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "files_in = list((proj_dir / 'data/consolidated').glob('*.ndjson'))\n",
+    "folder_out = proj_dir / 'data/processed'\n",
+    "folder_out_str = str(folder_out)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "509f41f9-a59f-4171-b61f-ae0cf756fc92",
+   "metadata": {},
+   "source": [
+    "# Analysis"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "f0cbd1c9-3105-4940-85dc-c01ccaa217c7",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "with open(files_in[0], 'r') as f:\n",
+    "    articles = [json.loads(line) for line in f]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "004aae7b-1a2f-4a0b-9450-5d80475258b1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'content': 'الماء مادةٌ شفافةٌ عديمة اللون والرائحة، وهو المكو...',\n",
+      " 'meta': {'id': '7',\n",
+      "          'revid': '2080427',\n",
+      "          'title': 'ماء',\n",
+      "          'url': 'https://ar.wikipedia.org/wiki?curid=7'}}\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pprint import pprint\n",
+    "article = articles[0].copy()\n",
+    "article['content'] = article['content'][:50] + '...'\n",
+    "pprint(article)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6a643cf2-abce-48a9-b4e0-478bcbee28c3",
+   "metadata": {},
+   "source": [
+    "# Preprocessing"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a8f9630e-447e-423e-9f6c-e1dbc654f2dd",
+   "metadata": {},
+   "source": [
+    "Its important to choose good pre-processing options. \n",
+    "\n",
+    "Clean whitespace helps each stage of RAG. It adds noise to the embeddings, and wastes space when we prompt with it.\n",
+    "\n",
+    "I chose to split by word as it would be tedious to tokenize here, and that doesnt scale well. The context length for most embedding models ends up being 512 tokens. We saw this within a good z-score is ~225 token.\n",
+    "\n",
+    "I like to respect the sentence boundary, thats why I gave a ~50 word buffer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "18807aea-24e4-4d74-bf10-55b24f3cb52c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "pp = PreProcessor(clean_whitespace = True,\n",
+    "             clean_header_footer = False,\n",
+    "             clean_empty_lines = True,\n",
+    "             remove_substrings = None,\n",
+    "             split_by='word',\n",
+    "             split_length = 225,\n",
+    "             split_overlap = 50,\n",
+    "             split_respect_sentence_boundary = True,\n",
+    "             tokenizer_model_folder = None,\n",
+    "             id_hash_keys = None,\n",
+    "             progress_bar = False,\n",
+    "             add_page_number = False,\n",
+    "             max_chars_check = 10_000)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3c1ab000-6574-485e-87f6-cc210f6e8a61",
+   "metadata": {},
+   "source": [
+    "When we break a wikipedia article up, we lose some of the context. The local context is somewhat preserved by the `split_overlap`. Im trying to preserve the global context by adding a prefix that has the article's title.\n",
+    "\n",
+    "You could enhance this with the summary as well. This is mostly to help the retrieval step of RAG. Note that the way Im doing it alters some of `haystack`'s features like the hash and the lengths, but those arent too necessary. \n",
+    "\n",
+    "A more advanced way for many business applications would be to summarize the document and add that as a prefix for sub-documents.\n",
+    "\n",
+    "One last thing to note, is that it would be prudent (in some use-cases) to preserve the original document without the summary to give to the reader (retrieve with the summary but prompt without), but since this is a demo use-case I wont be doing that."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "63871bdd-0369-4dd7-a65e-ccba29baed44",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(files_in[0], 'r', encoding='utf-8') as f:\n",
+    "    articles = [json.loads(line) for line in f]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "5c3b48b7-3c0f-41ba-a423-b716649efcaa",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "We found one or more sentences whose word count is higher than the split length.\n",
+      "Document e3e2bf8b3399979cb16219b175041b4d is 11336 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 91ad1d1a24e93abacabd5a5478a96977 is 14251 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 1625c431c0fcfaf81c13e0da59071a81 is 13395 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 790d3b2d94a68cbec6d77f3c15d0e679 is 13484 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 e2dcf80a1f9dfc118aed059255f9b90b is 13217 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"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 3min 31s, sys: 95.1 ms, total: 3min 31s\n",
+      "Wall time: 3min 31s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "documents = pp.process(articles)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "de6e1690-131a-41d1-a473-c908c2e40939",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Document 91ad1d1a24e93abacabd5a5478a96977 is 14251 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 e3e2bf8b3399979cb16219b175041b4d is 11336 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 1625c431c0fcfaf81c13e0da59071a81 is 13395 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 790d3b2d94a68cbec6d77f3c15d0e679 is 13484 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 e2dcf80a1f9dfc118aed059255f9b90b is 13217 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"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 6.86 s, sys: 1.31 s, total: 8.16 s\n",
+      "Wall time: 1min 33s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "import os\n",
+    "import concurrent.futures\n",
+    "\n",
+    "def parallel_preprocessing(articles):\n",
+    "    # Utility function to divide the articles into smaller chunks\n",
+    "    def chunkify(lst, n):\n",
+    "        \"\"\"Yield successive n-sized chunks from lst.\"\"\"\n",
+    "        for i in range(0, len(lst), n):\n",
+    "            yield lst[i:i + n]\n",
+    "\n",
+    "    # Size of each chunk. Adjust based on your needs.\n",
+    "    CHUNK_SIZE = 10_000  \n",
+    "    article_chunks = list(chunkify(articles, CHUNK_SIZE))\n",
+    "\n",
+    "    # Number of processes to run in parallel.\n",
+    "    # Use all available CPUs, but you can reduce the number if you wish to leave some CPUs free.\n",
+    "    NUM_PROCESSES = os.cpu_count()  \n",
+    "\n",
+    "    with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_PROCESSES) as executor:\n",
+    "        documents_list = list(executor.map(pp.process, article_chunks))\n",
+    "\n",
+    "    # Flatten the documents_list to get a single list of documents\n",
+    "    documents = [doc for sublist in documents_list for doc in sublist]\n",
+    "    return documents\n",
+    "\n",
+    "documents = parallel_preprocessing(articles)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "dab1658a-79a7-40f2-9a8c-1798e0d124bf",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a4d4ade8158144c6a06f072b550157c3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/23 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Document 91ad1d1a24e93abacabd5a5478a96977 is 14251 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 e3e2bf8b3399979cb16219b175041b4d is 11336 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 1625c431c0fcfaf81c13e0da59071a81 is 13395 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 790d3b2d94a68cbec6d77f3c15d0e679 is 13484 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 e2dcf80a1f9dfc118aed059255f9b90b is 13217 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"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bcb8fe6999ec408f99ef25c96c85c7b3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/243068 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Document 2693d73124e824fb45633e34189a6226 is 14375 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"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6ac231651a214181be49584511c97fe9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/104065 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c43b5ac1f3b84cac8f2715c5dcf21fb8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/123154 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a1b8201fc0ba44cf9ac938556c7effc8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/135965 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "58ce68b1b5044b0eb1c457495159cb6d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/99138 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b80a37f6986e478d85397f7295b67079",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/83678 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8197d71a396e47588699f22e44c843af",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/30573 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "39024c19258b41c4bcd7b6e5c0617d54",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/78957 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "aa26a27b522448b3bc2808f8ad650a28",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/86327 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3aae87d4a40d4e9e937646f02c969bb2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/83111 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "16038478b51a422d8ed7d597c9812999",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/92664 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0b9ea31fc970402f8368b4be64d91546",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/66404 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "710bf094621c4472bc6d3973619627df",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/62844 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0bae815b7b7448978fd042b6381b63ba",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/59349 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "483c99d5547a4bedb10802f9bbe414d3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/52554 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Document c55a744420239d9865b23f9ff7ab37cc is 23179 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 991fef35ead593c41f1dc73224ed7799 is 13179 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"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "34288f7c2437440ab07134116a6f5bf2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/34240 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c6b95ccc5799488980af6f73ab564c8b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/35933 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Document c4ae5d66606dcc644293946f8f5c7cab is 10612 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"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c6794c5150274e5cb2a09c52bfcf4c1a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/64575 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5aaf9358178d4daf9169a14cd71aa20f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/94244 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Document 1d1357a32c6cd115525ceb0590e6d67d is 11314 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 eca1c79de6ca898903e5db0a44a54803 is 24659 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 8a59ff0086cc7772046f54a065566ff9 is 14659 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"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "102f3b5a0c5e46f69742b9f41cc79928",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/124472 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3e4803f9fc3d405eac97a64489c0d7e0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/121849 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d2b4c133d207446c96b4e576aac15bf7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/147110 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1cbd39af41ef4b20a8bc1bea7fe8186f",
+       "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": [
+      "CPU times: user 2min 21s, sys: 20.1 s, total: 2min 41s\n",
+      "Wall time: 13min 36s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "for file_in in tqdm(files_in):\n",
+    "    # Load articles\n",
+    "    with open(file_in, 'r', encoding='utf-8') as f:\n",
+    "        articles = [json.loads(line) for line in f]\n",
+    "        \n",
+    "    # Preprocess articles\n",
+    "    documents = parallel_preprocessing(articles)\n",
+    "    \n",
+    "    # Prefix each document's content\n",
+    "    for document in tqdm(documents):\n",
+    "        if document.meta['_split_id'] != 0:\n",
+    "            document.content = f'عنوان: {document.meta[\"title\"]}. ' + document.content\n",
+    "            \n",
+    "    processed_articles = [document.to_dict() for document in documents]\n",
+    "    with open(folder_out/file_in.name, 'w', encoding='utf-8') as f:\n",
+    "        for article in processed_articles:\n",
+    "            json_str = json.dumps(article, ensure_ascii=False)\n",
+    "            f.write(json_str + '\\n')\n",
+    "        "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "72c1849c-1f4d-411f-b74b-6208b1e48217",
+   "metadata": {},
+   "source": [
+    "## Pre-processing Examples"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "02c1c6c8-6283-49a8-9d29-c355f1b08540",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<Document: {'content': 'عشاء كارين هو من سلسلة مطاعم استرالية يهدف عمداً عن تجربة تناول طعام غير سارَة ويتم توجيه الموظفين لإهانة العملاء طوال وجباتهم.\\nاقتبس اسم المطعم من المصطلح العامي على الإنترنت (كارين) والذي يستخدم لوصف امرأة بيضاء مسنة وقحة بشكل نمطي.\\nتاريخ المطعم.\\nتم إنشاء السلسلة في أستراليا (سيدني) في عام 2021 من قبل إيدين ليفن وجيمس فاريل. المطعم ذو طابع خاص يعتمد على خدمة تجربة طعام غير سارة حيث يدفع العملاء للموظفين لإهانتهم وكان من المفترض ان يكون المطعم مطعماً منبثقاً لمدة ستة أشهر في وورلد سكوير.\\nاثارت فكرة المطعم في البداية ردات فعل متغايرة مما أثار الخوف بشأن ما إذا كانت الإهانات المتبادلة من الممكن ان تعرض الموظفين لسوء المعاملة من قبل العملاء.\\nاسم (كارين) هو إشارة إلى الإسم المستخدم في الميمات (النكت التي تنشهر بسرعة في مواقع التواصل) لوصف امرأة بيضاء في منتصف العمر ووقحة بشكل نمطي.\\nيطلب من الموظفين ارتداء شخصية وقحة والسخرية من العملاء بشكل هزلي اثناء تناول وجباتهم ومن المتوقع ان يعيد العملاء هذا السلوك من خلال التصرف بوقاحة تجاه الموظفين ومع ذلك يُحظر على العملاء والموظفين استخدام الإهانات العنصرية أو التحيز الجنسي أو رهاب المثلية الجنسية.\\nتتضمن العديد من هذه التبادلات لغة نابية ويجب ان يكون برفقة الاشخاص اللذين يقلون عن 16 عاماََ بالغين.\\nكما يمكن لمالكي بطاقة هوية تظهر ان اسمهم كارين الحصول على مشروب مجاني.\\n', 'content_type': 'text', 'score': None, 'meta': {'id': '8974231', 'revid': '593870', 'url': 'https://ar.wikipedia.org/wiki?curid=8974231', 'title': 'مطعم عشاء كارين', '_split_id': 0, '_split_overlap': [{'doc_id': '288196225044b53e6ff86f2485257a0a', 'range': (790, 1225)}]}, 'id_hash_keys': ['content'], 'embedding': None, 'id': '1af84f3b4cc6a9f1018f2f80b4fd3ba7'}>"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "documents[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "b34890bf-9dba-459a-9b0d-aa4b5929cbe8",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<Document: {'content': 'عنوان: مطعم عشاء كارين. يطلب من الموظفين ارتداء شخصية وقحة والسخرية من العملاء بشكل هزلي اثناء تناول وجباتهم ومن المتوقع ان يعيد العملاء هذا السلوك من خلال التصرف بوقاحة تجاه الموظفين ومع ذلك يُحظر على العملاء والموظفين استخدام الإهانات العنصرية أو التحيز الجنسي أو رهاب المثلية الجنسية.\\nتتضمن العديد من هذه التبادلات لغة نابية ويجب ان يكون برفقة الاشخاص اللذين يقلون عن 16 عاماََ بالغين.\\nكما يمكن لمالكي بطاقة هوية تظهر ان اسمهم كارين الحصول على مشروب مجاني.\\nيرتكز المطعم علو وجبات العشاء الأمريكية في خمسينات القرن الماضي وتتميز القائمة بالهامبرغر وأجنحة الدجاج.\\nأصبح محتوى شائع لوسائل التواصل الإجتماعي خصوصاً على منصة (تيك توك) حيث نشر العملاء مقاطع فيديو لتفاعلاتهم مع الموظفين.\\nفتحت السلسلة في مواقع في المملكة المتحدة والولايات المتحدة ونيوزيلندا.\\nفي شهر أغسطس سنة 2022 اثار المطعم جدلاً بعد ان انتشر مقطع يُظهر فيه أحد موظفي فريق العمل في منطقة بريزبين يتصرف بشكل غير لائق على منصة تيك توك حيث القى تعليقات غير لائقة موجهة إلى زبونة قاصر ووالدها الذي كان يشاركها الطعام بإتهامه انه يمارس الرذيلة مع الأطفال، فقام المتحدث باسم السلسلة بالرد بإنهم اصيبو بخيبة أمل بسبب السلوك وأن الحادث يتعارض مع إرشاداتهم.', 'content_type': 'text', 'score': None, 'meta': {'id': '8974231', 'revid': '593870', 'url': 'https://ar.wikipedia.org/wiki?curid=8974231', 'title': 'مطعم عشاء كارين', '_split_id': 1, '_split_overlap': [{'doc_id': '1af84f3b4cc6a9f1018f2f80b4fd3ba7', 'range': (0, 435)}]}, 'id_hash_keys': ['content'], 'embedding': None, 'id': '288196225044b53e6ff86f2485257a0a'}>"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "documents[1]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "e6f50c27-a486-47e9-ba60-d567f5e530db",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<Document: {'content': 'نيكولاي فليفا والمعروف أيضًا باسم نيكو فليفا (1840 - 4 أغسطس عام 1920) هو سياسي وصحفي سياسي ومحامي أفلاقي ومن ثم روماني. اشتهر بانخراطه في المجريات السياسية ونزعته الوطنية الصريحة التي كادت تصل إلى حد الديماغوجية. اختبر كافة الصيغ السياسية التي يسمح بها نظام الحزبين في رومانيا. دام حضوره على الساحة العامة عقودًا من الزمن، شغل خلالها مقعدًا في جمعية النواب وتولى منصب عمدة مدينة بوخارست خلال الفترة الممتدة من عام 1884 حتى عام 1886.\\nباشر فليفا مسيرته السياسية مع الحزب الليبرالي الوطني الذي ساعد على تأسيسه وتمثيله أمام القضاء، ولكنه اتجه في ما بعد إلى معارضة احتكار الحزب للسلطة. حاول إنشاء حزب ثالث ودخل في مفاوضات من أجل اعتماد برامج سياسية مشتركة خاصة بقوى المعارضة المختلفة ومن بينها حزب المحافظين وجمعية جونيما في ظل الإدارات الليبرالية الوطنية المتعاقبة. ذاع صيته عندما تورط في فضيحتين كبيرتين خلال ثمانينيات القرن التاسع عشر حين أدى استهزائه بسلطة الحزب الليبرالي الوطني إلى اندلاع معارك في الشوارع ووقوع حادثتي إطلاق النار منفصلتين. اعتُبرت الجماعات الموالية لفليفا الصوت الرائد المعبر عن سخط الطبقة الوسطى وقتذاك، وشكلت إحدى التيارات التي دفعت باتجاه تبني حق الاقتراع العام للذكور.\\nعاد فليفا إلى المعسكر الليبرالي الوطني بعد منعه من تولي حقائب وزارية ريادية في الحكومات المحافظة، وأصبح وزيرًا للشؤون الداخلية خلال الفترة من عام 1895 حتى عام 1896. ', 'content_type': 'text', 'score': None, 'meta': {'id': '9044009', 'revid': '1673186', 'url': 'https://ar.wikipedia.org/wiki?curid=9044009', 'title': 'نيكولاي فليفا', '_split_id': 0, '_split_overlap': [{'doc_id': '188181b1026773d720383c7e7307b241', 'range': (943, 1257)}]}, 'id_hash_keys': ['content'], 'embedding': None, 'id': 'af5cda4722fa2a961bef66de8a6b3e17'}>"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "documents[10102]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "5485cc27-3d3f-4b96-8884-accf5324da2d",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2094596\n"
+     ]
+    }
+   ],
+   "source": [
+    "!cat \"$folder_out_str\"/*.ndjson | wc -l"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c5833dba-1bf6-48aa-be6f-0d70c71e54aa",
+   "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.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/04_get_embeddings.ipynb

@@ -0,0 +1,926 @@
+{
+ "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": null,
+   "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": null,
+   "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. Its cheap and fast. It took me less than 5 minutes to get it up and running!\n",
+    "\n",
+    "Check here for a [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\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": 1,
+   "id": "a69e2ee1-67f2-4f0a-b496-02f5415a52ca",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "What is your BEARER TOKEN? Check your endpoint. ········\n",
+      "What is your API_URL? ········\n"
+     ]
+    }
+   ],
+   "source": [
+    "import getpass\n",
+    "bearer_token = getpass.getpass(prompt='What is your BEARER TOKEN? Check your endpoint.')\n",
+    "API_URL = getpass.getpass(prompt='What is your API_URL?')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "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": 3,
+   "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": 4,
+   "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": 5,
+   "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": 6,
+   "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/'"
+   ]
+  },
+  {
+   "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": 7,
+   "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": 11,
+   "id": "f0d17264-72dc-40be-aa46-17cde38c8189",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1db8949409284a7cbeec2638ed197f59",
+       "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": "5945500ccf8649988918e2633269cb7b",
+       "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": "0cf8121a116f49fba72095fee46ef49d",
+       "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": "8f94983077854b5f9ab512f7d429eb55",
+       "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": "3d2932212e6b4323a377ff23758e7af7",
+       "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": "3de41d88c8bb439591925de045d8afe8",
+       "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": "604a4f3b1baf429687ac00aa63778cdf",
+       "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": "ecc69a2b763c4296b3a1fa35b15477aa",
+       "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": "bbdc03d2ca5a4099b412c1767b3d394c",
+       "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": "d4c64dfc612c4b7986de5385f5d88ba7",
+       "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": "d05e2884021143e0baf595a86725466a",
+       "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": "47f3537175d740aba9e2dc7de4c89fec",
+       "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": "02b89207ee8c407db4ad8045b3634243",
+       "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": "e1a96dfd66644007a39a5fef38e008ab",
+       "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": "3bef95b8dff044fa922e52e1e88b9813",
+       "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": "eb74103e549c4b0386d19f5d475af812",
+       "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": "41228a5bf8294c1e95f34b9376714543",
+       "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": "ba6eb3a975514d33ae3acac8859278d1",
+       "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": "67ecca34d2a1414c9e9817d835fe2083",
+       "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": "7fb632776adb4933b92c48f852b0ae6b",
+       "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": "17d73c88d4334357854f852c9783bfdb",
+       "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": "4dd7c3477a244d43b1d85417d4549eaa",
+       "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": "e8b657d57f584128ae5a7ee2ecf23c7f",
+       "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": "f3465378528a425e8dc9d040a003588a",
+       "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",
+      "6250.827601939993\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",
+    "print(time.perf_counter() - start)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cdee2b1c-0493-4b3e-8ecb-9d79109c756e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "documents[0]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f90a0ed7-b5e9-4ae4-9e87-4c04875ebcc9",
+   "metadata": {},
+   "source": [
+    "Lets double check that we got all the embeddings we expected!"
+   ]
+  },
+  {
+   "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": "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
+}

notebooks/05_vector_db.ipynb

@@ -0,0 +1,904 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "6a151ade-7d86-4a2e-bfe7-462089f4e04c",
+   "metadata": {},
+   "source": [
+    "# Approach\n",
+    "There are a number of aspects of choosing a vector db that might be unique to your situation. You should think through your HW, utilization, latency requirements, scale, etc before choosing. \n",
+    "\n",
+    "Im targeting a demo (low utilization, latency can be relaxed) that will live on a huggingface space. I have a small scale that could even fit in memory. I like [Qdrant](https://qdrant.tech) for this. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b1b28232-b65d-41ce-88de-fd70b93a528d",
+   "metadata": {},
+   "source": [
+    "# Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "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": 2,
+   "id": "abb5186b-ee67-4e1e-882d-3d8d5b4575d4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "import json\n",
+    "\n",
+    "from tqdm.notebook import tqdm\n",
+    "import lancedb"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "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": 4,
+   "id": "f6f74545-54a7-4f41-9f02-96964e1417f0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "files_in = list((proj_dir / 'data/embedded/').glob('*.ndjson'))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d2dd0df0-4274-45b3-9ee5-0205494e4d75",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Setup\n",
+    "Read in our list of dictionaries. This is the upper end for the machine Im using. This takes ~10GB of RAM. We could easily do this in batches of ~100k and be fine in most machines. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "3c08e039-3686-4eca-9f87-7c469e3f19bc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "with open(files_in[0], 'r') as f:\n",
+    "    first_line = f.readline().strip()  # read only the first line\n",
+    "    document = json.loads(first_line)\n",
+    "    document['vector'] = document.pop('embedding')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "98aec715-8d97-439e-99c0-0eff63df386b",
+   "metadata": {},
+   "source": [
+    "Convert the dictionaries to `Documents`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "4821e3c1-697d-4b69-bae3-300168755df9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'content': 'الماء مادةٌ شفافةٌ عديمة اللون والرائحة، وهو المكوّن الأساسي للجداول والبحيرات والبحار والمحيطات وكذلك للسوائل في جميع الكائنات الحيّة، وهو أكثر المركّبات الكيميائيّة انتشاراً على سطح الأرض. يتألّف جزيء الماء من ذرّة أكسجين مركزية ترتبط بها ذرّتا هيدروجين على طرفيها برابطة تساهميّة بحيث تكون صيغته الكيميائية H2O. عند الظروف القياسية من الضغط ودرجة الحرارة يكون الماء سائلاً؛ أمّا الحالة الصلبة فتتشكّل عند نقطة التجمّد، وتدعى بالجليد؛ أمّا الحالة الغازية فتتشكّل عند نقطة الغليان، وتسمّى بخار الماء.\\nإنّ الماء هو أساس وجود الحياة على كوكب الأرض، وهو يغطّي 71% من سطحها، وتمثّل مياه البحار والمحيطات أكبر نسبة للماء على الأرض، حيث تبلغ حوالي 96.5%. وتتوزّع النسب الباقية بين المياه الجوفيّة وبين جليد المناطق القطبيّة (1.7% لكليهما)، مع وجود نسبة صغيرة على شكل بخار ماء معلّق في الهواء على هيئة سحاب (غيوم)، وأحياناً أخرى على هيئة ضباب أو ندى، بالإضافة إلى الزخات المطريّة أو الثلجيّة. تبلغ نسبة الماء العذب حوالي 2.5% فقط من الماء الموجود على الأرض، ��أغلب هذه الكمّيّة (حوالي 99%) موجودة في الكتل الجليديّة في المناطق القطبيّة، في حين تتواجد 0.3% من الماء العذب في الأنهار والبحيرات وفي الغلاف الجوّي.\\nأما في الطبيعة، فتتغيّر حالة الماء بين الحالات الثلاثة للمادة على سطح الأرض باستمرار من خلال ما يعرف باسم الدورة المائيّة (أو دورة الماء)، والتي تتضمّن حدوث تبخّر ونتح (نتح تبخّري) ثم تكثيف فهطول ثم جريان لتصل إلى المصبّ في المسطّحات المائيّة.\\n',\n",
+       " 'content_type': 'text',\n",
+       " 'score': None,\n",
+       " 'meta': {'id': '7',\n",
+       "  'revid': '2080427',\n",
+       "  'url': 'https://ar.wikipedia.org/wiki?curid=7',\n",
+       "  'title': 'ماء',\n",
+       "  '_split_id': 0,\n",
+       "  '_split_overlap': [{'doc_id': '725ec671057ef790ad582509a8653584',\n",
+       "    'range': [887, 1347]}]},\n",
+       " 'id_hash_keys': ['content'],\n",
+       " 'id': '109a29bb227b1aaa5b784e972d8e1e3e',\n",
+       " 'vector': [-0.07318115,\n",
+       "  0.087646484,\n",
+       "  0.03274536,\n",
+       "  0.034942627,\n",
+       "  0.097961426,\n",
+       "  '...']}"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "doc = document.copy()\n",
+    "doc['vector'] = doc['vector'][:5] + ['...']\n",
+    "doc"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "676f644c-fb09-4d17-89ba-30c92aad8777",
+   "metadata": {},
+   "source": [
+    "Instantiate our `DocumentStore`. Note that Im saving this to disk, this is for portability which is good considering I want to move from this ec2 instance into a Hugging Face Space. \n",
+    "\n",
+    "Note that if you are doing this at scale, you should use a proper instance and not saving to file. You should also take a [measured ingestion](https://qdrant.tech/documentation/tutorials/bulk-upload/) approach instead of using a convenient loader. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "78033b87-8f68-4a44-899e-36fa8167cacf",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from lancedb.embeddings.registry import EmbeddingFunctionRegistry\n",
+    "from lancedb.embeddings.sentence_transformers import SentenceTransformerEmbeddings\n",
+    "\n",
+    "\n",
+    "db = lancedb.connect(proj_dir/\".lancedb\")\n",
+    "tbl = db.create_table('arabic-wiki', [document])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "21d5fa58-519e-4a23-9fc6-eed31e4723b5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "789efc342218412aa31d5a5a74b34c52",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Wiki Files:   0%|          | 0/23 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_1.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5bcf4a53ecaf471589d35588047969ea",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_1.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_2.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "763e45a7657e40288807556b1eb6c598",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_2.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_3.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6df7e0f7acfd409dbbaa04f9a21f6e18",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_3.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_4.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "87a7b93d85964254b7ae106795e55cbe",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_4.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_5.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b4c9eebb1be24e42885df1d4444c4856",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_5.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_6.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "feb88bc980534ca2bb7545784f9080df",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_6.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_7.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7ae8940509a94c17b209a783319af0b5",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_7.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_8.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "edbbeb596b034859856b525755331021",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_8.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_9.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1c0c7b32f0c14e77b3567dae6dea7c0f",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_9.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_10.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "193cc9f4f4ce4e22baf682bb31d77f13",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_10.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_11.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e5c2e514cae1462896f856e625d61c2b",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_11.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_12.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "818e7ebd889b4bbfa745904b5f132e48",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_12.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_13.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fe6268b0d85045da839a15f0dce577c9",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_13.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_14.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f67b03d9829545d8bfcbcdc55daac3a7",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_14.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_15.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "95bd941e1516437ca7e7cb59a55971ca",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_15.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_16.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "684b1247cf614817be09f34a2b349632",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_16.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_17.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d8c7f992f5064e5e83ec292a8e7c33bf",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_17.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_18.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5f7582025ad9462da77c0a57ba4a24f7",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_18.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_19.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e574e5ee70564cf78684eae0ca1ea6c4",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_19.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_20.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "21e50920d25047038bbad906e428ab5b",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_20.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_21.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "da30c5cd7cc14c47a739a1a2e410b012",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_21.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_22.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8efe4d8bf6f5437e88415819b0b88114",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_22.ndjson\n",
+      "Added documents\n",
+      "Reading documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_23.ndjson\n",
+      "Read documents\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7bda66b6033346559907a3bb56a95957",
+       "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": [
+      "Adding documents /home/ec2-user/arabic-wiki/data/embedded/ar_wiki_23.ndjson\n",
+      "Added documents\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                                                                                | 0/50 [00:00<?, ?it/s]/opt/conda/envs/arwiki/lib/python3.10/site-packages/lance/torch/kmeans.py:232: UserWarning: index_reduce() is in beta and the API may change at any time. (Triggered internally at ../aten/src/ATen/native/cuda/Indexing.cu:1047.)\n",
+      "  new_centroids.index_reduce_(\n",
+      " 34%|███████████████████████████████████                                                                    | 17/50 [01:26<02:47,  5.06s/it]\n",
+      "[2023-10-31T18:47:43Z WARN  lance_linalg::kmeans] KMeans: cluster 108 is empty\n",
+      "[2023-10-31T18:52:24Z WARN  lance_linalg::kmeans] KMeans: cluster 227 is empty\n",
+      "[2023-10-31T18:57:24Z WARN  lance_linalg::kmeans] KMeans: cluster 167 is empty\n",
+      "[2023-10-31T19:14:19Z WARN  lance_linalg::kmeans] KMeans: cluster 160 is empty\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 2h 2min 51s, sys: 1min 38s, total: 2h 4min 30s\n",
+      "Wall time: 42min 56s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "for file_in in tqdm(files_in, desc='Wiki Files: '):\n",
+    "\n",
+    "    tqdm.write(f\"Reading documents {str(file_in)}\")\n",
+    "    with open(file_in, 'r') as f:\n",
+    "        documents = [json.loads(line) for line in f]\n",
+    "    tqdm.write(f\"Read documents\")\n",
+    "\n",
+    "    for doc in tqdm(documents):\n",
+    "        if 'embedding' in doc:\n",
+    "            doc['vector'] = doc.pop('embedding')\n",
+    "    \n",
+    "    tqdm.write(f\"Adding documents {str(file_in)}\")\n",
+    "    tbl.add(documents)\n",
+    "    tqdm.write(f\"Added documents\")\n",
+    "tbl.create_index(\n",
+    "     num_partitions=1024,\n",
+    "     num_sub_vectors=384,\n",
+    "     accelerator=\"cuda\"\n",
+    ")\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "8ad72ca5-6ca3-43e3-bf2c-7461906576b9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from sentence_transformers import SentenceTransformer\n",
+    "\n",
+    "name=\"sentence-transformers/paraphrase-multilingual-minilm-l12-v2\"\n",
+    "model = SentenceTransformer(name)\n",
+    "\n",
+    "# used for both training and querying\n",
+    "def embed_func(batch):\n",
+    "    return [model.encode(sentence) for sentence in batch]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "41ab5a84-8984-4726-acd8-57ca0fce9e76",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['بكين',\n",
+       " 'كونمينغ',\n",
+       " 'نينغشيا',\n",
+       " 'تاي يوان',\n",
+       " 'تشنغتشو',\n",
+       " 'شانغهاي',\n",
+       " 'سنغافورة',\n",
+       " 'دلتا نهر يانغتسي',\n",
+       " 'تشانغتشون',\n",
+       " 'بكين']"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "query = \"What is the capital of China? I think it's Singapore.\"\n",
+    "query_vector = embed_func([query])[0]\n",
+    "[doc['meta']['title'] for doc in tbl.search(query_vector).limit(10).to_list()]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c0abad86-652a-4d7d-b118-21dc23a7a5c5",
+   "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.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}