Build a RAG chain for NVIDIA Triton documentation website

In this notebook we demonstrate how to build a RAG using NVIDIA AI Endpoints for LangChain. We create a vector store by downloading web pages and generating their embeddings using FAISS. We then showcase two different chat chains for querying the vector store. For this example, we use the NVIDIA Triton documentation website, though the code can be easily modified to use any other source.

First stage is to load NVIDIA Triton documentation from the web, chunkify the data, and generate embeddings using FAISS

To run this notebook, you need to complete the setup and generate an API key.

Let's install the prerequisite libraries and import the necessary packages to run this notebook.

Provide the API key by running the cell below.

Helper functions for loading html files, which we'll use to generate the embeddings. We'll use this later to load the relevant html documents from the Triton documentation website and convert to a vector store.

Read html files and split text in preparation for embedding generation Note chunk_size value must match the specific LLM used for embedding genetation

Make sure to pay attention to the chunk_size parameter in TextSplitter. Setting the right chunk size is critical for RAG performance, as much of a RAG’s success is based on the retrieval step finding the right context for generation. The entire prompt (retrieved chunks + user query) must fit within the LLM’s context window. Therefore, you should not specify chunk sizes too big, and balance them out with the estimated query size. For example, while OpenAI LLMs have a context window of 8k-32k tokens, Llama3 is limited to 8k tokens. Experiment with different chunk sizes, but typical values should be 100-600, depending on the LLM.

Generate embeddings using NVIDIA AI Endpoints for LangChain and save embeddings to offline vector store in the ./data/nv_embedding directory for future re-use

Second stage is to load the embeddings from the vector store and build a RAG using NVIDIAEmbeddings

Create the embeddings model using NVIDIA Retrieval QA Embedding endpoint. This model represents words, phrases, or other entities as vectors of numbers and understands the relation between words and phrases. See here for reference: https://build.nvidia.com/nvidia/embed-qa-4

Load documents from vector database using FAISS

Create a ConversationalRetrievalChain chain using NeMoLLM. In this chain we demonstrate the use of 2 LLMs: one for summarization and another for chat. This improves the overall result in more complicated scenarios. We'll use Llama3 70B for the first LLM and Mixtral for the Chat element in the chain. We add a question_generator to generate relevant query prompt. See here for reference: https://python.langchain.com/docs/modules/chains/popular/chat_vector_db#conversationalretrievalchain-with-streaming-to-stdout

Ask any question about Triton

Ask another question about Triton

Finally showcase chat capabilites by asking a question about the previous query

Now we demonstrate a simpler chain using a single LLM only, a chat LLM

Now try asking a question about Triton with the simpler chain. Compare the answer to the result with previous complex chain model

Ask another question about Triton

Finally showcase chat capabilites by asking a question about the previous query