notebooks/02_Option(1)_NVIDIA_AI_endpoint_simple.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "01537915",
   "metadata": {},
   "source": [
    "# NVIDIA API Catalog with LangChain\n",
    "\n",
    "In this notebook, we are going to use the [**ai-mixtral-8x7b-instruct as LLM**](https://build.nvidia.com/mistralai/mixtral-8x7b-instruct) and the [**ai-embed-qa-4 embedding**](https://build.nvidia.com/nvidia/embed-qa-4) provided by [NVIDIA_AI_Catelog](https://build.nvidia.com/explore/discover) and build a simply RAG example with faiss as vectorstore\n",
    "\n",
    "### Prerequisite \n",
    "In order to successfully run this notebook, you will need the following -\n",
    "\n",
    "1. Already successfully gone through the [setup](https://python.langchain.com/docs/integrations/text_embedding/nvidia_ai_endpoints#setup) and generated an API key.\n",
    "2. install necesary python dependencies in [requirements.txt](https://github.com/NVIDIA/GenerativeAIExamples/blob/3d29acf677466c5c301370cab5867cb09e04e318/notebooks/requirements.txt) : then upgrade the langchain-core with the below  \n",
    "\n",
    "Note: change **faiss-gpu --> faiss-cpu** in pre-requisite 2\n",
    "if you do not have access to a GPU.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09268061",
   "metadata": {},
   "source": [
    "### Step 1  - Export the NVIDIA_API_KEY\n",
    "You can supply the NVIDIA_API_KEY directly in this notebook when you run the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa57a6ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install faiss-cpu # replace with faiss-gpu if you are using GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd44077c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "## API Key can be found by going to NVIDIA NGC -> AI Foundation Models -> (some model) -> Get API Code or similar.\n",
    "## 10K free queries to any endpoint (which is a lot actually).\n",
    "\n",
    "# del os.environ['NVIDIA_API_KEY']  ## delete key and reset\n",
    "if os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    print(\"Valid NVIDIA_API_KEY already in environment. Delete to reset\")\n",
    "else:\n",
    "    nvapi_key = getpass.getpass(\"NVAPI Key (starts with nvapi-): \")\n",
    "    assert nvapi_key.startswith(\"nvapi-\"), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aec407b1",
   "metadata": {},
   "source": [
    "### Step 2 - initialize the LLM \n",
    "Here we will use **ai-mixtral-8x7b-instruct** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87e26a58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test run and see that you can genreate a respond successfully\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "\n",
    "llm = ChatNVIDIA(model=\"ai-mixtral-8x7b-instruct\", nvidia_api_key=nvapi_key, max_tokens=1024)\n",
    "\n",
    "result = llm.invoke(\"Write a ballad about LangChain.\")\n",
    "print(result.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5f20907",
   "metadata": {},
   "source": [
    "### Step 3 - We intiatlize the embedding as well \n",
    "We selected **ai-embed-qa-4** as the embedding \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0608482f",
   "metadata": {},
   "source": [
    "## first we initialize the embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c84ae59",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings\n",
    "\n",
    "embedder = NVIDIAEmbeddings(model=\"ai-embed-qa-4\")\n",
    "\n",
    "# Alternatively, if you want to specify whether it will use the query or passage type\n",
    "# embedder = NVIDIAEmbeddings(model=\"ai-embed-qa-4\", model_type=\"passage\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a6bb1b",
   "metadata": {},
   "source": [
    "### Step 4 - Obtain some toy text dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a19ad289",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "\n",
    "# Here we read in the text data and prepare them into vectorstore\n",
    "ps = os.listdir(\"./toy_data/\")\n",
    "data = []\n",
    "sources = []\n",
    "for p in ps:\n",
    "    if p.endswith('.txt'):\n",
    "        path2file=\"./toy_data/\"+p\n",
    "        with open(path2file,encoding=\"utf-8\") as f:\n",
    "            lines=f.readlines()\n",
    "            for line in lines:\n",
    "                if len(line)>=1:\n",
    "                    data.append(line)\n",
    "                    sources.append(path2file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b87f583b",
   "metadata": {},
   "source": [
    "### Step 5 - Do some basic cleaning and remove empty lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a967ae86",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents=[d for d in data if d != '\\n']\n",
    "len(data), len(documents), data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a061bf4",
   "metadata": {},
   "source": [
    "### Step 6a (optional) - Speed test: check how fast ( in seconds) processing 1 document vs. a batch of 10 documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ab3429a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "print(\"Single Document Embedding: \")\n",
    "s = time.perf_counter()\n",
    "q_embedding  = embedder.embed_documents([documents[0]])\n",
    "elapsed = time.perf_counter() - s\n",
    "print(\"\\033[1m\" + f\"Executed in {elapsed:0.2f} seconds.\" + \"\\033[0m\")\n",
    "print(\"Shape:\", (len(q_embedding),))\n",
    "\n",
    "print(\"\\nBatch Document Embedding: \")\n",
    "s = time.perf_counter()\n",
    "d_embeddings = embedder.embed_documents(documents[:10])\n",
    "elapsed = time.perf_counter() - s\n",
    "print(\"\\033[1m\" + f\"Executed in {elapsed:0.2f} seconds.\" + \"\\033[0m\")\n",
    "print(\"Shape:\",len(d_embeddings[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99905aa0",
   "metadata": {},
   "source": [
    "### Step 6b - Process the documents into faiss vectorstore and save it to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fb397bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we create a vector store from the documents and save it to disk.\n",
    "from operator import itemgetter\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "import faiss\n",
    "# create my own uuid\n",
    "text_splitter = CharacterTextSplitter(chunk_size=400, separator=\" \")\n",
    "docs = []\n",
    "metadatas = []\n",
    "\n",
    "for i, d in enumerate(documents):\n",
    "    splits = text_splitter.split_text(d)\n",
    "    #print(len(splits))\n",
    "    docs.extend(splits)\n",
    "    metadatas.extend([{\"source\": sources[i]}] * len(splits))\n",
    "\n",
    "store = FAISS.from_texts(docs, embedder , metadatas=metadatas)\n",
    "store.save_local('./toy_data/nv_embedding')\n",
    "\n",
    "# you will only need to do this once, later on we will restore the already saved vectorstore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1421512a",
   "metadata": {},
   "source": [
    "### Step 6c - Read the previously processed & saved Faiss vectore store back"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a35b5004",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the vectorestore back.\n",
    "\n",
    "store = FAISS.load_local(\"./toy_data/nv_embedding\", embedder)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "572ca074",
   "metadata": {},
   "source": [
    "### Step 7- Wrap the restored vectorsore into a retriever and ask our question "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba866e74",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "retriever = store.as_retriever()\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"Answer solely based on the following context:\\n<Documents>\\n{context}\\n</Documents>\",\n",
    "        ),\n",
    "        (\"user\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chain = (\n",
    "    {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "chain.invoke(\"Tell me about Sweden.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "048e4200",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}