datasetParsing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download the datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 21 Mil WikiDPR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"wiki_dpr\", \"psgs_w100.multiset.no_index\", streaming= True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# print(next(iter(dataset['train'])))\n",
    "# embeds = np.array([doc['emb'] for doc in dataset['train']])\n",
    "# print(\"Embeddings:\")\n",
    "# print(embeds)\n",
    "\n",
    "# get the document of id 500\n",
    "# temp = dataset['train'].filter(lambda x: x['id'] == 0)\n",
    "# print(\"Document of id 500:\")\n",
    "# print(next(iter(temp)))\n",
    "\n",
    "# dataset_head = dataset['train'].skip(10000)\n",
    "# dataset_head = dataset_head.take(1)\n",
    "# print(\"Head:\")\n",
    "# print(next(iter(dataset_head)))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cohere wiki 485k dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas\n",
    "from datasets import load_dataset\n",
    "\n",
    "docs = load_dataset(f\"Cohere/wikipedia-22-12-simple-embeddings\", split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for doc in docs:\n",
    "\n",
    "    with open(\"embeds2.txt\", \"a\") as f:\n",
    "        f.write(str(doc['emb']))\n",
    "        f.write('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clustering Using Faiss KMeans"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reading pre-downloaded parquet files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read embeddings  (535424,)\n",
      "Read ids  (535424,)\n",
      "(535424,)\n"
     ]
    }
   ],
   "source": [
    "# Read the dataset from the files in /batches/dataset in one big numpy array\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "\n",
    "path = \"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/Data_6/dataset/\"\n",
    "\n",
    "parq = pd.read_parquet(path, engine=\"fastparquet\", columns=[\"id\", \"embeddings\"])\n",
    "# print(embeds[\"embeddings\"][0])\n",
    "# normalize the embeddings\n",
    "parq['embeddings'] = parq['embeddings'].apply(lambda x: np.array(x) / np.linalg.norm(x) if x is not None else None)\n",
    "# print(embeds[\"embeddings\"][0])\n",
    "embeds = parq['embeddings'].values\n",
    "print(\"Read embeddings \", embeds.shape)\n",
    "ids = parq['id'].values\n",
    "print(\"Read ids \", ids.shape)\n",
    "del parq\n",
    "print(embeds.shape)\n",
    "embeds = np.vstack(embeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Written 0-embeds-batch.txt\n",
      "Written 1-embeds-batch.txt\n",
      "Written 2-embeds-batch.txt\n",
      "Written 3-embeds-batch.txt\n",
      "Written 4-embeds-batch.txt\n",
      "Written 5-embeds-batch.txt\n"
     ]
    }
   ],
   "source": [
    "# write the embeds in normalizedDataset each file is 100000\n",
    "c = 0\n",
    "for i in range(0, len(embeds), 100000):\n",
    "    with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/Data_6/normalizedDataset/{c}-embeds-batch.txt\", \"w\") as f:\n",
    "        for j in range(i, min(i + 100000, len(embeds))):\n",
    "            f.write(str(ids[j]) + \" \")\n",
    "            f.write(str(embeds[j]).replace('  ',' ').replace('[','').replace(']',''))\n",
    "            f.write('\\n')\n",
    "    print(f\"Written {c}-embeds-batch.txt\")\n",
    "    c += 1\n",
    "# print(ids[82050])\n",
    "# with open(\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/Data_1/normalizedDataset/0-embeds-batch.txt\", \"a\") as f:\n",
    "#     for i in range(82050,100000):\n",
    "#         f.write(str(ids[i]) + \" \")\n",
    "#         f.write(str(embeds[i]).replace('  ',' ').replace('[','').replace(']',''))\n",
    "#         f.write('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reading the dataset from the pre-downloaded data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read the dataset from the files in /batches/dataset in one big numpy array\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "\n",
    "embeds = []\n",
    "path = \"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/Data_0/dataset/\"\n",
    "\n",
    "pd.read_parquet(path, engine=\"fastparquet\", columns=[\"id\", \"embeddings\"])\n",
    "\n",
    "\n",
    "for i in range(5):\n",
    "   \n",
    "    with open(f\"./Data/dataset/{i}-embeds-batch.txt\", \"r\") as f:\n",
    "        lines = [np.fromstring(line, sep=\" \") for line in f]\n",
    "        # normalize the embeddings\n",
    "        lines = [line / np.linalg.norm(line) for line in lines]\n",
    "        embeds.append(lines)\n",
    "        f.close()\n",
    "    # # write the normalized embeddings to a file line by line\n",
    "    # with open(f\"./Data/normalizedDataset/{i}-embeds-batch.txt\", \"a\") as f:\n",
    "    #     for line in lines:\n",
    "    #         line = np.array(line)\n",
    "    #         f.write(str(line).replace('[', '').replace(']', '').replace('  ',' '))\n",
    "    #         f.write('\\n')\n",
    "    #     f.close()\n",
    "    \n",
    "\n",
    "# Concatenate all lists into a\n",
    "embeds = np.concatenate(embeds, axis=0)\n",
    "\n",
    "# print the norm of the first 100 embeddings\n",
    "for i in range(10):\n",
    "    print(np.linalg.norm(embeds[i]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Constructing the clusters using Faiss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(535424, 768)\n"
     ]
    }
   ],
   "source": [
    "!pip install faiss-cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.06638588e-03  4.65310067e-02 -2.50247270e-02 -5.58813885e-02 -4.46244627e-02  4.14618179e-02  5.71691506e-02 -6.84761908e-03 -1.98977273e-02 -3.59311737e-02 -5.13926297e-02  9.38886683e-03 -1.76828336e-02  2.40578894e-02  7.11571425e-03  1.45968189e-03  2.79998016e-02  4.30624187e-02 -5.72291389e-02  2.21943315e-02 -5.91637846e-03  4.58485214e-03 -2.09718961e-02  9.58977919e-03  3.36066857e-02  2.57384889e-02  3.45021971e-02 -2.65295338e-02 -3.16395946e-02 -2.35705748e-02  1.50086200e-02 -8.02448206e-03  4.15298045e-02 -1.67061538e-02 -1.72677170e-02 -1.42657869e-02 -9.58224572e-03  6.86428696e-03  2.52715982e-02  2.52384823e-02 -4.59874049e-03 -4.61478457e-02  2.02011690e-02 -2.98269722e-03 -3.66607383e-02 -2.08327528e-02 -5.81736639e-02  6.76988438e-02 -1.98959652e-03  9.70022473e-03 -3.93440761e-02 -4.56899358e-03 -4.33134958e-02  5.28937541e-02 -8.20746704e-04 -6.58629090e-03 -3.06480564e-02 -1.73739670e-03  2.04576682e-02 -2.56287213e-02  2.68921610e-02  2.23283973e-02  7.54206628e-02  5.93048800e-03 -2.28522439e-03  2.93510351e-02 -2.87275203e-02  6.40924787e-03 -4.80551682e-02  1.49306310e-02  8.26281216e-03 -1.98361706e-02 -7.59598985e-02 -3.02249892e-03  3.88769223e-03 -3.11482251e-02  1.24735050e-02 -1.29203824e-03  2.77004531e-03  2.79771835e-02 -4.23028171e-02 -9.96377505e-03  2.33756173e-02  1.52565204e-02  4.40801121e-03  2.70224903e-02  1.29643073e-02  1.33023644e-02 -7.12170964e-03 -7.02917427e-02 -1.60790887e-02  1.88135356e-02  4.11877548e-03 -2.24210019e-03 -3.28616835e-02  1.60126444e-02  2.52769887e-02 -1.47258528e-02 -8.28830898e-03 -3.65636125e-02  9.34224389e-03  8.32028594e-03  4.71876003e-02 -8.18489213e-03 -1.05804410e-02 -1.19326795e-02  1.51963243e-02 -2.57353093e-02 -2.10155156e-02  2.13261135e-02 -1.87764838e-02  4.62610312e-02  1.32681374e-02 -1.91860623e-03 -5.07834880e-03  3.15115638e-02 -3.86581495e-02 -8.25165510e-02 -3.08829267e-02 -5.55286463e-03 -5.57590276e-03  5.02912141e-02  3.84913199e-02  6.66047111e-02 -1.29078664e-02 -4.66472805e-02  3.95395001e-03  3.54073048e-02 -4.76902835e-02  7.62749114e-04  7.70211359e-03  4.85905027e-03  3.61961275e-02 -4.67846356e-03 -4.47231941e-02 -2.60779522e-02 -2.57086474e-02  1.83285736e-02 -1.22016121e-03 -1.24039277e-02 -8.78970027e-02 -3.55161130e-02  9.66316182e-03 -8.96683894e-03  6.43652305e-03 -1.96397547e-02 -2.34360970e-03 -2.03897581e-02 -2.22887900e-02  1.35448240e-02 -2.25546490e-02  1.30993593e-02  7.88297970e-03 -3.87431867e-02 -1.44980904e-02  2.84461994e-02  1.48538407e-02  1.82092879e-02  8.97703134e-03  1.12585071e-02  1.34722730e-02  1.31027941e-02  5.10097621e-03 -2.74189357e-02  1.30135519e-02  1.32531729e-02  2.69728731e-02 -3.34250443e-02  2.02757787e-04 -4.34620492e-02 -1.28273088e-02 -2.17009410e-02  1.98685080e-02 -4.19310257e-02 -8.22116714e-03  2.31031235e-03  3.93724367e-02 -1.38133811e-02  9.83173773e-03 -3.85078602e-02 -5.73380329e-02  2.91792187e-03 -4.47613485e-02 -4.61597033e-02  7.53113767e-03  2.07212735e-02 -2.38919780e-02  1.25317862e-02 -1.06590921e-02  7.68612884e-03 -4.56861518e-02 -3.02816685e-02 -1.38591947e-02 -3.81639833e-03  2.74923798e-02 -4.80068475e-02  1.24806294e-03  5.48756123e-02  2.35883035e-02  3.57459188e-02 -4.40169685e-02  4.12164256e-03 -9.38005222e-04  3.08944564e-02 -3.77361737e-02  3.86108123e-02 -4.65151370e-02  4.00322862e-02  4.36247177e-02  1.25654312e-02  3.65298316e-02 -5.15242033e-02  6.86935848e-04 -1.80211291e-02  7.07198121e-03  3.63881961e-02  8.88750423e-03 -3.60064656e-02  4.73027118e-03  1.21512683e-02  4.76738885e-02 -2.15622112e-02 -1.24320751e-02  1.94059070e-02 -2.13006139e-02  1.13598257e-02  2.61519086e-02  9.84603260e-03 -1.64013896e-02 -1.93100479e-02  2.30080131e-02 -1.32036489e-02  5.43861650e-02 -1.16153499e-02 -3.13236304e-02 -1.45061398e-02  2.94751916e-02 -1.38078835e-02  5.12982905e-02 -2.20459350e-03  3.09881661e-02  2.04888023e-02 -2.62643043e-02 -3.10497861e-02  4.82752323e-02  1.58509221e-02 -2.15131547e-02 -1.94505360e-02 -3.42750214e-02  6.33830577e-02 -3.57603319e-02 -4.11630096e-03  4.74693738e-02 -3.78093943e-02  5.09781614e-02 -3.82447653e-02 -2.40907772e-03  1.90645475e-02  6.50643632e-02  1.59391519e-02 -8.09784178e-05  4.60834727e-02 -4.62939367e-02 -2.49499716e-02  2.88693293e-04 -2.82553927e-04  7.61555834e-03  2.04344429e-02  2.38504577e-02  2.42455732e-02  8.00277665e-02 -1.37750683e-02 -3.03374305e-02  1.77873555e-03  2.63471901e-02 -4.06999551e-02  3.81079130e-02  3.48770805e-02  4.41873074e-02 -4.42135632e-02 -1.02947867e-02  1.80779174e-02 -2.51882542e-02 -3.17575573e-03 -2.11982168e-02 -3.83161269e-02 -2.29231492e-02 -7.23808780e-02  8.27296078e-03  8.93022772e-03  3.15006711e-02  6.28977194e-02  4.93034534e-02 -6.31357217e-03  2.72838883e-02  4.60004956e-02  2.29189713e-02 -1.77909452e-02  1.75746810e-02 -2.22696662e-02 -3.60655338e-02 -1.43553847e-02 -1.68754514e-02  3.65821794e-02  4.14746366e-02  3.47368047e-02 -1.77744087e-02  5.09160198e-02 -4.59251463e-01 -3.97156179e-02 -3.06934752e-02  1.76206650e-03  5.17813079e-02 -1.09907277e-02  5.39624058e-02 -3.07337139e-02  5.53303994e-02  2.44358815e-02  2.76747923e-02 -2.01415978e-02  2.25076042e-02  1.18354959e-02  7.41290161e-03  3.65038924e-02  2.69415379e-02 -5.16565926e-02  1.66550446e-02  2.00316845e-03  1.28739802e-02 -4.20524403e-02 -7.66411200e-02  6.16761371e-02  5.69906493e-04 -4.32366179e-03 -4.70126681e-02  2.50594574e-03 -5.90763837e-02 -1.52328052e-02 -1.65860094e-02 -2.46960241e-02 -1.07343765e-02 -3.14804167e-02 -1.64745077e-02 -2.80346535e-02  3.72149423e-02 -1.22913113e-02  6.46753842e-03 -1.82984900e-02 -2.92817038e-02 -2.49831006e-02 -1.99497733e-02 -3.93009465e-03  1.15054790e-02 -1.46775674e-02 -6.52187467e-02 -5.72870299e-03 -5.61502017e-03  1.90818496e-02 -2.38545649e-02  1.48982145e-02  4.93834168e-03 -1.60587803e-02  2.14631855e-02 -3.71806882e-03  2.56405082e-02  3.26585323e-02 -1.06972782e-02  6.79324847e-03  2.07224376e-02 -4.68575619e-02 -1.35148233e-02 -6.44337153e-03  3.60724181e-02 -1.03052873e-02  1.15047991e-02 -4.77374904e-02  1.54225426e-02 -9.66934673e-03 -6.17276393e-02  2.14356650e-02  7.96907488e-03 -7.89307654e-02  2.23150011e-02 -1.49256196e-02  1.47885261e-02  3.62548530e-02 -3.86935771e-02 -4.00781780e-02 -3.38700376e-02 -4.97177765e-02  4.87275384e-02  1.30576489e-03  2.18466986e-02  3.46562862e-02  3.72738764e-03 -2.81844288e-02  1.81895401e-02  3.67801785e-02  6.97180480e-02 -3.22017260e-02  2.38486733e-02  2.47704554e-02  5.43326372e-04  4.35124040e-02 -4.09716815e-02  4.34940867e-02  6.42509162e-02  2.76623871e-02 -1.73324011e-02  1.89179395e-04  4.75937594e-03 -8.59102793e-03  1.68946553e-02  3.64033058e-02 -3.44986022e-02  2.00884864e-02  2.42669564e-02  4.93530463e-03 -1.92682352e-02  2.48522982e-02 -2.64476612e-02 -7.58881215e-03  4.98129204e-02  1.61886457e-02  2.81854812e-02  8.02835524e-02  6.74304441e-02  9.44981351e-02 -4.74918596e-02  2.42031310e-02 -3.03716380e-02 -8.88713263e-03 -6.93832943e-03 -2.71902103e-02 -6.31242618e-03 -7.06999702e-03  3.70524563e-02 -4.96705202e-03  5.72625473e-02  5.06692845e-03 -1.34559674e-02  8.89257621e-03  9.59218387e-03 -4.23642620e-02 -4.58639953e-03 -6.98388219e-02 -1.37693165e-02 -4.60867723e-03 -7.41797164e-02 -3.28920456e-03  1.62249207e-02  3.48545797e-02  8.05838127e-03 -1.95881091e-02 -2.70744953e-02  2.67664297e-03 -2.20929142e-02  4.31022570e-02  3.83416079e-02  1.23970292e-03  7.80535489e-02  2.32641175e-02 -2.54782531e-02  5.63883549e-03 -2.32594926e-02  3.72168981e-02 -5.09582721e-02  1.98740624e-02 -2.55218484e-02 -6.44714907e-02 -4.58236597e-02 -3.95013858e-03  1.87514853e-02  2.11250912e-02 -8.73581599e-03 -3.40146758e-02  2.84994096e-02  3.03051043e-02  1.25740021e-02 -1.98738240e-02 -2.92037986e-02 -2.74385680e-02  5.82823157e-02  2.93486863e-02 -1.11864880e-02  4.12185825e-02 -1.28866965e-02 -2.63788849e-02 -1.74545255e-02  6.20517088e-03 -3.45682390e-02 -1.90633778e-02  2.86185015e-02  6.41741650e-03 -2.66466085e-02 -8.28558505e-02 -5.58461100e-02  1.94943114e-03  2.22886205e-02 -5.28162252e-03 -3.59157361e-02  4.36805822e-02 -1.03601292e-02 -3.41477282e-02  3.78927328e-02 -3.95827852e-02  1.85145773e-02 -1.83031149e-02 -7.41270278e-03  1.23769892e-02 -4.21715081e-02  2.15793680e-03 -1.01877302e-02 -4.79500704e-02  2.43456401e-02  2.67727971e-02  1.18140867e-02 -8.43158364e-03  6.25065528e-03 -2.97430456e-02 -4.74497378e-02  3.43840085e-02  1.49865113e-02 -1.62406750e-02  2.72465181e-02 -1.19435210e-02  1.31491264e-02 -1.39700733e-02  2.07126718e-02 -3.90679799e-02  2.44531110e-02  3.25472280e-02  3.86319831e-02  3.81677710e-02  7.98764639e-03 -1.78932603e-02  3.20394188e-02 -2.64742319e-02  2.60311682e-02  5.13408296e-02  1.22884428e-02  5.96662834e-02 -5.99489324e-02 -1.25901997e-02  2.27843728e-02  3.86472568e-02 -1.55183561e-02  3.17636468e-02  2.44503692e-02  1.99540220e-02 -3.16410698e-03 -3.36151980e-02  5.35080358e-02 -3.42061045e-03 -2.81536877e-02 -4.64856513e-02 -5.55966224e-04 -1.98756903e-02  3.90093066e-02 -1.07923662e-02 -8.70320667e-03 -1.30228866e-02  6.84209093e-02 -1.28619699e-02 -1.16327731e-02 -4.11020555e-02 -2.93337107e-02 -2.64872294e-02  4.47598426e-03 -3.25084068e-02 -1.12765692e-02 -2.44318284e-02 -1.09370202e-02 -1.61682721e-02 -1.85636822e-02 -2.62791514e-02  5.60284033e-03 -3.40544581e-02 -5.80023974e-02 -6.45743832e-02  2.62594670e-02  3.47176492e-02 -2.71486640e-02  3.94119658e-02 -5.33050895e-02  1.02327866e-02 -9.66343563e-03  1.07219182e-02 -7.41349757e-02 -1.63307041e-02 -2.53311475e-03  2.39288546e-02  3.29732858e-02 -2.11625919e-02 -1.36818849e-02  8.72594584e-03  2.71956436e-02  6.76753582e-04 -4.66498882e-02  3.62443477e-02 -1.49152977e-02  1.77372228e-02 -3.73987630e-02 -1.33736515e-02  5.55564761e-02  4.98019680e-02 -5.46147395e-03 -2.33486816e-02  1.26596969e-02  2.51028948e-02  4.17576134e-02  3.90054844e-02 -2.24296246e-02  4.41033244e-02  9.14904568e-03 -7.53453560e-03 -4.29231189e-02  3.02391462e-02  2.07852758e-02 -1.02980128e-02  5.25167845e-02 -8.15811916e-04 -5.33150658e-02  1.34045607e-03  5.91069683e-02  2.30129412e-03 -1.37382932e-02  4.49629268e-04  1.30582890e-02 -7.47572957e-03  1.75305996e-02  1.02643510e-02  7.59967486e-04  4.58112583e-02  4.10522521e-02  1.25804953e-02 -7.41753727e-03  4.82532149e-03  2.78516840e-02  2.36507114e-02 -4.13814187e-02 -1.29771689e-02  1.67835839e-02  1.53311454e-02  1.64327212e-02 -1.37600424e-02  1.23005807e-02  5.79211973e-02  7.21627381e-03  6.61403313e-02 -5.89747727e-02  1.59436222e-02 -2.17061862e-03  5.44503778e-02  1.95026603e-02 -2.54723858e-02  3.68227047e-04  6.43685311e-02 -2.19820179e-02 -1.01059081e-03  2.62502227e-02 -3.33848931e-02  8.37497320e-03  7.36760795e-02  4.29018140e-02  1.57072600e-02 -2.66952831e-02  9.60265379e-03  1.46882059e-02 -5.75055135e-03  4.36292887e-02 -1.57237183e-02 -2.38313824e-02 -5.69211552e-03  1.43865235e-02 -8.47134143e-02 -2.65024845e-02 -9.16605629e-03  2.04189331e-03  1.62579492e-03 -6.75124824e-02 -1.84973944e-02  3.43609378e-02  1.11558558e-02  7.81714320e-02 -5.49522275e-03 -1.51996361e-02  4.76938970e-02  1.68492533e-02 -3.04164756e-02 -1.65272392e-02 -7.19350874e-02 -3.92437950e-02  4.13833447e-02 -3.05935973e-03 -1.75968558e-02 -1.19854569e-01  1.87387038e-02  1.21218693e-02 -2.49273274e-02  3.28705758e-02  4.41757403e-02  6.98454585e-03  8.09968915e-03 -3.59516442e-02  6.08325787e-02 -2.27745436e-02  2.11376194e-02  3.23967868e-03  2.22185068e-02  2.37153284e-02 -1.77512988e-02  2.61204340e-03  7.71679450e-03 -3.08509786e-02 -8.34045187e-02 -6.29286990e-02  4.29803273e-03  2.88508553e-02 -1.15858456e-02  2.17792168e-02 -4.04994898e-02 -1.84838306e-02 -3.33074331e-02  2.66448911e-02 -4.43037413e-02  1.92902349e-02 -8.74877535e-03  1.84421762e-04 -2.02794485e-02 -4.32577468e-02  4.79730107e-02 -2.86451634e-02  1.03001520e-02  3.35061178e-02 -3.62282409e-03 -4.65558060e-02  3.37383077e-02  3.71970534e-02 -6.26126453e-02 -3.27363126e-02 -4.46447171e-02 -1.11249518e-02 -6.53849989e-02  3.09160519e-02  7.00028688e-02  4.30897884e-02 -6.10660948e-02 -1.23231001e-02  7.72167882e-03 -1.32871531e-02  1.33951856e-02 -4.43044752e-02  4.29534987e-02 -5.41009791e-02 -4.23156843e-02 -1.02046803e-02  6.48994818e-02  2.56961752e-02  7.90675916e-03 -1.38908927e-03  2.88802991e-03  3.53993066e-02 -1.85436178e-02  2.64718328e-02 -4.30478230e-02  6.46922365e-02 -3.79009321e-02]\n"
     ]
    }
   ],
   "source": [
    "# embeds = np.vstack(embeds)\n",
    "print(embeds[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use faiss to cluster the embeddings\n",
    "import faiss\n",
    "import numpy as np\n",
    "\n",
    "ncentroids = 100\n",
    "niter = 300\n",
    "verbose = True\n",
    "gpu = True\n",
    "dim = 768\n",
    "kmeans = faiss.Kmeans(dim, ncentroids, niter = niter, verbose = verbose, \n",
    "                      min_points_per_centroid = 100, max_points_per_centroid = 100000, \n",
    "                      nredo = 2, spherical = True)\n",
    "\n",
    "kmeans.train(embeds)\n",
    "\n",
    "faiss.write_index(kmeans.index, \"./Data/Data_6/KMeansFAISS/index.bin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mapping each point to its nearest centroid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import faiss\n",
    "import numpy as np\n",
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "\n",
    "\n",
    "# read the faiss index\n",
    "index = faiss.read_index(\"./Data/Data_6/KMeansFAISS/index.bin\")\n",
    "\n",
    "# mapping each point to its nearest centroid\n",
    "_, I = index.search(embeds, 2)\n",
    "\n",
    "# flatten the array of arrays\n",
    "I_flat = [item for sublist in I for item in sublist]\n",
    "\n",
    "# write the labels to a file where I_flat is the centroid index for each embedding\n",
    "with open(\"./Data/Data_6/labels/faiss_labels.txt\", \"a\") as f:\n",
    "    # remove the string array representation of the list\n",
    "    f.write(str(I_flat).replace(\"[\", \"\").replace(\"]\", \"\").replace(\",\", \"\").replace(\"  \", \" \").strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting the nearest clusters to the given query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read queries and compute nearest cloisters\n",
    "import faiss \n",
    "import numpy as np\n",
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "\n",
    "# read the faiss index\n",
    "index = faiss.read_index(\"./Data/Data_6/KMeansFAISS/index.bin\")\n",
    "\n",
    "with open(\"./Data/query.txt\", \"r\") as f:\n",
    "    lines = [np.fromstring(line, sep=\" \") for line in f]\n",
    "    queries = np.array(lines)\n",
    "\n",
    "#Normalize the query\n",
    "queries = queries / np.linalg.norm(queries)\n",
    "\n",
    "_, I = index.search(queries, 10)\n",
    "I_flat = [i for i in I]\n",
    "I_flat = np.array(I_flat)\n",
    "\n",
    "# write the labels to a file where I_flat is the centroid index for each embedding\n",
    "with open(\"./Data/Data_6/labels/query_faiss_labels.txt\", \"w\") as f:\n",
    "    f.write(str(I_flat).replace(\"[\", \"\").replace(\"]\", \"\").replace(\",\", \"\").replace(\"  \", \" \").strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the labels file, where each item from dataset (in batches) is mapped to its closest two centroids\n",
    "Write the centroids to a file\n",
    "\n",
    "For each centroid file, we need to manually remove square brackets \"[\" or \"]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "centroids_filename = \"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/centroids/final-centroids.txt\"\n",
    "with open(centroids_filename, \"a\") as f:\n",
    "    # print the centroids line by line\n",
    "    # make the printing so that it execludes any squared brackets or double spacing also remove that first space \n",
    "    # also remove all the commas\n",
    "    # also remove all the white spaces at the very beginig and end of each line\n",
    "\n",
    "    for i in range(ncentroids):\n",
    "        f.write(str(kmeans.centroids[i]).replace(\"[\", \"\").replace(\"]\", \"\").replace(\"  \", \" \").strip())\n",
    "        f.write(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each label file, we need to manually remove double spaces \"  \" and any square brackets \"[\" or \"]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def assign_labels(X, centroids):\n",
    "    # // make this of dimension x.shape[0][num]\n",
    "    assigned_custers = 2\n",
    "    labels = np.empty(X.shape[0]*assigned_custers, dtype=np.int64)\n",
    "    for i, x in enumerate(X):\n",
    "        # L2\n",
    "        # distances = np.linalg.norm(x - centroids, axis=1)\n",
    "\n",
    "        # cosine similarity\n",
    "        # distances = np.dot(x, centroids.T) / (np.linalg.norm(x) * np.linalg.norm(centroids))\n",
    "\n",
    "        # use inner product to calculate the distances\n",
    "        distances = np.inner(centroids, x)\n",
    "\n",
    "        sorted_indices = np.argsort(distances)\n",
    "        min_index = sorted_indices[0]\n",
    "        second_min_index = sorted_indices[1]\n",
    "        labels[2 * i] = min_index\n",
    "        labels[2*i + 1] = second_min_index\n",
    "    return labels\n",
    "\n",
    "import numpy as np\n",
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "centroids_filename = \"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/centroids/final-centroids.txt\"\n",
    "\n",
    "# #print the norm of the centroids\n",
    "# centroids = np.loadtxt(centroids_filename)\n",
    "# for i in range(100):\n",
    "#     print(np.linalg.norm(centroids[i]))\n",
    "\n",
    "\n",
    "for i in range(5):\n",
    "    with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/dataset/{i}-embeds-batch.txt\", \"r\") as f:\n",
    "        # read each line and convert to numpy array\n",
    "        lines = [np.fromstring(line, sep=\" \") for line in f]\n",
    "        # convert to numpy array\n",
    "        embeds = np.array(lines)\n",
    "        # print(embeds.shape)\n",
    "        centroids = np.loadtxt(centroids_filename)\n",
    "        # #normalize the centroids\n",
    "        # centroids = centroids / np.linalg.norm(centroids, axis=1)[:, np.newaxis]\n",
    "        # #normalize the embeddings\n",
    "        # embeds = embeds / np.linalg.norm(embeds, axis=1)[:, np.newaxis]\n",
    "        labels = assign_labels(embeds, centroids)\n",
    "        with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/labels/labels{i}.txt\", \"a\") as f:\n",
    "            f.write(str(labels).replace('[', '').replace(']', '').replace('  ', ' ').replace('  ', ' ').strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attempts\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KD tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Constructing KD tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neighbors import KDTree\n",
    "from scipy.spatial import cKDTree\n",
    "tree = KDTree(embeds) # leafSize = 16\n",
    "# tree = cKDTree(embeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_data, index, tree_nodes, node_bounds = tree.get_arrays()\n",
    "tree_nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## K means"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running K means on each batch\n",
    "- 100k batches\n",
    "- each batch 100 clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loop on folders\n",
    "import os\n",
    "import numpy as np\n",
    "from sklearn.cluster import KMeans\n",
    "\n",
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "\n",
    "with open(\"./cmake-build-debug/batches/1-embeds-batch.txt\", \"r\") as f:\n",
    "    # read each line and convert to numpy array\n",
    "    lines = [np.fromstring(line, sep=\" \") for line in f]\n",
    "    # convert to numpy array\n",
    "    embeds = np.array(lines)\n",
    "    # apply kmeans on embeds and get the centroids\n",
    "    kmeans = KMeans(n_clusters = 100).fit(embeds)\n",
    "    with open(\"./cmake-build-debug/batches/centroids1.txt\", \"a\") as f:\n",
    "        f.write(str(kmeans.cluster_centers_))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Collecting all centroid files into one array\n",
    "* Make sure that centroids files don't have any squared brackets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "centroids = np.loadtxt(\"./cmake-build-debug/batches/centroids0.txt\")\n",
    "centroids = np.concatenate((centroids, np.loadtxt(\"./cmake-build-debug/batches/centroids1.txt\")), axis=0)\n",
    "centroids = np.concatenate((centroids, np.loadtxt(\"./cmake-build-debug/batches/centroids2.txt\")), axis=0)\n",
    "centroids = np.concatenate((centroids, np.loadtxt(\"./cmake-build-debug/batches/centroids3.txt\")), axis=0)\n",
    "centroids = np.concatenate((centroids, np.loadtxt(\"./cmake-build-debug/batches/centroids4.txt\")), axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Applying K-means on all the centroids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply kmeans on embeds and get the centroids\n",
    "kmeans = KMeans(n_clusters = 100).fit(centroids)\n",
    "with open(\"./cmake-build-debug/batches/final-centroids.txt\", \"a\") as f:\n",
    "    f.write(str(kmeans.cluster_centers_))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assigning labels to each batch according to the final centroids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def assign_labels(X, centroids):\n",
    "    # // make this of dimension x.shape[0][num]\n",
    "    assigned_custers = 2\n",
    "    labels = np.empty(X.shape[0]*assigned_custers, dtype=np.int64)\n",
    "    for i, x in enumerate(X):\n",
    "        distances = np.linalg.norm(x - centroids, axis=1)\n",
    "        sorted_indices = np.argsort(distances)\n",
    "        min_index = sorted_indices[0]\n",
    "        second_min_index = sorted_indices[1]\n",
    "        labels[2 * i] = min_index\n",
    "        labels[2*i + 1] = second_min_index\n",
    "    return labels\n",
    "\n",
    "import numpy as np\n",
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "\n",
    "\n",
    "for i in range(5):\n",
    "    with open(f\"./cmake-build-debug/batches/dataset/{i}-embeds-batch.txt\", \"r\") as f:\n",
    "        # read each line and convert to numpy array\n",
    "        lines = [np.fromstring(line, sep=\" \") for line in f]\n",
    "        # convert to numpy array\n",
    "        embeds = np.array(lines)\n",
    "        # print(embeds.shape)\n",
    "        centroids = np.loadtxt(\"./cmake-build-debug/batches/centroids/final-centroids.txt\")\n",
    "        labels = assign_labels(embeds, centroids)\n",
    "        with open(f\"./cmake-build-debug/batches/labels/labels{i}.txt\", \"a\") as f:\n",
    "            f.write(str(labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Constructing clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# opening the embeds file\n",
    "with open(\"./cmake-build-debug/batches/dataset/1-embeds-batch.txt\", \"r\") as f:\n",
    "    lines = [np.fromstring(line, sep=\" \") for line in f]\n",
    "    embeds = np.array(lines)\n",
    "\n",
    "# read the labels and convert it to a list\n",
    "with open(\"./cmake-build-debug/batches/labels/labels1.txt\", \"r\") as f:\n",
    "    labels =[int(x) for x in next(f).split()]   \n",
    "\n",
    "# loop on the labels \n",
    "for i in range(100):\n",
    "    indices = [j//2 for j, x in enumerate(labels) if x == i]\n",
    "    current_embeds = embeds[indices]\n",
    "    with open(f\"./cmake-build-debug/batches/clusters/{i}.txt\", \"a\") as f:\n",
    "        f.write(str(current_embeds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(100):\n",
    "    with open(f\"./cmake-build-debug/batches/clusters/cluster{i}.txt\", \"a\") as f:\n",
    "        f.write(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clustering using Faiss IVF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reading the dataset from the pre-downloaded data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read the dataset from the files in /batches/dataset in one big numpy array\n",
    "import numpy as np\n",
    "\n",
    "embeds = []\n",
    "for i in range(5):\n",
    "    with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/dataset/{i}-embeds-batch.txt\", \"r\") as f:\n",
    "        lines = [np.fromstring(line, sep=\" \") for line in f]\n",
    "        # normalize the embeddings\n",
    "        # lines = [line / np.linalg.norm(line) for line in lines]\n",
    "        embeds.append(np.array(lines))\n",
    "\n",
    "# Concatenate all lists into a\n",
    "embeds = np.concatenate(embeds, axis=0)\n",
    "\n",
    "# print the norm of the first 100 embeddings\n",
    "for i in range(100):\n",
    "    print(np.linalg.norm(embeds[i]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Constructing the clusters using Faiss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install faiss-cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use faiss to cluster the embeddings\n",
    "import faiss\n",
    "\n",
    "# create an inverted index\n",
    "nlist = 256\n",
    "m = 1\n",
    "k = 5\n",
    "d = 768\n",
    "coarse_quantizer = faiss.IndexFlatL2(d)\n",
    "index = faiss.IndexIVFPQ(coarse_quantizer, d, nlist, m, 8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.train(embeds)\n",
    "index.add(embeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if there is a pre-transform, you can also use\n",
    "# invlists = faiss.extract_index_ivf(index).invlists\n",
    "invlists = index.invlists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_invlist(invlists, l):\n",
    "    \"\"\" returns the inverted lists content as a pair of (list_ids, list_codes).\n",
    "    The codes are reshaped to a proper size\n",
    "    \"\"\"\n",
    "    invlists = faiss.downcast_InvertedLists(invlists)\n",
    "    ls = invlists.list_size(l)\n",
    "    list_ids = np.zeros(ls, dtype='int64')\n",
    "    ids = codes = None\n",
    "    try:\n",
    "        ids = invlists.get_ids(l)\n",
    "        if ls > 0:\n",
    "            faiss.memcpy(faiss.swig_ptr(list_ids), ids, list_ids.nbytes)\n",
    "        codes = invlists.get_codes(l)\n",
    "        if invlists.code_size != faiss.InvertedLists.INVALID_CODE_SIZE:\n",
    "            list_codes = np.zeros((ls, invlists.code_size), dtype='uint8')\n",
    "        else:\n",
    "            # it's a BlockInvertedLists\n",
    "            npb = invlists.n_per_block\n",
    "            bs = invlists.block_size\n",
    "            ls_round = (ls + npb - 1) // npb\n",
    "            list_codes = np.zeros((ls_round, bs // npb, npb), dtype='uint8')\n",
    "        if ls > 0:\n",
    "            faiss.memcpy(faiss.swig_ptr(list_codes), codes, list_codes.nbytes)\n",
    "    finally:\n",
    "        if ids is not None:\n",
    "            invlists.release_ids(l, ids)\n",
    "        if codes is not None:\n",
    "            invlists.release_codes(l, codes)\n",
    "    return list_ids, list_codes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get content of inverted list #123\n",
    "list_ids, list_codes = get_invlist(invlists, 255)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "np.set_printoptions(threshold=np.inf, linewidth=np.inf)  # ensure numpy array won't truncate\n",
    "\n",
    "# use the get_invlist function to get the list_ids and save them the the labels files each in a file indexed from 0 to 255\n",
    "for i in range(256):\n",
    "    list_ids, list_codes = get_invlist(invlists, i)\n",
    "    with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/labels/labels{i}.txt\", \"a\") as f:\n",
    "        # make the printing so that it execludes any squared brackets or double spacing also remove that first space \n",
    "        # also remove all the white spaces at the very beginig and end of the string\n",
    "        f.write(str(list_ids).replace('[', '').replace(']', '').replace('  ', ' ').replace('  ', ' ').strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use list_ids to map each vector from embeds to its corresponding label using the indices in list_ids and output the clusters to the clusters folder\n",
    "for i in range(256):\n",
    "    with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/clusters/{i}.txt\", \"a\") as f:\n",
    "\n",
    "        with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/labels/labels{i}.txt\", \"r\") as f2:\n",
    "            labels =[int(x) for x in next(f2).split()]   \n",
    "\n",
    "        for j in range(len(labels)):\n",
    "            # make the printing so that it execludes any squared brackets or double spacing also remove that first space \n",
    "            # also remove all the white spaces at the very beginig and end of the string\n",
    "            f.write(str(embeds[labels[j]]).replace('[', '').replace(']', '').replace('  ', ' ').replace('  ', ' ').strip())\n",
    "            f.write(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read the clusters from the folder file by file and calculate the medoid for each cluster file and output the medoids to one file in centroids folder\n",
    "import numpy as np\n",
    "\n",
    "for i in range(256):\n",
    "    with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/clusters/{i}.txt\", \"r\") as f:\n",
    "        lines = [np.fromstring(line, sep=\" \") for line in f]\n",
    "        embeds = np.array(lines)\n",
    "        medoid = embeds[np.argmin(np.sum(np.abs(embeds[:, np.newaxis] - embeds), axis=2), 0)]\n",
    "        with open(f\"D:/Boody/GP/Indexer/RAGn-Roll-Indexer/Data/centroids/test-centroids.txt\", \"a\") as f:\n",
    "            f.write(str(medoid).replace('[', '').replace(']', '').replace('  ', ' ').replace('  ', ' ').strip())\n",
    "            f.write(\"\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}