examples/notebooks/rag/rag_1B_load_data_into_milvus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Processed Data into Vector Database\n",
    "\n",
    "This notebook loads output from data prep kit into Milvus\n",
    "\n",
    "**Step-4 in this workflow**\n",
    "\n",
    "![](../media/rag-overview-2.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-1: Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from my_config import MY_CONFIG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-2: Load Parquet Data\n",
    "\n",
    "Load all  `.parquet` files in the given dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data from :  output/output_final\n",
      "Number of parquet files to read :  2\n",
      "\n",
      "Read file: 'output/output_final/granite.parquet'.  number of rows = 123\n",
      "Read file: 'output/output_final/attension.parquet'.  number of rows = 88\n",
      "\n",
      "Total number of rows = 211\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import glob\n",
    "\n",
    "print ('Loading data from : ', MY_CONFIG.OUTPUT_FOLDER_FINAL)\n",
    "\n",
    "# Get a list of all Parquet files in the directory\n",
    "parquet_files = glob.glob(f'{MY_CONFIG.OUTPUT_FOLDER_FINAL}/*.parquet')\n",
    "print (\"Number of parquet files to read : \", len(parquet_files))\n",
    "print ()\n",
    "\n",
    "# Create an empty list to store the DataFrames\n",
    "dfs = []\n",
    "\n",
    "# Loop through each Parquet file and read it into a DataFrame\n",
    "for file in parquet_files:\n",
    "    df = pd.read_parquet(file)\n",
    "    print (f\"Read file: '{file}'.  number of rows = {df.shape[0]}\")\n",
    "    dfs.append(df)\n",
    "\n",
    "# Concatenate all DataFrames into a single DataFrame\n",
    "data_df = pd.concat(dfs, ignore_index=True)\n",
    "\n",
    "print (f\"\\nTotal number of rows = {data_df.shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding length:  384\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 211 entries, 0 to 210\n",
      "Data columns (total 20 columns):\n",
      " #   Column              Non-Null Count  Dtype  \n",
      "---  ------              --------------  -----  \n",
      " 0   filename            211 non-null    object \n",
      " 1   num_pages           211 non-null    int64  \n",
      " 2   num_tables          211 non-null    int64  \n",
      " 3   num_doc_elements    211 non-null    int64  \n",
      " 4   ext                 211 non-null    object \n",
      " 5   hash                211 non-null    object \n",
      " 6   size                211 non-null    int64  \n",
      " 7   date_acquired       211 non-null    object \n",
      " 8   pdf_convert_time    211 non-null    float64\n",
      " 9   source_filename     211 non-null    object \n",
      " 10  source_document_id  211 non-null    object \n",
      " 11  text                211 non-null    object \n",
      " 12  doc_jsonpath        211 non-null    object \n",
      " 13  page_number         211 non-null    int64  \n",
      " 14  bbox                211 non-null    object \n",
      " 15  document_id         211 non-null    object \n",
      " 16  chunk_id            211 non-null    int64  \n",
      " 17  removed             211 non-null    object \n",
      " 18  chunk_hash          211 non-null    int64  \n",
      " 19  vector              211 non-null    object \n",
      "dtypes: float64(1), int64(7), object(12)\n",
      "memory usage: 33.1+ KB\n",
      "None\n"
     ]
    },
    {
     "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>filename</th>\n",
       "      <th>num_pages</th>\n",
       "      <th>num_tables</th>\n",
       "      <th>num_doc_elements</th>\n",
       "      <th>ext</th>\n",
       "      <th>hash</th>\n",
       "      <th>size</th>\n",
       "      <th>date_acquired</th>\n",
       "      <th>pdf_convert_time</th>\n",
       "      <th>source_filename</th>\n",
       "      <th>source_document_id</th>\n",
       "      <th>text</th>\n",
       "      <th>doc_jsonpath</th>\n",
       "      <th>page_number</th>\n",
       "      <th>bbox</th>\n",
       "      <th>document_id</th>\n",
       "      <th>chunk_id</th>\n",
       "      <th>removed</th>\n",
       "      <th>chunk_hash</th>\n",
       "      <th>vector</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>28</td>\n",
       "      <td>17</td>\n",
       "      <td>348</td>\n",
       "      <td>pdf</td>\n",
       "      <td>79c53d694df467391e94f279af2fa6a9a7e45c3922546e...</td>\n",
       "      <td>655054</td>\n",
       "      <td>2024-10-02T00:28:23.836369</td>\n",
       "      <td>167.768806</td>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>81bc331a-69cf-49bd-84b9-afedcab1344a</td>\n",
       "      <td>Granite Code Models: A Family of Open Foundati...</td>\n",
       "      <td>$.main-text[3]</td>\n",
       "      <td>1</td>\n",
       "      <td>[142.70646667, 672.96929932, 468.58251953, 711...</td>\n",
       "      <td>b773445f7cf4cc9a5bf6ec296c74504f93c9c179028ac6...</td>\n",
       "      <td>88</td>\n",
       "      <td>[]</td>\n",
       "      <td>-1</td>\n",
       "      <td>[-0.015789315, -0.07841933, -0.032271657, 0.00...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>28</td>\n",
       "      <td>17</td>\n",
       "      <td>348</td>\n",
       "      <td>pdf</td>\n",
       "      <td>79c53d694df467391e94f279af2fa6a9a7e45c3922546e...</td>\n",
       "      <td>655054</td>\n",
       "      <td>2024-10-02T00:28:23.836369</td>\n",
       "      <td>167.768806</td>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>81bc331a-69cf-49bd-84b9-afedcab1344a</td>\n",
       "      <td>Granite Code Models: A Family of Open Foundati...</td>\n",
       "      <td>$.main-text[4]</td>\n",
       "      <td>1</td>\n",
       "      <td>[107.61845398, 535.62896729, 503.99923706, 647...</td>\n",
       "      <td>7353bcc8d99c279335eaf120c793ca6a08f9a4fddcbb5b...</td>\n",
       "      <td>89</td>\n",
       "      <td>[]</td>\n",
       "      <td>-1</td>\n",
       "      <td>[-0.059480786, -0.056680508, -0.042864937, -0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>28</td>\n",
       "      <td>17</td>\n",
       "      <td>348</td>\n",
       "      <td>pdf</td>\n",
       "      <td>79c53d694df467391e94f279af2fa6a9a7e45c3922546e...</td>\n",
       "      <td>655054</td>\n",
       "      <td>2024-10-02T00:28:23.836369</td>\n",
       "      <td>167.768806</td>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>81bc331a-69cf-49bd-84b9-afedcab1344a</td>\n",
       "      <td>Granite Code Models: A Family of Open Foundati...</td>\n",
       "      <td>$.main-text[5]</td>\n",
       "      <td>1</td>\n",
       "      <td>[220.87228394, 484.46414185, 390.87872314, 529...</td>\n",
       "      <td>389267895ca214924a0a071df8379c2b15fcf374f232a6...</td>\n",
       "      <td>90</td>\n",
       "      <td>[]</td>\n",
       "      <td>-1</td>\n",
       "      <td>[-0.07557265, -0.07152908, -0.048923455, -0.04...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      filename  num_pages  num_tables  num_doc_elements  ext  \\\n",
       "0  granite.pdf         28          17               348  pdf   \n",
       "1  granite.pdf         28          17               348  pdf   \n",
       "2  granite.pdf         28          17               348  pdf   \n",
       "\n",
       "                                                hash    size  \\\n",
       "0  79c53d694df467391e94f279af2fa6a9a7e45c3922546e...  655054   \n",
       "1  79c53d694df467391e94f279af2fa6a9a7e45c3922546e...  655054   \n",
       "2  79c53d694df467391e94f279af2fa6a9a7e45c3922546e...  655054   \n",
       "\n",
       "                date_acquired  pdf_convert_time source_filename  \\\n",
       "0  2024-10-02T00:28:23.836369        167.768806     granite.pdf   \n",
       "1  2024-10-02T00:28:23.836369        167.768806     granite.pdf   \n",
       "2  2024-10-02T00:28:23.836369        167.768806     granite.pdf   \n",
       "\n",
       "                     source_document_id  \\\n",
       "0  81bc331a-69cf-49bd-84b9-afedcab1344a   \n",
       "1  81bc331a-69cf-49bd-84b9-afedcab1344a   \n",
       "2  81bc331a-69cf-49bd-84b9-afedcab1344a   \n",
       "\n",
       "                                                text    doc_jsonpath  \\\n",
       "0  Granite Code Models: A Family of Open Foundati...  $.main-text[3]   \n",
       "1  Granite Code Models: A Family of Open Foundati...  $.main-text[4]   \n",
       "2  Granite Code Models: A Family of Open Foundati...  $.main-text[5]   \n",
       "\n",
       "   page_number                                               bbox  \\\n",
       "0            1  [142.70646667, 672.96929932, 468.58251953, 711...   \n",
       "1            1  [107.61845398, 535.62896729, 503.99923706, 647...   \n",
       "2            1  [220.87228394, 484.46414185, 390.87872314, 529...   \n",
       "\n",
       "                                         document_id  chunk_id removed  \\\n",
       "0  b773445f7cf4cc9a5bf6ec296c74504f93c9c179028ac6...        88      []   \n",
       "1  7353bcc8d99c279335eaf120c793ca6a08f9a4fddcbb5b...        89      []   \n",
       "2  389267895ca214924a0a071df8379c2b15fcf374f232a6...        90      []   \n",
       "\n",
       "   chunk_hash                                             vector  \n",
       "0          -1  [-0.015789315, -0.07841933, -0.032271657, 0.00...  \n",
       "1          -1  [-0.059480786, -0.056680508, -0.042864937, -0....  \n",
       "2          -1  [-0.07557265, -0.07152908, -0.048923455, -0.04...  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "## Shape the data\n",
    "\n",
    "MY_CONFIG.EMBEDDING_LENGTH =  len(data_df.iloc[0]['embeddings'])\n",
    "print ('embedding length: ', MY_CONFIG.EMBEDDING_LENGTH)\n",
    "\n",
    "# rename 'embeddings' columns as 'vector' to match default schema\n",
    "# if 'vector' not in data_df.columns and 'embeddings' in data_df.columns:\n",
    "#     data_df = data_df.rename( columns= {'embeddings' : 'vector'})\n",
    "# if 'text' not in data_df.columns and 'contents' in data_df.columns:\n",
    "#     data_df = data_df.rename( columns= {'contents' : 'text'})\n",
    "\n",
    "data_df = data_df.rename( columns= {'embeddings' : 'vector', 'contents' : 'text'})\n",
    "\n",
    "print (data_df.info())\n",
    "data_df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-3: Connect to Vector Database\n",
    "\n",
    "Milvus can be embedded and easy to use.\n",
    "\n",
    "<span style=\"color:blue;\">Note: If you encounter an error about unable to load database, try this: </span>\n",
    "\n",
    "- <span style=\"color:blue;\">In **vscode** : **restart the kernel** of previous notebook. This will release the db.lock </span>\n",
    "- <span style=\"color:blue;\">In **Jupyter**: Do `File --> Close and Shutdown Notebook` of previous notebook. This will release the db.lock</span>\n",
    "- <span style=\"color:blue;\">Re-run this cell again</span>\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Connected to Milvus instance: ./rag_1_dpk.db\n"
     ]
    }
   ],
   "source": [
    "from pymilvus import MilvusClient\n",
    "\n",
    "milvus_client = MilvusClient(MY_CONFIG.DB_URI)\n",
    "\n",
    "print (\"✅ Connected to Milvus instance:\", MY_CONFIG.DB_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step-4: Create A Collection\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Created collection : dpk_papers\n"
     ]
    }
   ],
   "source": [
    "# if we already have a collection, clear it first\n",
    "if milvus_client.has_collection(collection_name=MY_CONFIG.COLLECTION_NAME):\n",
    "    milvus_client.drop_collection(collection_name=MY_CONFIG.COLLECTION_NAME)\n",
    "    print ('✅ Cleared collection :', MY_CONFIG.COLLECTION_NAME)\n",
    "\n",
    "\n",
    "milvus_client.create_collection(\n",
    "    collection_name=MY_CONFIG.COLLECTION_NAME,\n",
    "    dimension=MY_CONFIG.EMBEDDING_LENGTH,\n",
    "    metric_type=\"IP\",  # Inner product distance\n",
    "    consistency_level=\"Strong\",  # Strong consistency level\n",
    "    auto_id=True\n",
    ")\n",
    "print (\"✅ Created collection :\", MY_CONFIG.COLLECTION_NAME)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inserted # rows 211\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'row_count': 211}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = milvus_client.insert(collection_name=MY_CONFIG.COLLECTION_NAME, data=data_df.to_dict('records'))\n",
    "\n",
    "print('inserted # rows', res['insert_count'])\n",
    "\n",
    "milvus_client.get_collection_stats(MY_CONFIG.COLLECTION_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-5: Close DB Connection\n",
    "\n",
    "Close the connection so the lock files are relinquished and other notebooks can access the db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ SUCCESS\n"
     ]
    }
   ],
   "source": [
    "milvus_client.close()\n",
    "\n",
    "print (\"✅ SUCCESS\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test your data by doing a Vector Search\n",
    "\n",
    "See notebook [vector_search.ipynb](vector_search.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}