Retrieval-Augmented Generation: Question Answering using Llama-2 and Text Embedding Models

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

In this notebook we will demonstrate how to use Llama-2-7b to answer questions using a library of documents as a reference, by using document embeddings and retrieval. Unlike other RAG solutions, embeddings will be generated and combined with the embedding model to identify the nearest neighbors, all from a single endpoint in this solution.

To perform inference on the Llama models, you need to pass custom_attributes='accept_eula=true' as part of header. This means you have read and accept the end-user-license-agreement (EULA) of the model. EULA can be found in model card description or from this webpage.

Note: Custom_attributes used to pass EULA are key/value pairs. The key and value are separated by '=' and pairs are separated by ';'. If the user passes the same key more than once, the last value is kept and passed to the script handler (i.e., in this case, used for conditional logic). For example, if 'accept_eula=false; accept_eula=true' is passed to the server, then 'accept_eula=true' is kept and passed to the script handler.

Other Retrieval Augmented Generation Solutions -

• Question Answering using LangChain and Cohere's Generate and Embedding Models from SageMaker JumpStart
• Question Answering based on Custom Dataset
• Question Answering based on Custom Dataset with Open-sourced LangChain Library
• Question Answering using LLama-2, Pinecone & Custom Dataset

Step 1. Deploy Llama-2 7 Billion Chat Model in SageMaker JumpStart

To begin, we will initialize all of the SageMaker session variables we'll need to use throughout the walkthrough.

We will use a ml.g5.4xlarge instance to deploy our Llama-2-7 billion model. We can find pricing for all instances here.

To gain an understanding of the necessity for a retrieval-augmented generation (RAG) approach in addressing the question and answering problem, please refer to this question_answering_pinecone_llama-2_jumpstart.ipynb

Step 2. Use Text Embedding to identify the correct documents, and use them along with prompt and question to query LLM

We plan to use document embeddings to fetch the most relevant documents in our document knowledge library and combine them with the prompt that we provide to LLM.

To achieve that, we will do following.

• Running a text embedding model training job. The training job will generate embeddings for dataset provided and save them along with the model. These embeddings will be utilized during inference to find the nearest neighbors for an input sentence. The nearest neighbor is based on the cosine similarity between the input sentence embedding and already computed sentence embeddings during the training job. To get more information please refer to text-embedding-sentence-similarity.ipynb
• Query the text embedding model endpoint created above to Identify top K most relevant documents based on user query
• Combine the retrieved documents with prompt and question and send them into LLM.

Note: We are saving the dataset here with the model only to get the most similar document unlike the other RAG solutions.

Note: The retrieved document/text should be large enough to contain enough information to answer a question; but small enough to fit into the LLM prompt -- maximum sequence length of 1024 tokens.

To train and host on Amazon Sagemaker, we need to setup and authenticate the use of AWS services. Here, we use the execution role associated with the current notebook instance as the AWS account role with SageMaker access. It has necessary permissions, including access to your data in S3.

2.1. Preparing Dataset

2.2. Getting the Embeddings for the Input data and Training Job

2.3. Deploy & run Inference on the model to get nearest neighbor

You can make queries to the endpoint using a JSON payload containing a batch of input texts, to find the nearest neighbors of the input text from the dataset which is provided during the training job.

• queries: Provide the list of inputs for which to find the closest match from the training data
• top_k: The number of closest match to find from the training data
• mode: Supply it as "nn_train_data" for getting the nearest neighbors to input queries within the dataset provided

2.4 Combine the retrieved documents, prompt, and question to query the LLM

Now we're ready begin querying our LLM with a Retrieval Augmented Generation (RAG) pipeline. Let's see how this will work step-by-step first.

Let's place all of this logic into a single RAG query function:

We can now ask the question:

We can also ask questions about things that are out of context (not contained within our dataset). From this we expect the model to not hallucinate and honestly tell us that it does not know the answer:

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.