nim-rag.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8dcac362-deed-420a-a1b0-86207dce7c04",
   "metadata": {},
   "source": [
    "Mostly stolen from:\n",
    "\n",
    "https://github.com/NVIDIA/GenerativeAIExamples/blob/main/RAG/notebooks/langchain/RAG_Langchain_with_Local_NIM.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c51fc9-1836-4044-a253-c1200ada535b",
   "metadata": {},
   "source": [
    "# Build a RAG using NVIDIA NIM microservices locally installed on MPC. \n",
    "\n",
    "Use a Llama3-8b-instruct model using NVIDIA NIM for LLMs that is locally hosted on SIH-MPC and connect to it using LangChain NVIDIA AI Endpoints package.\n",
    "\n",
    "Create a vector store by downloading web pages and generating their embeddings using FAISS. The embedding model, uses the GPU accelerated NV-Embed-QA model from NVIDIA API Catalog.\n",
    "\n",
    "This example creates a RAG from the VAST Pipeline documentation web pages.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "470f85ba-12c4-469e-b610-ce5d6e8bfe44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bunch of dependencies from LangChain, BeautifulSoup, FAISS\n",
    "import re\n",
    "\n",
    "import requests\n",
    "from bs4 import BeautifulSoup\n",
    "import getpass\n",
    "import os\n",
    "from langchain.chains import ConversationalRetrievalChain, LLMChain\n",
    "from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT, QA_PROMPT\n",
    "from langchain.chains.question_answering import load_qa_chain\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16ad09f1-b7ef-4b14-b1e1-d1ffd7ac83a9",
   "metadata": {},
   "source": [
    "Set the NVIDIA API key as `NVIDIA_API_KEY` environment variable if `NVIDIA_API_KEY` doesn't already exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6221565b-7fcf-4ad7-ab85-3ff02350be65",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_api_key():\n",
    "    if not os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "        nvapi_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "        assert nvapi_key.startswith(\"nvapi-\"), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "        os.environ[\"NVIDIA_API_KEY\"] = nvapi_key\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e679c4c-d548-41ad-86b3-2cd6ae9d7445",
   "metadata": {},
   "source": [
    "Check to see if we can connect to the local NIM on MPC (10.167.67.78)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3200146-be38-4eb9-a3e7-97bd2056d6e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_api_key()\n",
    "llm = ChatNVIDIA(base_url=\"http://10.167.67.78:8000/v1\", model=\"meta/llama3-8b-instruct\", temperature=0.1, max_tokens=64, top_p=1.0)\n",
    "result = llm.invoke(\"What is VAST?\")\n",
    "print(result.content)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1bcbff8-69fd-4f06-a527-84f5162f8917",
   "metadata": {},
   "outputs": [],
   "source": [
    "def html_document_loader(url):\n",
    "    try:\n",
    "        response = requests.get(url)\n",
    "        html_content = response.text\n",
    "    except Exception as e:\n",
    "        print(f\"Failed to load {url} due to exception {e}\")\n",
    "        return \"\"\n",
    "\n",
    "    try:\n",
    "        # Create a Beautiful Soup object to parse html\n",
    "        soup = BeautifulSoup(html_content, \"html.parser\")\n",
    "\n",
    "        # Remove script and style tags\n",
    "        for script in soup([\"script\", \"style\"]):\n",
    "            script.extract()\n",
    "\n",
    "        # Get the plain text from the HTML document\n",
    "        text = soup.get_text()\n",
    "\n",
    "        # Remove excess whitespace and newlines\n",
    "        text = re.sub(\"\\s+\", \" \", text).strip()\n",
    "\n",
    "        return text\n",
    "    except Exception as e:\n",
    "        print(f\"Exception {e} while loading document\")\n",
    "        return \"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2235ffe0-cdc7-4487-b96e-e31ccdcca451",
   "metadata": {},
   "outputs": [],
   "source": [
    "html_document_loader(\"https://vast-survey.org/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "901e6d6a-d1b1-4a9b-9e99-0c9f40c33faa",
   "metadata": {},
   "source": [
    "The `create_embeddings` function creates embeddings from a list of input URLs using the embedding model specified and saves them to disk locally.\n",
    "\n",
    "The model represents words, phrases, or other entities as vectors of numbers and understands the relation between words and phrases.\n",
    "\n",
    "kwargs to `RecursiveCharacterTextSplitter`:\n",
    "\n",
    "chunk_size (int) – Maximum size of chunks to return\n",
    "\n",
    "chunk_overlap (int) – Overlap in characters between chunks\n",
    "\n",
    "length_function (Callable[[str], int]) – Function that measures the length of given chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fcd02c9-6da5-4cdb-b28e-7836423c60b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_embeddings(embedding_path, embedding_model, in_urls):\n",
    "\n",
    "    print(f\"Storing embeddings to {embedding_path}\")\n",
    "\n",
    "    documents = []\n",
    "    for url in in_urls:\n",
    "        document = html_document_loader(url)\n",
    "        documents.append(document)\n",
    "\n",
    "\n",
    "    text_splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size=1000,\n",
    "        chunk_overlap=0,\n",
    "        length_function=len,\n",
    "    )\n",
    "    texts = text_splitter.create_documents(documents)\n",
    "    index_docs(text_splitter, texts, embedding_path, embedding_model)\n",
    "    print(\"Generated embedding successfully\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "541138c5-0787-43f9-b7f6-2380e34c1294",
   "metadata": {},
   "source": [
    "`index_docs` (called by `create_embeddings` above) does the work of splitting the text into chunks and creating the embeddings.\n",
    "\n",
    "The embedding file is created if it doesn't already exist otherwise each input text chunk is appended to the already existing output file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fc0864b-da9b-42cc-a453-9c0763eeb29e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def index_docs(splitter, documents, dest_embed_dir, embedding_model):\n",
    "    \"\"\"\n",
    "    Split the document into chunks and create embeddings for the document\n",
    "    \"\"\"\n",
    "\n",
    "    for document in documents:\n",
    "        texts = splitter.split_text(document.page_content)\n",
    "\n",
    "        # metadata to attach to document\n",
    "        metadatas = [document.metadata]\n",
    "\n",
    "        # create embeddings and add to vector store\n",
    "        if os.path.exists(dest_embed_dir):\n",
    "            update = FAISS.load_local(folder_path=dest_embed_dir, embeddings=embedding_model, allow_dangerous_deserialization=True)\n",
    "            update.add_texts(texts, metadatas=metadatas)\n",
    "            update.save_local(folder_path=dest_embed_dir)\n",
    "        else:\n",
    "            docsearch = FAISS.from_texts(texts, embedding=embedding_model, metadatas=metadatas)\n",
    "            docsearch.save_local(folder_path=dest_embed_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba697c33-cfda-4d07-bbd5-08d55ecbdca7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bd02c57-d294-49c2-b261-6637fa61f002",
   "metadata": {},
   "outputs": [],
   "source": [
    "VAST_URLS = [\"https://vast-survey.org/vast-pipeline/dev/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/gettingstarted/installation/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/gettingstarted/configuration/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/gettingstarted/deployment/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/overview/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/imageingest/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/association/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/newsources/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/monitor/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/sourcestats/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/architecture/intro/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/faq/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/help_and_acknowledgements/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/overview/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/imageingest/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/association/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/newsources/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/monitor/\",\n",
    "             \"https://vast-survey.org/vast-pipeline/1.2.0/design/sourcestats/\",\n",
    "             \"https://vast-survey.org/\",\n",
    "             \"https://vast-survey.org/Survey/\",\n",
    "             \"https://vast-survey.org/Team/\"\n",
    "            ]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "760f16c5-7d5c-4925-ae98-145a51f42e41",
   "metadata": {},
   "source": [
    "Create the Embeddings from the list if URLs and save them locally.\n",
    "\n",
    "NVIDIAEmbeddings is from the NVIDIA AI Endpoints for LangChain library.\n",
    "\n",
    "The \"NV-Embed-QA\" model is described here: https://build.nvidia.com/nvidia/embed-qa-4/modelcard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5cdad2d-1e73-4d38-b93c-200da04d7842",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_api_key()\n",
    "embedding_model = NVIDIAEmbeddings(model=\"NV-Embed-QA\", truncate=\"END\")\n",
    "embedding_path = \"./data/vast_embedding_all\"\n",
    "create_embeddings(embedding_path, embedding_model, VAST_URLS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fe7b14c-0acb-4675-a636-2b7fb6b253cf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "23222890-19ce-42e3-8a4d-1f2603e7a992",
   "metadata": {},
   "source": [
    "Use `ConversationRetrievalChain` to create a Question-Answer chat-bot from the locally hosted LLM, providing the embedded URL text as context. Also provide chat memory using `ConversationBufferMemory`.\n",
    "\n",
    "https://python.langchain.com/api_reference/langchain/chains/langchain.chains.conversational_retrieval.base.ConversationalRetrievalChain.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b59fdfe-e667-44ca-b9ad-98ae42e9e4fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_api_key()\n",
    "llm = ChatNVIDIA(base_url=\"http://10.167.67.78:8000/v1\", model=\"meta/llama3-8b-instruct\", temperature=0.1, max_tokens=1000, top_p=1.0)\n",
    "\n",
    "embedding_model = NVIDIAEmbeddings(model=\"NV-Embed-QA\", truncate=\"END\")\n",
    "embedding_path = \"./data/vast_embedding_all_bak\"\n",
    "docsearch = FAISS.load_local(folder_path=embedding_path, embeddings=embedding_model, allow_dangerous_deserialization=True)\n",
    "\n",
    "qa_prompt=QA_PROMPT\n",
    "\n",
    "memory = ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True)\n",
    "\n",
    "qa = ConversationalRetrievalChain.from_llm(\n",
    "    llm=llm,\n",
    "    retriever=docsearch.as_retriever(),\n",
    "    chain_type=\"stuff\",\n",
    "    memory=memory,\n",
    "    combine_docs_chain_kwargs={'prompt': qa_prompt},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16f24df2-d704-4664-87c9-826525e45380",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_query(query):\n",
    "    result = qa({\"question\": query})\n",
    "    print(result.get(\"answer\"))\n",
    "\n",
    "query = \"What is VAST?\"\n",
    "get_query(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0555428b-81be-4705-aab3-d96aa1a2c7ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_query(\"What is the VAST pipeline?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0587e4e-de7a-4bba-8dc1-29348cdd99e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_query(\"What are the steps of it?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5341051-9ebf-4e64-9e79-f9ff1c1d85f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_query(\"Describe the final step in more detail.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb5ca004-157f-448b-996f-ae7c36a86ef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try another embedding model\n",
    "from langchain_huggingface import HuggingFaceEmbeddings\n",
    "\n",
    "embedding_model = HuggingFaceEmbeddings(model_name=\"snowflake/arctic-embed-xl\")\n",
    "embedding_path = \"./data/vast_embedding_snowflake\"\n",
    "create_embeddings(embedding_path, embedding_model, VAST_URLS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75d25919-4db7-4ce6-a51f-2846a2b2807f",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatNVIDIA(base_url=\"http://10.167.67.78:8000/v1\", model=\"meta/llama3-8b-instruct\", temperature=0.1, max_tokens=1000, top_p=1.0)\n",
    "\n",
    "embedding_path = \"./data/vast_embedding_snowflake_bak\"\n",
    "docsearch = FAISS.load_local(folder_path=embedding_path, embeddings=embedding_model, allow_dangerous_deserialization=True)\n",
    "\n",
    "qa_prompt=QA_PROMPT\n",
    "\n",
    "memory = ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True)\n",
    "\n",
    "qa = ConversationalRetrievalChain.from_llm(\n",
    "    llm=llm,\n",
    "    retriever=docsearch.as_retriever(),\n",
    "    chain_type=\"stuff\",\n",
    "    memory=memory,\n",
    "    combine_docs_chain_kwargs={'prompt': qa_prompt},\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fd0c235-f8c1-48dd-b044-64f867d691ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_query(query):\n",
    "    result = qa({\"question\": query})\n",
    "    print(result.get(\"answer\"))\n",
    "\n",
    "query = \"What is VAST?\"\n",
    "get_query(query)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c3fc7a5-2a00-4768-adb3-ea12004a349e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}