examples/notebooks/rag/rag_2A_llamaindex_process.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG With llama-index  + Milvus + LLama\n",
    "\n",
    "References\n",
    "- https://docs.llamaindex.ai/en/stable/examples/vector_stores/MilvusIndexDemo/\n",
    "- https://docs.llamaindex.ai/en/stable/api_reference/storage/vector_store/milvus/?h=milvusvectorstore#llama_index.vector_stores.milvus.MilvusVectorStore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-1: Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from my_config import MY_CONFIG\n",
    "\n",
    "MY_CONFIG.DB_URI = './rag_2_llamaindex.db'\n",
    "MY_CONFIG.COLLECTION_NAME = 'llamaindex_papers'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-2: Read documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 43 chunks\n",
      "CPU times: user 3.9 s, sys: 869 ms, total: 4.77 s\n",
      "Wall time: 2.76 s\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "\n",
    "from llama_index.core import SimpleDirectoryReader\n",
    "import pprint\n",
    "\n",
    "# load documents\n",
    "documents = SimpleDirectoryReader(\n",
    "    # input_dir = './data/10k/input/'\n",
    "    input_dir = MY_CONFIG.INPUT_DATA_DIR\n",
    ").load_data()\n",
    "\n",
    "print (f\"Loaded {len(documents)} chunks\")\n",
    "\n",
    "# print(\"Document [0].doc_id:\", documents[0].doc_id)\n",
    "# pprint.pprint (documents[0], indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-3: Setup Embedding Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If connection to https://huggingface.co/ failed, uncomment the following path\n",
    "import os\n",
    "os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "08aefb8116a540678e28c78accd09648",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "modules.json:   0%|          | 0.00/349 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "930ad0af631d4f308978973d4bf84fcf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config_sentence_transformers.json:   0%|          | 0.00/116 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c46860ae2634b32b0564b3c3966447d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f2e4b4e0ba174ff9ae8542631fdd53dd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "sentence_bert_config.json:   0%|          | 0.00/53.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sujee/apps/anaconda3/envs/data-prep-kit-4-021/lib/python3.11/site-packages/huggingface_hub/file_download.py:1142: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2a6d591a5f564b11a84648d57f24a418",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/360 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a458eff1e8c46c9ab9eb4c26cdfac04",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/90.9M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b754e347a67a40008346c4d928b06b50",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/350 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "680021888456420abd1eecf438c880ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.txt: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d303a919b7d14ca68b13fd85747461ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2bab4bf697c843fd9dd636c793dba92a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/112 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac75b94e156b4562a9073889a137551d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "1_Pooling/config.json:   0%|          | 0.00/190 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.embed_model = HuggingFaceEmbedding(\n",
    "    model_name = MY_CONFIG.EMBEDDING_MODEL\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-4: Connect to Milvus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Connected to Milvus instance:  ./rag_2_llamaindex.db\n"
     ]
    }
   ],
   "source": [
    "from pymilvus import MilvusClient\n",
    "\n",
    "milvus_client = MilvusClient(MY_CONFIG.DB_URI)\n",
    "print (\"✅ Connected to Milvus instance: \", MY_CONFIG.DB_URI )\n",
    "\n",
    "# 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",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Connected Llama-index to Milvus instance:  ./rag_2_llamaindex.db\n"
     ]
    }
   ],
   "source": [
    "# connect to vector db\n",
    "from llama_index.core import VectorStoreIndex, StorageContext\n",
    "from llama_index.vector_stores.milvus import MilvusVectorStore\n",
    "\n",
    "vector_store = MilvusVectorStore(\n",
    "    uri = MY_CONFIG.DB_URI ,\n",
    "    dim = MY_CONFIG.EMBEDDING_LENGTH , \n",
    "    collection_name = MY_CONFIG.COLLECTION_NAME,\n",
    "    overwrite=True\n",
    ")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "print (\"✅ Connected Llama-index to Milvus instance: \", MY_CONFIG.DB_URI )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-5: Create Index and Save to DB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Created index: <llama_index.core.indices.vector_store.base.VectorStoreIndex object at 0x7ac2d9f4e450>\n",
      "✅ Saved index to db  ./rag_2_llamaindex.db\n",
      "CPU times: user 912 ms, sys: 155 ms, total: 1.07 s\n",
      "Wall time: 1.03 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# create an index\n",
    "\n",
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context\n",
    ")\n",
    "print (\"✅ Created index:\", index )\n",
    "print (\"✅ Saved index to db \", MY_CONFIG.DB_URI )"
   ]
  }
 ],
 "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
}