add: notebook

semantic_search_using_word2vec_runpod_latest.ipynb

@@ -0,0 +1,2358 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "z9VPy5izbVuc"
+   },
+   "source": [
+    "### Semantic Search on arxive research papers using word2vec skipgram : https://arxiv.org/abs/1301.3781\n",
+    "- **Use Case:** Build a semantic search engine for research paper\n",
+    "- **Model Used:** Skipgram based word2vec for generating embeddings\n",
+    "- **Search Technique:** Cosine Similiarity, FAISS\n",
+    "- **Tokenizer Used:** sent_tokizer from nltk library"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "mKW7m--JbVue",
+    "outputId": "8464bdee-2f79-4696-9c12-957d17a0fdd7"
+   },
+   "outputs": [],
+   "source": [
+    "!pip install -q langdetect\n",
+    "!pip install -q gensim\n",
+    "!pip install -q nltk\n",
+    "!pip install -q spacy\n",
+    "!pip install -q kagglehub\n",
+    "!pip install -q gensim\n",
+    "!pip install -q pandas\n",
+    "!pip install -q numpy\n",
+    "!pip install -q tqdm\n",
+    "!pip install -q gdown\n",
+    "!pip install -q fastparquet\n",
+    "!pip install -q pyarrow\n",
+    "!pip install -q faiss-cpu"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloading...\n",
+      "From (original): https://drive.google.com/uc?id=1wTB4q0o2qREFNXcS2HZ3Y5tBlu97kF2o\n",
+      "From (redirected): https://drive.google.com/uc?id=1wTB4q0o2qREFNXcS2HZ3Y5tBlu97kF2o&confirm=t&uuid=13f5bd40-8f1d-469a-9b4e-4410b512530b\n",
+      "To: /workspace/data.pkl\n",
+      "100%|██████████████████████████████████████| 3.99G/3.99G [00:57<00:00, 69.8MB/s]\n",
+      "Downloading...\n",
+      "From: https://drive.google.com/uc?id=1lqo3Gyubefop6DlSnXEru65-sEnx2m9y\n",
+      "To: /workspace/word2vec_arxiv_skipgram.model\n",
+      "100%|██████████████████████████████████████| 11.3M/11.3M [00:00<00:00, 36.8MB/s]\n",
+      "Downloading...\n",
+      "From (original): https://drive.google.com/uc?id=1LrZaN2h5TVLC30Yz-VzYRsTQFNk7zSfC\n",
+      "From (redirected): https://drive.google.com/uc?id=1LrZaN2h5TVLC30Yz-VzYRsTQFNk7zSfC&confirm=t&uuid=3b8d4172-5e26-48ff-8593-0da1d51240f9\n",
+      "To: /workspace/word2vec_arxiv_skipgram.model.syn1neg.npy\n",
+      "100%|████████████████████████████████████████| 135M/135M [00:04<00:00, 27.8MB/s]\n",
+      "Downloading...\n",
+      "From (original): https://drive.google.com/uc?id=1QA_7_YwC6p2_7HhJ19-mqyIKRV-yfFiG\n",
+      "From (redirected): https://drive.google.com/uc?id=1QA_7_YwC6p2_7HhJ19-mqyIKRV-yfFiG&confirm=t&uuid=7186517e-df73-4662-85b3-84890b962fc2\n",
+      "To: /workspace/word2vec_arxiv_skipgram.model.wv.vectors.npy\n",
+      "100%|████████████████████████████████████████| 135M/135M [00:03<00:00, 35.4MB/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "!gdown 1wTB4q0o2qREFNXcS2HZ3Y5tBlu97kF2o\n",
+    "!gdown 1lqo3Gyubefop6DlSnXEru65-sEnx2m9y\n",
+    "!gdown 1LrZaN2h5TVLC30Yz-VzYRsTQFNk7zSfC\n",
+    "!gdown 1QA_7_YwC6p2_7HhJ19-mqyIKRV-yfFiG"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "V_EtGue3cMXF",
+    "outputId": "7914db9f-8050-429c-ebbe-1234080ae59c"
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[nltk_data] Downloading package punkt_tab to /root/nltk_data...\n",
+      "[nltk_data]   Unzipping tokenizers/punkt_tab.zip.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import nltk\n",
+    "nltk.download('punkt_tab')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "id": "0O2gP26JbVuf"
+   },
+   "outputs": [],
+   "source": [
+    "import spacy\n",
+    "import re\n",
+    "\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "import kagglehub\n",
+    "\n",
+    "from tqdm.notebook import tqdm\n",
+    "from gensim.models import Word2Vec\n",
+    "from langdetect import DetectorFactory, detect\n",
+    "from nltk.tokenize import sent_tokenize, word_tokenize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "id": "kXES0Z6ybVuf"
+   },
+   "outputs": [],
+   "source": [
+    "tqdm.pandas()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Building Word2Vec model for arxive dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "A0QAH1qubfoT",
+    "outputId": "76a9f028-3706-4dc5-e8bd-3bf24b69ac53"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Dataset downloaded at: /root/.cache/kagglehub/datasets/Cornell-University/arxiv/versions/259\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Download the dataset locally (returns the local path)\n",
+    "dataset_path = kagglehub.dataset_download(\"Cornell-University/arxiv\")\n",
+    "\n",
+    "print(\"✅ Dataset downloaded at:\", dataset_path)\n",
+    "\n",
+    "# Load using pandas with lines=True\n",
+    "file_path = f\"{dataset_path}/arxiv-metadata-oai-snapshot.json\"\n",
+    "\n",
+    "chunks = pd.read_json(file_path, lines=True, chunksize=10000)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chunks = pd.read_json(\n",
+    "    \"/kaggle/input/arxiv/arxiv-metadata-oai-snapshot.json\",\n",
+    "    lines=True,\n",
+    "    chunksize=10000  # tune this based on memory\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 695
+    },
+    "id": "LDMbjWV1bVug",
+    "outputId": "a75a4470-8f17-4d66-ced6-4b02ddce8df3"
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>id</th>\n",
+       "      <th>submitter</th>\n",
+       "      <th>authors</th>\n",
+       "      <th>title</th>\n",
+       "      <th>comments</th>\n",
+       "      <th>journal-ref</th>\n",
+       "      <th>doi</th>\n",
+       "      <th>report-no</th>\n",
+       "      <th>categories</th>\n",
+       "      <th>license</th>\n",
+       "      <th>abstract</th>\n",
+       "      <th>versions</th>\n",
+       "      <th>update_date</th>\n",
+       "      <th>authors_parsed</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>704.0001</td>\n",
+       "      <td>Pavel Nadolsky</td>\n",
+       "      <td>C. Bal\\'azs, E. L. Berger, P. M. Nadolsky, C.-...</td>\n",
+       "      <td>Calculation of prompt diphoton production cros...</td>\n",
+       "      <td>37 pages, 15 figures; published version</td>\n",
+       "      <td>Phys.Rev.D76:013009,2007</td>\n",
+       "      <td>10.1103/PhysRevD.76.013009</td>\n",
+       "      <td>ANL-HEP-PR-07-12</td>\n",
+       "      <td>hep-ph</td>\n",
+       "      <td>None</td>\n",
+       "      <td>A fully differential calculation in perturba...</td>\n",
+       "      <td>[{'version': 'v1', 'created': 'Mon, 2 Apr 2007...</td>\n",
+       "      <td>2008-11-26</td>\n",
+       "      <td>[[Balázs, C., ], [Berger, E. L., ], [Nadolsky,...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>704.0002</td>\n",
+       "      <td>Louis Theran</td>\n",
+       "      <td>Ileana Streinu and Louis Theran</td>\n",
+       "      <td>Sparsity-certifying Graph Decompositions</td>\n",
+       "      <td>To appear in Graphs and Combinatorics</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "      <td>math.CO cs.CG</td>\n",
+       "      <td>http://arxiv.org/licenses/nonexclusive-distrib...</td>\n",
+       "      <td>We describe a new algorithm, the $(k,\\ell)$-...</td>\n",
+       "      <td>[{'version': 'v1', 'created': 'Sat, 31 Mar 200...</td>\n",
+       "      <td>2008-12-13</td>\n",
+       "      <td>[[Streinu, Ileana, ], [Theran, Louis, ]]</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>704.0003</td>\n",
+       "      <td>Hongjun Pan</td>\n",
+       "      <td>Hongjun Pan</td>\n",
+       "      <td>The evolution of the Earth-Moon system based o...</td>\n",
+       "      <td>23 pages, 3 figures</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "      <td>physics.gen-ph</td>\n",
+       "      <td>None</td>\n",
+       "      <td>The evolution of Earth-Moon system is descri...</td>\n",
+       "      <td>[{'version': 'v1', 'created': 'Sun, 1 Apr 2007...</td>\n",
+       "      <td>2008-01-13</td>\n",
+       "      <td>[[Pan, Hongjun, ]]</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>704.0004</td>\n",
+       "      <td>David Callan</td>\n",
+       "      <td>David Callan</td>\n",
+       "      <td>A determinant of Stirling cycle numbers counts...</td>\n",
+       "      <td>11 pages</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "      <td>math.CO</td>\n",
+       "      <td>None</td>\n",
+       "      <td>We show that a determinant of Stirling cycle...</td>\n",
+       "      <td>[{'version': 'v1', 'created': 'Sat, 31 Mar 200...</td>\n",
+       "      <td>2007-05-23</td>\n",
+       "      <td>[[Callan, David, ]]</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>704.0005</td>\n",
+       "      <td>Alberto Torchinsky</td>\n",
+       "      <td>Wael Abu-Shammala and Alberto Torchinsky</td>\n",
+       "      <td>From dyadic $\\Lambda_{\\alpha}$ to $\\Lambda_{\\a...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>Illinois J. Math. 52 (2008) no.2, 681-689</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "      <td>math.CA math.FA</td>\n",
+       "      <td>None</td>\n",
+       "      <td>In this paper we show how to compute the $\\L...</td>\n",
+       "      <td>[{'version': 'v1', 'created': 'Mon, 2 Apr 2007...</td>\n",
+       "      <td>2013-10-15</td>\n",
+       "      <td>[[Abu-Shammala, Wael, ], [Torchinsky, Alberto, ]]</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         id           submitter  \\\n",
+       "0  704.0001      Pavel Nadolsky   \n",
+       "1  704.0002        Louis Theran   \n",
+       "2  704.0003         Hongjun Pan   \n",
+       "3  704.0004        David Callan   \n",
+       "4  704.0005  Alberto Torchinsky   \n",
+       "\n",
+       "                                             authors  \\\n",
+       "0  C. Bal\\'azs, E. L. Berger, P. M. Nadolsky, C.-...   \n",
+       "1                    Ileana Streinu and Louis Theran   \n",
+       "2                                        Hongjun Pan   \n",
+       "3                                       David Callan   \n",
+       "4           Wael Abu-Shammala and Alberto Torchinsky   \n",
+       "\n",
+       "                                               title  \\\n",
+       "0  Calculation of prompt diphoton production cros...   \n",
+       "1           Sparsity-certifying Graph Decompositions   \n",
+       "2  The evolution of the Earth-Moon system based o...   \n",
+       "3  A determinant of Stirling cycle numbers counts...   \n",
+       "4  From dyadic $\\Lambda_{\\alpha}$ to $\\Lambda_{\\a...   \n",
+       "\n",
+       "                                  comments  \\\n",
+       "0  37 pages, 15 figures; published version   \n",
+       "1    To appear in Graphs and Combinatorics   \n",
+       "2                      23 pages, 3 figures   \n",
+       "3                                 11 pages   \n",
+       "4                                     None   \n",
+       "\n",
+       "                                 journal-ref                         doi  \\\n",
+       "0                   Phys.Rev.D76:013009,2007  10.1103/PhysRevD.76.013009   \n",
+       "1                                       None                        None   \n",
+       "2                                       None                        None   \n",
+       "3                                       None                        None   \n",
+       "4  Illinois J. Math. 52 (2008) no.2, 681-689                        None   \n",
+       "\n",
+       "          report-no       categories  \\\n",
+       "0  ANL-HEP-PR-07-12           hep-ph   \n",
+       "1              None    math.CO cs.CG   \n",
+       "2              None   physics.gen-ph   \n",
+       "3              None          math.CO   \n",
+       "4              None  math.CA math.FA   \n",
+       "\n",
+       "                                             license  \\\n",
+       "0                                               None   \n",
+       "1  http://arxiv.org/licenses/nonexclusive-distrib...   \n",
+       "2                                               None   \n",
+       "3                                               None   \n",
+       "4                                               None   \n",
+       "\n",
+       "                                            abstract  \\\n",
+       "0    A fully differential calculation in perturba...   \n",
+       "1    We describe a new algorithm, the $(k,\\ell)$-...   \n",
+       "2    The evolution of Earth-Moon system is descri...   \n",
+       "3    We show that a determinant of Stirling cycle...   \n",
+       "4    In this paper we show how to compute the $\\L...   \n",
+       "\n",
+       "                                            versions update_date  \\\n",
+       "0  [{'version': 'v1', 'created': 'Mon, 2 Apr 2007...  2008-11-26   \n",
+       "1  [{'version': 'v1', 'created': 'Sat, 31 Mar 200...  2008-12-13   \n",
+       "2  [{'version': 'v1', 'created': 'Sun, 1 Apr 2007...  2008-01-13   \n",
+       "3  [{'version': 'v1', 'created': 'Sat, 31 Mar 200...  2007-05-23   \n",
+       "4  [{'version': 'v1', 'created': 'Mon, 2 Apr 2007...  2013-10-15   \n",
+       "\n",
+       "                                      authors_parsed  \n",
+       "0  [[Balázs, C., ], [Berger, E. L., ], [Nadolsky,...  \n",
+       "1           [[Streinu, Ileana, ], [Theran, Louis, ]]  \n",
+       "2                                 [[Pan, Hongjun, ]]  \n",
+       "3                                [[Callan, David, ]]  \n",
+       "4  [[Abu-Shammala, Wael, ], [Torchinsky, Alberto, ]]  "
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "next(chunks).head()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "tiIgyjkK4q3t"
+   },
+   "source": [
+    "#### Using subset of the given dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "rsA5QxtHbVug",
+    "outputId": "3b7820bc-f46e-4328-bd62-ee04fda760bc"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Processed 0 records\n",
+      "Processed 100000 records\n",
+      "Processed 200000 records\n",
+      "Processed 300000 records\n",
+      "Processed 400000 records\n",
+      "Processed 500000 records\n",
+      "Processed 600000 records\n",
+      "Processed 700000 records\n",
+      "Processed 800000 records\n",
+      "Processed 900000 records\n",
+      "Processed 1000000 records\n",
+      "(1010000, 4)\n",
+      "         id                                              title  \\\n",
+      "0  706.1314               When Did Cosmic Acceleration Start ?   \n",
+      "1  706.1315    The Dirac system on the Anti-de Sitter Universe   \n",
+      "2  706.1316  Coupling of Optical Lumped Nanocircuit Element...   \n",
+      "3  706.1317  A model for learning to segment temporal seque...   \n",
+      "4  706.1318  Constructing a maximum utility slate of on-lin...   \n",
+      "\n",
+      "                                            abstract               categories  \n",
+      "0    A precise determination, and comparison, of ...    astro-ph gr-qc hep-ph  \n",
+      "1    We investigate the global solutions of the D...  math-ph math.AP math.MP  \n",
+      "2    We present here a model for the coupling amo...        cond-mat.mtrl-sci  \n",
+      "3    This paper proposes a novel learning method ...                  nlin.AO  \n",
+      "4    We present an algorithm for constructing an ...              cs.DM cs.DS  \n"
+     ]
+    }
+   ],
+   "source": [
+    "data = []\n",
+    "for i, chunk in enumerate(chunks):\n",
+    "    # Example: keep only title and abstract\n",
+    "    subset = chunk[[\"id\", \"title\", \"abstract\", \"categories\"]] # taking only abstract for generating word embeddings\n",
+    "    data.append(subset)\n",
+    "\n",
+    "    if i % 10 == 0:\n",
+    "        print(f\"Processed {i*10000} records\")\n",
+    "\n",
+    "    # Stop early if you only need part\n",
+    "    if i == 100:  # ~1 million records\n",
+    "        break\n",
+    "\n",
+    "df = pd.concat(data, ignore_index=True)\n",
+    "print(df.shape)\n",
+    "print(df.head())\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 458
+    },
+    "id": "g1jwkMB_bVuh",
+    "outputId": "ba0c426a-3829-4d05-c28f-f1573672c657"
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>id</th>\n",
+       "      <th>title</th>\n",
+       "      <th>abstract</th>\n",
+       "      <th>categories</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>706.13140</td>\n",
+       "      <td>When Did Cosmic Acceleration Start ?</td>\n",
+       "      <td>A precise determination, and comparison, of ...</td>\n",
+       "      <td>astro-ph gr-qc hep-ph</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>706.13150</td>\n",
+       "      <td>The Dirac system on the Anti-de Sitter Universe</td>\n",
+       "      <td>We investigate the global solutions of the D...</td>\n",
+       "      <td>math-ph math.AP math.MP</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>706.13160</td>\n",
+       "      <td>Coupling of Optical Lumped Nanocircuit Element...</td>\n",
+       "      <td>We present here a model for the coupling amo...</td>\n",
+       "      <td>cond-mat.mtrl-sci</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>706.13170</td>\n",
+       "      <td>A model for learning to segment temporal seque...</td>\n",
+       "      <td>This paper proposes a novel learning method ...</td>\n",
+       "      <td>nlin.AO</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>706.13180</td>\n",
+       "      <td>Constructing a maximum utility slate of on-lin...</td>\n",
+       "      <td>We present an algorithm for constructing an ...</td>\n",
+       "      <td>cs.DM cs.DS</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1009995</th>\n",
+       "      <td>1808.10849</td>\n",
+       "      <td>On sets defining few ordinary hyperplanes</td>\n",
+       "      <td>Let $P$ be a set of $n$ points in real proje...</td>\n",
+       "      <td>math.CO math.AG</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1009996</th>\n",
+       "      <td>1808.10850</td>\n",
+       "      <td>Quantum walks in external gauge fields</td>\n",
+       "      <td>Describing a particle in an external electro...</td>\n",
+       "      <td>quant-ph cond-mat.other math-ph math.MP</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1009997</th>\n",
+       "      <td>1808.10851</td>\n",
+       "      <td>Solving the muon g-2 anomaly in CMSSM extensio...</td>\n",
+       "      <td>We propose to generate non-universal gaugino...</td>\n",
+       "      <td>hep-ph hep-ex</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1009998</th>\n",
+       "      <td>1808.10852</td>\n",
+       "      <td>Towards Asynchronous Motor Imagery-Based Brain...</td>\n",
+       "      <td>In this paper, the deep learning (DL) approa...</td>\n",
+       "      <td>eess.SP cs.HC q-bio.NC</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1009999</th>\n",
+       "      <td>1808.10853</td>\n",
+       "      <td>Charged impurity scattering in two-dimensional...</td>\n",
+       "      <td>The singular density of states and the two F...</td>\n",
+       "      <td>cond-mat.mes-hall</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>1010000 rows × 4 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                 id                                              title  \\\n",
+       "0         706.13140               When Did Cosmic Acceleration Start ?   \n",
+       "1         706.13150    The Dirac system on the Anti-de Sitter Universe   \n",
+       "2         706.13160  Coupling of Optical Lumped Nanocircuit Element...   \n",
+       "3         706.13170  A model for learning to segment temporal seque...   \n",
+       "4         706.13180  Constructing a maximum utility slate of on-lin...   \n",
+       "...             ...                                                ...   \n",
+       "1009995  1808.10849          On sets defining few ordinary hyperplanes   \n",
+       "1009996  1808.10850             Quantum walks in external gauge fields   \n",
+       "1009997  1808.10851  Solving the muon g-2 anomaly in CMSSM extensio...   \n",
+       "1009998  1808.10852  Towards Asynchronous Motor Imagery-Based Brain...   \n",
+       "1009999  1808.10853  Charged impurity scattering in two-dimensional...   \n",
+       "\n",
+       "                                                  abstract  \\\n",
+       "0          A precise determination, and comparison, of ...   \n",
+       "1          We investigate the global solutions of the D...   \n",
+       "2          We present here a model for the coupling amo...   \n",
+       "3          This paper proposes a novel learning method ...   \n",
+       "4          We present an algorithm for constructing an ...   \n",
+       "...                                                    ...   \n",
+       "1009995    Let $P$ be a set of $n$ points in real proje...   \n",
+       "1009996    Describing a particle in an external electro...   \n",
+       "1009997    We propose to generate non-universal gaugino...   \n",
+       "1009998    In this paper, the deep learning (DL) approa...   \n",
+       "1009999    The singular density of states and the two F...   \n",
+       "\n",
+       "                                      categories  \n",
+       "0                          astro-ph gr-qc hep-ph  \n",
+       "1                        math-ph math.AP math.MP  \n",
+       "2                              cond-mat.mtrl-sci  \n",
+       "3                                        nlin.AO  \n",
+       "4                                    cs.DM cs.DS  \n",
+       "...                                          ...  \n",
+       "1009995                          math.CO math.AG  \n",
+       "1009996  quant-ph cond-mat.other math-ph math.MP  \n",
+       "1009997                            hep-ph hep-ex  \n",
+       "1009998                   eess.SP cs.HC q-bio.NC  \n",
+       "1009999                        cond-mat.mes-hall  \n",
+       "\n",
+       "[1010000 rows x 4 columns]"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df.dropna()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "n2J6AYWQbVuh"
+   },
+   "source": [
+    "#### Cleaning abstracts"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 325,
+     "referenced_widgets": [
+      "7d5629fd92c543bcaafe7fd352c079c6",
+      "a815fb290af249cbb23679e34d1fdaa3",
+      "b042bdc131654d25bac957a4e921f59d",
+      "5510fb4117fd4954aeeb0de2214b1ed4",
+      "4de066fc521f493cab302c8d5ca255e2",
+      "06d6872e3d1f4086a4f16b8caababccb",
+      "88ba2081fc004bbdb6f438c5031458d5",
+      "52f12329e31645888592d0120a05f6ea",
+      "214288b8cb6c453084c3cce01982eff7",
+      "4df933c8001c4a27b351c5b66df23eb5",
+      "80071f18a3fd4ddc865a8a536d83cd2e"
+     ]
+    },
+    "id": "6Lx5uI74bVui",
+    "outputId": "97a945eb-38a5-4e23-94ea-4e2bb5e17d6d"
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0c66c831aedf446ea20c323688c43d4b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1010000 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>id</th>\n",
+       "      <th>title</th>\n",
+       "      <th>abstract</th>\n",
+       "      <th>categories</th>\n",
+       "      <th>processed_abstract</th>\n",
+       "      <th>tokenized_abstract</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>706.1314</td>\n",
+       "      <td>When Did Cosmic Acceleration Start ?</td>\n",
+       "      <td>A precise determination, and comparison, of ...</td>\n",
+       "      <td>astro-ph gr-qc hep-ph</td>\n",
+       "      <td>a precise determination and comparison of the ...</td>\n",
+       "      <td>[[a, precise, determination, ,, and, compariso...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>706.1315</td>\n",
+       "      <td>The Dirac system on the Anti-de Sitter Universe</td>\n",
+       "      <td>We investigate the global solutions of the D...</td>\n",
+       "      <td>math-ph math.AP math.MP</td>\n",
+       "      <td>we investigate the global solutions of the dir...</td>\n",
+       "      <td>[[we, investigate, the, global, solutions, of,...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>706.1316</td>\n",
+       "      <td>Coupling of Optical Lumped Nanocircuit Element...</td>\n",
+       "      <td>We present here a model for the coupling amo...</td>\n",
+       "      <td>cond-mat.mtrl-sci</td>\n",
+       "      <td>we present here a model for the coupling among...</td>\n",
+       "      <td>[[we, present, here, a, model, for, the, coupl...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>706.1317</td>\n",
+       "      <td>A model for learning to segment temporal seque...</td>\n",
+       "      <td>This paper proposes a novel learning method ...</td>\n",
+       "      <td>nlin.AO</td>\n",
+       "      <td>this paper proposes a novel learning method fo...</td>\n",
+       "      <td>[[this, paper, proposes, a, novel, learning, m...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>706.1318</td>\n",
+       "      <td>Constructing a maximum utility slate of on-lin...</td>\n",
+       "      <td>We present an algorithm for constructing an ...</td>\n",
+       "      <td>cs.DM cs.DS</td>\n",
+       "      <td>we present an algorithm for constructing an op...</td>\n",
+       "      <td>[[we, present, an, algorithm, for, constructin...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         id                                              title  \\\n",
+       "0  706.1314               When Did Cosmic Acceleration Start ?   \n",
+       "1  706.1315    The Dirac system on the Anti-de Sitter Universe   \n",
+       "2  706.1316  Coupling of Optical Lumped Nanocircuit Element...   \n",
+       "3  706.1317  A model for learning to segment temporal seque...   \n",
+       "4  706.1318  Constructing a maximum utility slate of on-lin...   \n",
+       "\n",
+       "                                            abstract               categories  \\\n",
+       "0    A precise determination, and comparison, of ...    astro-ph gr-qc hep-ph   \n",
+       "1    We investigate the global solutions of the D...  math-ph math.AP math.MP   \n",
+       "2    We present here a model for the coupling amo...        cond-mat.mtrl-sci   \n",
+       "3    This paper proposes a novel learning method ...                  nlin.AO   \n",
+       "4    We present an algorithm for constructing an ...              cs.DM cs.DS   \n",
+       "\n",
+       "                                  processed_abstract  \\\n",
+       "0  a precise determination and comparison of the ...   \n",
+       "1  we investigate the global solutions of the dir...   \n",
+       "2  we present here a model for the coupling among...   \n",
+       "3  this paper proposes a novel learning method fo...   \n",
+       "4  we present an algorithm for constructing an op...   \n",
+       "\n",
+       "                                  tokenized_abstract  \n",
+       "0  [[a, precise, determination, ,, and, compariso...  \n",
+       "1  [[we, investigate, the, global, solutions, of,...  \n",
+       "2  [[we, present, here, a, model, for, the, coupl...  \n",
+       "3  [[this, paper, proposes, a, novel, learning, m...  \n",
+       "4  [[we, present, an, algorithm, for, constructin...  "
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import re\n",
+    "\n",
+    "def pre_processor(sentence):\n",
+    "    if not isinstance(sentence, str):\n",
+    "        return \"\"\n",
+    "    # lowercase\n",
+    "    sentence = sentence.lower()\n",
+    "    # remove punctuation using regex\n",
+    "    sentence = re.sub(r\"[^\\w\\s]\", \"\", sentence)\n",
+    "    # remove extra multiple spaces\n",
+    "    sentence = re.sub(r\"\\s+\", \" \", sentence).strip()\n",
+    "    return sentence\n",
+    "\n",
+    "df[\"processed_abstract\"] = df[\"abstract\"].progress_apply(pre_processor)\n",
+    "df.head()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "d-CU5YOfbVui"
+   },
+   "source": [
+    "#### Converting sentences into list of words"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 49,
+     "referenced_widgets": [
+      "3048b4405eb44748b38b5237facafade",
+      "26fbd3641bb142d0a42a3e9e692b1689",
+      "d1404cdae4624c619c7793ac55e7917f",
+      "1362965a367146ff94a368a343b39ee8",
+      "99baec9dffac45eea0d573e9b61f81c1",
+      "a5482757a62545f19ade21ca6f6c4240",
+      "b0828f5f2f204222a40c7e37f827b5b8",
+      "93a30d3ff5084c58a264f1e1cb6e214c",
+      "421f7f03e06546d390c74bd2069883bd",
+      "87140c7e325a4fef97173c83ef14dcea",
+      "35a3cef4a7ca461a8edc089a9db655a1"
+     ]
+    },
+    "id": "Kz0i9reWbVui",
+    "outputId": "7b79cddb-82a0-4e0a-f5d0-af8e5c1a0e94"
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7dfca3237c814703aa915efee7af1032",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1010000 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "abstracts = df['processed_abstract'].values\n",
+    "\n",
+    "def tokenize_sentences(text):\n",
+    "    if not isinstance(text, str) or not text.strip():\n",
+    "        return []\n",
+    "    sentences = sent_tokenize(text)\n",
+    "    return [word_tokenize(sent) for sent in sentences]\n",
+    "\n",
+    "df['tokenized_abstract'] = df['processed_abstract'].progress_apply(lambda x: tokenize_sentences(x))\n",
+    "\n",
+    "corpus_data = df['tokenized_abstract'].to_list()\n",
+    "word2vec_corpus = [item for items in corpus_data for item in items]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "K1MFdpWibVui"
+   },
+   "source": [
+    "#### Using Gensim's API to train Word2Vec\n",
+    "\n",
+    "#### Parameters:\n",
+    "- **input** : Corpus data\n",
+    "- **min_count** : Ignores all words with total frequency lower than this.\n",
+    "- **vector_size** : Dimensionality of the output word vectors.\n",
+    "- **workers** : Use these many worker threads to train the model (=faster training with multicore machines).\n",
+    "- **sg** : Training algorithm: 1 for skip-gram; otherwise CBOW.\n",
+    "- **negative** : If > 0, negative sampling will be used, the int for negative specifies how many \"noise words\"."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "metadata": {
+    "id": "iC7gIDHwbVui"
+   },
+   "outputs": [],
+   "source": [
+    "model = Word2Vec(word2vec_corpus, min_count=3, vector_size= 100, workers=4, window =5, sg = 1, negative=5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.save(\"word2vec_arxiv_skipgram.model\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Loading model weights"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = Word2Vec.load(\"word2vec_arxiv_skipgram.model\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Finding embeddings of each word in abstract and calculating centroid of whole document"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "metadata": {
+    "id": "lfe_zzszbVui"
+   },
+   "outputs": [],
+   "source": [
+    "# initialize centroid column\n",
+    "df[\"centroid\"] = [ [0.0]*100 ] * df.shape[0]\n",
+    "\n",
+    "for index, row in df.iterrows():\n",
+    "    tokenized = row['tokenized_abstract']   # list of list of tokens\n",
+    "    \n",
+    "    # flatten the list of lists\n",
+    "    words = [w for sent in tokenized for w in sent]\n",
+    "\n",
+    "    centroid = np.zeros(100, dtype=float)\n",
+    "    word_count = 0\n",
+    "\n",
+    "    for word in words:\n",
+    "        try:\n",
+    "            vec = model.wv[word]    # get vector from trained Word2Vec\n",
+    "            centroid += vec\n",
+    "            word_count += 1\n",
+    "        except KeyError:\n",
+    "            # word not in vocabulary\n",
+    "            continue\n",
+    "\n",
+    "    # average\n",
+    "    if word_count > 0:\n",
+    "        centroid = centroid / word_count\n",
+    "\n",
+    "    df.at[index, \"centroid\"] = centroid.tolist()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>id</th>\n",
+       "      <th>title</th>\n",
+       "      <th>abstract</th>\n",
+       "      <th>categories</th>\n",
+       "      <th>processed_abstract</th>\n",
+       "      <th>tokenized_abstract</th>\n",
+       "      <th>centroid</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>706.1314</td>\n",
+       "      <td>When Did Cosmic Acceleration Start ?</td>\n",
+       "      <td>A precise determination, and comparison, of ...</td>\n",
+       "      <td>astro-ph gr-qc hep-ph</td>\n",
+       "      <td>a precise determination and comparison of the ...</td>\n",
+       "      <td>[[a, precise, determination, and, comparison, ...</td>\n",
+       "      <td>[-0.08764898034384136, 0.11816140904464356, -0...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>706.1315</td>\n",
+       "      <td>The Dirac system on the Anti-de Sitter Universe</td>\n",
+       "      <td>We investigate the global solutions of the D...</td>\n",
+       "      <td>math-ph math.AP math.MP</td>\n",
+       "      <td>we investigate the global solutions of the dir...</td>\n",
+       "      <td>[[we, investigate, the, global, solutions, of,...</td>\n",
+       "      <td>[-0.2019515800796124, 0.14643365251905624, -0....</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>706.1316</td>\n",
+       "      <td>Coupling of Optical Lumped Nanocircuit Element...</td>\n",
+       "      <td>We present here a model for the coupling amo...</td>\n",
+       "      <td>cond-mat.mtrl-sci</td>\n",
+       "      <td>we present here a model for the coupling among...</td>\n",
+       "      <td>[[we, present, here, a, model, for, the, coupl...</td>\n",
+       "      <td>[-0.10665342995032136, 0.15606039330525942, -0...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>706.1317</td>\n",
+       "      <td>A model for learning to segment temporal seque...</td>\n",
+       "      <td>This paper proposes a novel learning method ...</td>\n",
+       "      <td>nlin.AO</td>\n",
+       "      <td>this paper proposes a novel learning method fo...</td>\n",
+       "      <td>[[this, paper, proposes, a, novel, learning, m...</td>\n",
+       "      <td>[-0.018284802690839336, 0.009443171508610248, ...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>706.1318</td>\n",
+       "      <td>Constructing a maximum utility slate of on-lin...</td>\n",
+       "      <td>We present an algorithm for constructing an ...</td>\n",
+       "      <td>cs.DM cs.DS</td>\n",
+       "      <td>we present an algorithm for constructing an op...</td>\n",
+       "      <td>[[we, present, an, algorithm, for, constructin...</td>\n",
+       "      <td>[-0.07965052640159434, 0.05054909288165002, -0...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         id                                              title  \\\n",
+       "0  706.1314               When Did Cosmic Acceleration Start ?   \n",
+       "1  706.1315    The Dirac system on the Anti-de Sitter Universe   \n",
+       "2  706.1316  Coupling of Optical Lumped Nanocircuit Element...   \n",
+       "3  706.1317  A model for learning to segment temporal seque...   \n",
+       "4  706.1318  Constructing a maximum utility slate of on-lin...   \n",
+       "\n",
+       "                                            abstract               categories  \\\n",
+       "0    A precise determination, and comparison, of ...    astro-ph gr-qc hep-ph   \n",
+       "1    We investigate the global solutions of the D...  math-ph math.AP math.MP   \n",
+       "2    We present here a model for the coupling amo...        cond-mat.mtrl-sci   \n",
+       "3    This paper proposes a novel learning method ...                  nlin.AO   \n",
+       "4    We present an algorithm for constructing an ...              cs.DM cs.DS   \n",
+       "\n",
+       "                                  processed_abstract  \\\n",
+       "0  a precise determination and comparison of the ...   \n",
+       "1  we investigate the global solutions of the dir...   \n",
+       "2  we present here a model for the coupling among...   \n",
+       "3  this paper proposes a novel learning method fo...   \n",
+       "4  we present an algorithm for constructing an op...   \n",
+       "\n",
+       "                                  tokenized_abstract  \\\n",
+       "0  [[a, precise, determination, and, comparison, ...   \n",
+       "1  [[we, investigate, the, global, solutions, of,...   \n",
+       "2  [[we, present, here, a, model, for, the, coupl...   \n",
+       "3  [[this, paper, proposes, a, novel, learning, m...   \n",
+       "4  [[we, present, an, algorithm, for, constructin...   \n",
+       "\n",
+       "                                            centroid  \n",
+       "0  [-0.08764898034384136, 0.11816140904464356, -0...  \n",
+       "1  [-0.2019515800796124, 0.14643365251905624, -0....  \n",
+       "2  [-0.10665342995032136, 0.15606039330525942, -0...  \n",
+       "3  [-0.018284802690839336, 0.009443171508610248, ...  \n",
+       "4  [-0.07965052640159434, 0.05054909288165002, -0...  "
+      ]
+     },
+     "execution_count": 42,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 43,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df.to_pickle(\"data.pkl\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Saving Dataset as parquet for efficient loading and compression"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df.to_parquet(\"data.parquet\", engine=\"fastparquet\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Loading dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = pd.read_pickle(\"data.pkl\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = pd.read_parquet(\"data.parquet\", engine=\"fastparquet\") ## loading parquet dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>id</th>\n",
+       "      <th>title</th>\n",
+       "      <th>abstract</th>\n",
+       "      <th>categories</th>\n",
+       "      <th>processed_abstract</th>\n",
+       "      <th>tokenized_abstract</th>\n",
+       "      <th>centroid</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>706.1314</td>\n",
+       "      <td>When Did Cosmic Acceleration Start ?</td>\n",
+       "      <td>A precise determination, and comparison, of ...</td>\n",
+       "      <td>astro-ph gr-qc hep-ph</td>\n",
+       "      <td>a precise determination and comparison of the ...</td>\n",
+       "      <td>[[a, precise, determination, and, comparison, ...</td>\n",
+       "      <td>[-0.08764898034384136, 0.11816140904464356, -0...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>706.1315</td>\n",
+       "      <td>The Dirac system on the Anti-de Sitter Universe</td>\n",
+       "      <td>We investigate the global solutions of the D...</td>\n",
+       "      <td>math-ph math.AP math.MP</td>\n",
+       "      <td>we investigate the global solutions of the dir...</td>\n",
+       "      <td>[[we, investigate, the, global, solutions, of,...</td>\n",
+       "      <td>[-0.2019515800796124, 0.14643365251905624, -0....</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>706.1316</td>\n",
+       "      <td>Coupling of Optical Lumped Nanocircuit Element...</td>\n",
+       "      <td>We present here a model for the coupling amo...</td>\n",
+       "      <td>cond-mat.mtrl-sci</td>\n",
+       "      <td>we present here a model for the coupling among...</td>\n",
+       "      <td>[[we, present, here, a, model, for, the, coupl...</td>\n",
+       "      <td>[-0.10665342995032136, 0.15606039330525942, -0...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>706.1317</td>\n",
+       "      <td>A model for learning to segment temporal seque...</td>\n",
+       "      <td>This paper proposes a novel learning method ...</td>\n",
+       "      <td>nlin.AO</td>\n",
+       "      <td>this paper proposes a novel learning method fo...</td>\n",
+       "      <td>[[this, paper, proposes, a, novel, learning, m...</td>\n",
+       "      <td>[-0.018284802690839336, 0.009443171508610248, ...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>706.1318</td>\n",
+       "      <td>Constructing a maximum utility slate of on-lin...</td>\n",
+       "      <td>We present an algorithm for constructing an ...</td>\n",
+       "      <td>cs.DM cs.DS</td>\n",
+       "      <td>we present an algorithm for constructing an op...</td>\n",
+       "      <td>[[we, present, an, algorithm, for, constructin...</td>\n",
+       "      <td>[-0.07965052640159434, 0.05054909288165002, -0...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         id                                              title  \\\n",
+       "0  706.1314               When Did Cosmic Acceleration Start ?   \n",
+       "1  706.1315    The Dirac system on the Anti-de Sitter Universe   \n",
+       "2  706.1316  Coupling of Optical Lumped Nanocircuit Element...   \n",
+       "3  706.1317  A model for learning to segment temporal seque...   \n",
+       "4  706.1318  Constructing a maximum utility slate of on-lin...   \n",
+       "\n",
+       "                                            abstract               categories  \\\n",
+       "0    A precise determination, and comparison, of ...    astro-ph gr-qc hep-ph   \n",
+       "1    We investigate the global solutions of the D...  math-ph math.AP math.MP   \n",
+       "2    We present here a model for the coupling amo...        cond-mat.mtrl-sci   \n",
+       "3    This paper proposes a novel learning method ...                  nlin.AO   \n",
+       "4    We present an algorithm for constructing an ...              cs.DM cs.DS   \n",
+       "\n",
+       "                                  processed_abstract  \\\n",
+       "0  a precise determination and comparison of the ...   \n",
+       "1  we investigate the global solutions of the dir...   \n",
+       "2  we present here a model for the coupling among...   \n",
+       "3  this paper proposes a novel learning method fo...   \n",
+       "4  we present an algorithm for constructing an op...   \n",
+       "\n",
+       "                                  tokenized_abstract  \\\n",
+       "0  [[a, precise, determination, and, comparison, ...   \n",
+       "1  [[we, investigate, the, global, solutions, of,...   \n",
+       "2  [[we, present, here, a, model, for, the, coupl...   \n",
+       "3  [[this, paper, proposes, a, novel, learning, m...   \n",
+       "4  [[we, present, an, algorithm, for, constructin...   \n",
+       "\n",
+       "                                            centroid  \n",
+       "0  [-0.08764898034384136, 0.11816140904464356, -0...  \n",
+       "1  [-0.2019515800796124, 0.14643365251905624, -0....  \n",
+       "2  [-0.10665342995032136, 0.15606039330525942, -0...  \n",
+       "3  [-0.018284802690839336, 0.009443171508610248, ...  \n",
+       "4  [-0.07965052640159434, 0.05054909288165002, -0...  "
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df.head()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Using Simple Cosine Similiarity to search top K papers"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {
+    "id": "YUtUbO4ebVuj"
+   },
+   "outputs": [],
+   "source": [
+    "def rank_docs(model, query, df_covid, num):\n",
+    "\n",
+    "    cosine_list = []\n",
+    "\n",
+    "    a = []\n",
+    "    query_words = query.lower().split(\" \") # Lowercase query words\n",
+    "    found_words = []\n",
+    "    not_found_words = []\n",
+    "\n",
+    "    for q_word in query_words:\n",
+    "        try:\n",
+    "            a.append(model.wv[q_word]) # Use model.wv for word vectors\n",
+    "            found_words.append(q_word)\n",
+    "        except KeyError:\n",
+    "            not_found_words.append(q_word)\n",
+    "            continue\n",
+    "\n",
+    "    if not_found_words:\n",
+    "        print(f\"Warning: The following query words were not found in the model's vocabulary: {', '.join(not_found_words)}\")\n",
+    "\n",
+    "    if not found_words:\n",
+    "        print(\"No query words were found in the model's vocabulary. Returning empty list.\")\n",
+    "        return []\n",
+    "\n",
+    "    for index, row in df_covid.iterrows():\n",
+    "        centroid = row['centroid']\n",
+    "        total_sim = 0\n",
+    "        # Ensure centroid is not all zeros to avoid division by zero\n",
+    "        if np.linalg.norm(centroid) == 0:\n",
+    "            continue\n",
+    "\n",
+    "        for a_i in a:\n",
+    "            # Ensure a_i is not all zeros to avoid division by zero\n",
+    "            if np.linalg.norm(a_i) == 0:\n",
+    "                continue\n",
+    "            cos_sim = np.dot(a_i, centroid)/(np.linalg.norm(a_i)*np.linalg.norm(centroid))\n",
+    "            total_sim += cos_sim\n",
+    "        cosine_list.append((row['title'], total_sim, row['abstract'], row['id']))\n",
+    "\n",
+    "    cosine_list.sort(key=lambda x:x[1], reverse=True) ## in Descending order\n",
+    "\n",
+    "    papers_list = []\n",
+    "    for item in cosine_list[:num]:\n",
+    "        papers_list.append((item[0], item[1], item[2], item[3])) # Only append title and total_sim\n",
+    "    return papers_list"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {
+    "id": "yCo-1JTmbVuj"
+   },
+   "outputs": [],
+   "source": [
+    "def query(query, top_matches=10):\n",
+    "    model_to_use = model\n",
+    "    df_to_use = df\n",
+    "    return rank_docs(model_to_use, query, df_to_use, top_matches)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "PMBh22TBbVuj",
+    "outputId": "07f8412a-47e3-4c5b-a2c9-1234a485085c"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[('Reversibility of laser filamentation', np.float64(1.2172915391534613), '  We investigate the reversibility of laser filamentation, a self-sustained,\\nnon-linear propagation regime including dissipation and time-retarded effects.\\nWe show that even losses related to ionization marginally affect the\\npossibility of reverse propagating ultrashort pulses back to the initial\\nconditions, although they make it prone to finite-distance blow-up susceptible\\nto prevent backward propagation.\\n', 1401.284), ('Droplet-shaped waves: Causal finite-support analogs of X-shaped waves', np.float64(1.206201484806961), '  A model of steady-state X-shaped wave generation by a superluminal\\n(supersonic) pointlike source infinitely moving along a straight line is\\nextended to a more realistic causal scenario of a source pulse launched at time\\nzero and propagating rectilinearly at constant superluminal speed. In the case\\nof infinitely short (delta) pulse, the new model yields an analytical solution,\\ncorresponding to the propagation-invariant X-shaped wave clipped by a\\ndroplet-shaped support, which perpetually expands along the propagation and\\ntransversal directions, thus tending the droplet-shaped wave to the X-shaped\\none.\\n', 1110.3494), ('Role of beam propagation in Goos-H\\\\\"{a}nchen and Imbert-Fedorov shifts', np.float64(1.1974150038736082), '  We derive the polarization-dependent displacements parallel and perpendicular\\nto the plane of incidence, for a Gaussian light beam reflected from a planar\\ninterface, taking into account the propagation of the beam. Using a\\nclassical-optics formalism we show that beam propagation may greatly affect\\nboth Goos-H\\\\\"{a}nchen and Imbert-Fedorov shifts when the incident beam is\\nfocussed.\\n', 804.1895), ('The falling slinky', np.float64(1.1963728894486865), '  The slinky, released from rest hanging under its own weight, falls in a\\npeculiar manner. The bottom stays at rest until a wave hits it from above. Two\\ncases -- one unphysical one where the slinky is able to pass through itself,\\nand the other where the coils of the slinky collide creating a shock wave\\ntravelling down the slinky -- are analysed. In the former case, the bottom\\nbegins to move much later than in the latter.\\n', 1110.4368), ('MHD Mode Conversion around a 2D Magnetic Null Point', np.float64(1.1959848595837723), '  Mode conversion occurs when a wave passes through a region where the sound\\nand Alfven speeds are equal. At this point there is a resonance, which allows\\nsome of the incident wave to be converted into a different mode. We study this\\nphenomenon in the vicinity of a two-dimensional, coronal null point. As a wave\\napproaches the null it passes from low- to high-beta plasma, allowing\\nconversion to take place. We simulate this numerically by sending in a slow\\nmagnetoacoustic wave from the upper boundary; as this passes through the\\nconversion layer a fast wave can clearly be seen propagating ahead. Numerical\\nsimulations combined with an analytical WKB investigation allow us to determine\\nand track both the incident and converted waves throughout the domain.\\n', 907.1541), ('An analog fluid model for some tachyonic effects in field theory', np.float64(1.1954942312587349), '  We consider the sound radiation from an acoustic point-like source moving\\nalong a supersonic (\"space-like\") trajectory in a fluid at rest. We call it an\\nacoustic \"tachyonic\" source. We describe the radiation emitted by this\\nsupersonic source. After quantizing the acoustic perturbations, we present the\\ndistribution of phonons generated by this classical tachyonic source and the\\nclassical wave interference pattern.\\n', 1109.098), ('Interactions in an acoustic world', np.float64(1.1949165744874333), '  The present paper aims to complete an earlier paper where the acoustic world\\nwas introduced. This is accomplished by analyzing the interactions which occur\\nbetween the inhomogeneities of the acoustic medium, which are induced by the\\nacoustic vibrations traveling in the medium. When a wave packet travels in a\\nmedium, the medium becomes inhomogeneous. The spherical wave packet behaves\\nlike an acoustic spherical lens for the acoustic plane waves. According to the\\nprinciple of causality, there is an interaction between the wave and plane wave\\npacket. In specific conditions the wave packet behaves as an acoustic black\\nhole.\\n', 1612.00294), ('Peculiar Behavior of Si Cluster Ions in Solid Al', np.float64(1.193792102194581), '  A peculiar ion behavior is found in a Si cluster, moving with a speed of\\n~0.22c (c: speed of light) in a solid Al plasma: the Si ion, moving behind the\\nforward moving Si ion closely in a several angstrom distance in the cluster,\\nfeels the wake field generated by the forward Si. The interaction potential on\\nthe rear Si may balance the deceleration backward force by itself with the\\nacceleration forward force by the forward Si in the longitudinal moving\\ndirection. The forward Si would be decelerated normally. However, the\\ndeceleration of the rear Si, moving behind closely, would be reduced\\nsignificantly, and the rear Si may catch up and overtake the forward moving Si\\nin the cluster during the Si cluster interaction with the high-density Al\\nplasma.\\n', 1808.02649), ('The effective Hamiltonian which governs the propagation dynamics of\\n  nonspreading wave packets', np.float64(1.1935298414558697), \"  We discuss the propagation dynamics of nonspreading wave packets. We\\ndecompose the Hamiltonian into two parts. The first part is such that wave\\npackets is its instantaneous eigenstate and is therefore irrelevant to the\\npropagation of the packet. The second part is shown to be the effective\\nHamiltonian governing the motion of the packet both classically and quantum\\nmechanically. Thus, analogous to Ehrenfest's theorem, nonspreading wave packets\\noffer another view point directly connecting quantum mechanics and classical\\nmechanics. This analysis also works for non-square-integrable packets, such as\\nAiry packets.\\n\", 1306.1311), ('Multiple trains of same-color surface plasmon-polaritons guided by the\\n  planar interface of a metal and a sculptured nematic thin film. Part IV:\\n  Canonical problem', np.float64(1.1905086316212172), '  The canonical problem of the propagation of surface-plasmon-polariton (SPP)\\nwaves localized to the planar interface of a metal and a sculptured nematic\\nthin film (SNTF) that is periodically nonhomogeneous along the direction normal\\nto the interface was formulated. Solution of the dispersion equation obtained\\nthereby confirmed the possibility of exciting multiple SPP waves of the same\\nfrequency or color. However, these SPP waves differ in phase speed, field\\nstructure, and the e-folding distance along the direction of propagation.\\n', 1002.2435)]\n",
+      "Total Time taken by query: 91.10928773880005\n"
+     ]
+    }
+   ],
+   "source": [
+    "import time\n",
+    "start_time = time.time()\n",
+    "print(query('back propagation'))\n",
+    "end_time = time.time()\n",
+    "print(f\"Total Time taken by query: {end_time-start_time}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Search Optimizations"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Using FAISS for getting top K Result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import faiss"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Convert centroid column to matrix\n",
+    "X = np.vstack(df['centroid'].values).astype('float32')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Normalize for cosine similarity\n",
+    "faiss.normalize_L2(X)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Build FAISS index\n",
+    "dim = X.shape[1]\n",
+    "flat_index = faiss.IndexFlatIP(dim)\n",
+    "index = faiss.IndexIDMap(flat_index)\n",
+    "\n",
+    "ids = np.arange(X.shape[0]).astype('int64')\n",
+    "index.add_with_ids(X, ids)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "FAISS index built with 1010000 vectors\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"FAISS index built with\", index.ntotal, \"vectors\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 60,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def faiss_query(query, top_k=10):\n",
+    "    words = word_tokenize(query.lower()) # use sent_tokenizer here\n",
+    "    print(words)\n",
+    "    vecs = []\n",
+    "    for w in words:\n",
+    "        if w in model.wv:\n",
+    "            vecs.append(model.wv[w])\n",
+    "    if len(vecs) == 0:\n",
+    "        return []\n",
+    "    qvec = np.mean(vecs, axis=0).astype('float32').reshape(1, -1)\n",
+    "    faiss.normalize_L2(qvec)\n",
+    "    scores, neighbors = index.search(qvec, top_k)\n",
+    "    return df.iloc[neighbors[0]]\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 63,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['black', 'holes']\n",
+      "                 id                                              title  \\\n",
+      "593935   1503.01221  Superradiantly stable non-extremal Reissner-No...   \n",
+      "174856   1004.29160          Area spectra of near extremal black holes   \n",
+      "245022   1104.14680  Exact solutions of higher dimensional black holes   \n",
+      "111775    905.01790  Thermodynamics of Ho\\v{r}ava-Lifshitz black holes   \n",
+      "112553    905.09570  Thermodynamics of black holes in the deformed ...   \n",
+      "75657     810.00780             Regular black hole in three dimensions   \n",
+      "34761     801.24340  Phase transition for black holes with scalar h...   \n",
+      "1001755  1808.02609  Instability of Reissner-Nordstr\\\"{o}m black ho...   \n",
+      "192983   1007.37450  A scalar field condensation instability of rot...   \n",
+      "470646   1311.69850  Thermodynamic and classical instability of AdS...   \n",
+      "\n",
+      "                                                  abstract    categories  \\\n",
+      "593935     The superradiant stability is investigated f...         gr-qc   \n",
+      "174856     Motivated by Maggiore's new interpretation o...         gr-qc   \n",
+      "245022     We review exact solutions of black holes in ...  hep-th gr-qc   \n",
+      "111775     We study black holes in the Ho\\v{r}ava-Lifsh...  hep-th gr-qc   \n",
+      "112553     We study thermodynamics of black holes in th...        hep-th   \n",
+      "75657      We find a new black hole in three dimensiona...         gr-qc   \n",
+      "34761      We study phase transitions between black hol...  hep-th gr-qc   \n",
+      "1001755    The scalarization of Reissner-Nordstr\\\"{o}m ...  gr-qc hep-th   \n",
+      "192983     Near-extreme Reissner-Nordstrom-anti-de Sitt...  hep-th gr-qc   \n",
+      "470646     We study thermodynamic and classical instabi...  hep-th gr-qc   \n",
+      "\n",
+      "                                        processed_abstract  \\\n",
+      "593935   the superradiant stability is investigated for...   \n",
+      "174856   motivated by maggiores new interpretation of q...   \n",
+      "245022   we review exact solutions of black holes in hi...   \n",
+      "111775   we study black holes in the hovravalifshitz gr...   \n",
+      "112553   we study thermodynamics of black holes in the ...   \n",
+      "75657    we find a new black hole in three dimensional ...   \n",
+      "34761    we study phase transitions between black holes...   \n",
+      "1001755  the scalarization of reissnernordstrom black h...   \n",
+      "192983   nearextreme reissnernordstromantide sitter bla...   \n",
+      "470646   we study thermodynamic and classical instabili...   \n",
+      "\n",
+      "                                        tokenized_abstract  \\\n",
+      "593935   [[the, superradiant, stability, is, investigat...   \n",
+      "174856   [[motivated, by, maggiores, new, interpretatio...   \n",
+      "245022   [[we, review, exact, solutions, of, black, hol...   \n",
+      "111775   [[we, study, black, holes, in, the, hovravalif...   \n",
+      "112553   [[we, study, thermodynamics, of, black, holes,...   \n",
+      "75657    [[we, find, a, new, black, hole, in, three, di...   \n",
+      "34761    [[we, study, phase, transitions, between, blac...   \n",
+      "1001755  [[the, scalarization, of, reissnernordstrom, b...   \n",
+      "192983   [[nearextreme, reissnernordstromantide, sitter...   \n",
+      "470646   [[we, study, thermodynamic, and, classical, in...   \n",
+      "\n",
+      "                                                  centroid  \n",
+      "593935   [-0.1548219741116343, 0.07305222587574674, -0....  \n",
+      "174856   [-0.15942226781044155, 0.09534141645417549, -0...  \n",
+      "245022   [-0.16427493500797188, -0.021989740924361872, ...  \n",
+      "111775   [-0.185662076366134, 0.1500076398253441, -0.01...  \n",
+      "112553   [-0.1874230935668143, 0.0645052320861186, -0.0...  \n",
+      "75657    [-0.140232238308366, 0.11955675896450853, -0.0...  \n",
+      "34761    [-0.18028071506373716, 0.18077764537577568, -0...  \n",
+      "1001755  [-0.1625649823457934, 0.08357451187912375, -0....  \n",
+      "192983   [-0.20853021891077397, 0.17061649511522387, -0...  \n",
+      "470646   [-0.22651504412189954, 0.09521500566215427, -0...  \n",
+      "Total Time taken by query: 0.0438535213470459\n"
+     ]
+    }
+   ],
+   "source": [
+    "import time\n",
+    "start_time = time.time()\n",
+    "print(faiss_query('black holes'))\n",
+    "end_time = time.time()\n",
+    "print(f\"Total Time taken by query: {end_time-start_time}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "As we can se using faiss drastically decreases search time from 91.1 seconds to 0.04 seconds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 64,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "faiss.write_index(index, \"faiss_search_index.bin\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kaggle": {
+   "accelerator": "gpu",
+   "dataSources": [
+    {
+     "datasetId": 612177,
+     "sourceId": 13661773,
+     "sourceType": "datasetVersion"
+    },
+    {
+     "isSourceIdPinned": true,
+     "modelId": 499524,
+     "modelInstanceId": 484022,
+     "sourceId": 641799,
+     "sourceType": "modelInstanceVersion"
+    }
+   ],
+   "dockerImageVersionId": 31192,
+   "isGpuEnabled": true,
+   "isInternetEnabled": true,
+   "language": "python",
+   "sourceType": "notebook"
+  },
+  "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.12.3"
+  },
+  "widgets": {
+   "application/vnd.jupyter.widget-state+json": {
+    "06d6872e3d1f4086a4f16b8caababccb": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "1362965a367146ff94a368a343b39ee8": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_87140c7e325a4fef97173c83ef14dcea",
+      "placeholder": "​",
+      "style": "IPY_MODEL_35a3cef4a7ca461a8edc089a9db655a1",
+      "value": " 110000/110000 [01:35&lt;00:00, 710.46it/s]"
+     }
+    },
+    "214288b8cb6c453084c3cce01982eff7": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "26fbd3641bb142d0a42a3e9e692b1689": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_a5482757a62545f19ade21ca6f6c4240",
+      "placeholder": "​",
+      "style": "IPY_MODEL_b0828f5f2f204222a40c7e37f827b5b8",
+      "value": "100%"
+     }
+    },
+    "3048b4405eb44748b38b5237facafade": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_26fbd3641bb142d0a42a3e9e692b1689",
+       "IPY_MODEL_d1404cdae4624c619c7793ac55e7917f",
+       "IPY_MODEL_1362965a367146ff94a368a343b39ee8"
+      ],
+      "layout": "IPY_MODEL_99baec9dffac45eea0d573e9b61f81c1"
+     }
+    },
+    "35a3cef4a7ca461a8edc089a9db655a1": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "421f7f03e06546d390c74bd2069883bd": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "4de066fc521f493cab302c8d5ca255e2": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "4df933c8001c4a27b351c5b66df23eb5": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "52f12329e31645888592d0120a05f6ea": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "5510fb4117fd4954aeeb0de2214b1ed4": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_4df933c8001c4a27b351c5b66df23eb5",
+      "placeholder": "​",
+      "style": "IPY_MODEL_80071f18a3fd4ddc865a8a536d83cd2e",
+      "value": " 110000/110000 [00:03&lt;00:00, 39633.11it/s]"
+     }
+    },
+    "7d5629fd92c543bcaafe7fd352c079c6": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_a815fb290af249cbb23679e34d1fdaa3",
+       "IPY_MODEL_b042bdc131654d25bac957a4e921f59d",
+       "IPY_MODEL_5510fb4117fd4954aeeb0de2214b1ed4"
+      ],
+      "layout": "IPY_MODEL_4de066fc521f493cab302c8d5ca255e2"
+     }
+    },
+    "80071f18a3fd4ddc865a8a536d83cd2e": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "87140c7e325a4fef97173c83ef14dcea": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "88ba2081fc004bbdb6f438c5031458d5": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "93a30d3ff5084c58a264f1e1cb6e214c": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "99baec9dffac45eea0d573e9b61f81c1": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "a5482757a62545f19ade21ca6f6c4240": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "a815fb290af249cbb23679e34d1fdaa3": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_06d6872e3d1f4086a4f16b8caababccb",
+      "placeholder": "​",
+      "style": "IPY_MODEL_88ba2081fc004bbdb6f438c5031458d5",
+      "value": "100%"
+     }
+    },
+    "b042bdc131654d25bac957a4e921f59d": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_52f12329e31645888592d0120a05f6ea",
+      "max": 110000,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_214288b8cb6c453084c3cce01982eff7",
+      "value": 110000
+     }
+    },
+    "b0828f5f2f204222a40c7e37f827b5b8": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "d1404cdae4624c619c7793ac55e7917f": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_93a30d3ff5084c58a264f1e1cb6e214c",
+      "max": 110000,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_421f7f03e06546d390c74bd2069883bd",
+      "value": 110000
+     }
+    }
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}