Lab 3b: Multimodal RAG

Lab 3b: Multimodal RAG#

Table of Contents#

This exercise shows how to implement a multimodal retrieval augmented generation (RAG) system. In retrieval augmented generation, an external source of information as well as the input prompt are used to generate the response. In a multimodal setting, one of the most popular use cases is to include the images in the response generation process.

This exercise implements the multimodal RAG system by using a PDF file that include images, text, and tables. This PDF file is the external source of information mentioned earlier in RAG definition. Once the system is setup, the model will be able to generate its responses considering the images, text, and tables from the PDF provided.

About This Lab#

Throughout this lab, you will encounter two types of interactive elements:

            ![Activity](../mlu_utils/images/activity.png)

            ![Challenge](../mlu_utils/images/challenge.png)

            No coding is needed for an activity. You try to understand a concept,

answer questions, or run a code cell.

            Challenges are where you test your understanding by implementing something new or taking a short quiz.

Please work through this notebook from top to bottom to avoid errors due to missing code or context.

1. Installing dependencies#

Note: the pip command below will output some error messages. You can disregard them as they are not affecting the notebook.

Installing the required libraries:

%%capture
!pip install -q -r ../requirements.txt

Importing the libraries used in this exercise:

import sys
sys.path.append('..')

import boto3
from botocore.exceptions import ClientError
import os
import json
import numpy as np
import base64
import pymupdf
import pandas as pd
from PIL import Image
import faiss
from tqdm import tqdm
from IPython import display

# Import utility functions that provide answers to challenges
%load_ext autoreload
%aimport mlu_utils.course_utils

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

2. Process PDF#

In this lab, we will read the sample PDF file of a the well-known paper “Attention Is All You Need” paper by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin. This research paper laid out the foundations of the transformers models that power many generative AI applications nowadays. Paper linked here. The paper is 11 pages long.

Challenge

Challenge: Test different documents#

    After observing how the RAG workflow is implemented in this lab, test it with other PDFs and observe how well it performds with varying document formats.

filename = "transformers_paper.pdf"

display.IFrame("data/"+ filename, width=600, height=600)

Extract text and images from each page#

The contents of the PDF need to be extracted and processed to be compatible with the RAG application.

The following are the steps we will follow to process the data in this section:

  1. Extract the data (text and images) from the PDF using pymupdf.

  2. Go through each page. Create smaller text chunks from the text of the page.

  3. Convert each page of the PDF into an image.

  4. For each text chunk, image and page, generate embeddings using Amazon Nova 2 Multimodal Embeddings.

  5. Save the information of each page in a list to store in a vector database.

from mlu_utils.multimodal_utils import pdf2imgs

doc = pymupdf.open("data/"+ filename)
num_pages = len(doc)

# Define the directories to store the extracted text, images and page images from each page
image_save_dir = "images/data"
text_save_dir = "text/data"
page_images_save_dir = "page_images"

# Chunk the text for effective retrieval
chunk_size = 700
overlap=200


items = []
# Process all pages of the PDF
for page_num in tqdm(range(num_pages), desc="Processing PDF pages"):
    page = page = doc[page_num]
    text = page.get_text()
    
    # Process chunks with overlap
    chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size-overlap)]
    
    # Generate an item to add to items
    for i,chunk in enumerate(chunks):
        text_file_name = f"{text_save_dir}/{filename}_text_{page_num}_{i}.txt"
        print("text_file_name", text_file_name)
        print("text_save_dir", text_save_dir)
        print("filename", filename)
        # If the text folder doesn't exist, create one
        os.makedirs(text_save_dir, exist_ok=True)
        with open(text_file_name, 'w') as f:
            f.write(chunk)
        
        item={}
        item["page"] = page_num
        item["type"] = "text"
        item["text"] = chunk
        item["path"] = text_file_name
        items.append(item)
    
    
    # Get all the images in the current page
    images = page.get_images()
    for idx, image in enumerate(images):        
        # Extract the image data
        xref = image[0]
        pix = pymupdf.Pixmap(doc, xref)
        pix.tobytes("png")
        # Create the image_name that includes the image path
        image_name = f"{image_save_dir}/{filename}_image_{page_num}_{idx}_{xref}.png"
        # If the image folder doesn't exist, create one
        os.makedirs(image_save_dir, exist_ok=True)
        # Save the image
        pix.save(image_name)
        
        # Produce base64 string
        with open(image_name, 'rb') as f:
            image = base64.b64encode(f.read()).decode('utf8')
        
        item={}
        item["page"] = page_num
        item["type"] = "image"
        item["path"] = image_name
        item["image"] = image
        items.append(item)

# Save pdf pages as images
page_images_save_dir = pdf2imgs("data/" + filename, page_images_save_dir)

for page_num in range(num_pages):
    page_path = os.path.join(page_images_save_dir,  f"page_{page_num:03d}.png")
    
    # Produce base64 string
    with open(image_name, 'rb') as f:
        page_image = base64.b64encode(f.read()).decode('utf8')
    
    item = {}
    item["page"] = page_num
    item["type"] = "page"
    item["path"] = page_path
    item["image"] = page_image
    items.append(item)

Processing PDF pages:   0%|          | 0/11 [00:00<?, ?it/s]

text_file_name text/data/transformers_paper.pdf_text_0_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_0_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_0_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_0_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_0_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_0_5.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_5.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_6.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_7.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_1_8.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_2_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_2_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_2_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_2_3.txt
text_save_dir text/data
filename transformers_paper.pdf

Processing PDF pages: 100%|██████████| 11/11 [00:00<00:00, 31.64it/s]

text_file_name text/data/transformers_paper.pdf_text_3_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_3_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_3_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_3_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_3_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_4_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_4_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_4_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_4_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_4_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_4_5.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_4_6.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_5_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_5_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_5_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_5_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_5_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_5_5.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_5_6.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_5_7.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_6_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_6_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_6_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_6_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_6_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_6_5.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_6_6.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_7_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_7_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_7_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_7_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_7_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_7_5.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_7_6.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_8_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_8_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_8_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_8_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_8_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_8_5.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_9_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_9_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_9_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_9_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_9_4.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_9_5.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_9_6.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_10_0.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_10_1.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_10_2.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_10_3.txt
text_save_dir text/data
filename transformers_paper.pdf
text_file_name text/data/transformers_paper.pdf_text_10_4.txt
text_save_dir text/data
filename transformers_paper.pdf

Challenge

Challenge: Types of chunking#

    In the cell above, we have used a simple character-based chunking solution. This can result in broken sentences and words, losing a lot of semantic and syntactic meaning. Try updating the chunking process using a different strategy to preserve the structure and meaning of the document. You can use modules offered by LangChain for this purpose.

3. Generate Multimodal Embeddings#

We will use the same function defined in Lab 2 to generate embeddings from text or image data.

The following function is used to generate multimodal embeddings using Amazon’s Nova 2 Multimodal Embeddings model. Embeddings can be generated with text data, image data or both.

def generate_multimodal_embeddings(prompt=None, image=None, output_embedding_length = 1024):
    """
    Invoke the Amazon Nova 2 Multimodal Embeddings model using AWS Bedrock runtime.

    Args:
        prompt (str): The text prompt to provide to the model.
        image (str): A base64-encoded image data.
        output_embedding_length (int): The dimension of the output embedding (default: 1024).
    Returns:
        list: The embedding vector.

    Raises:
        ValueError: If neither prompt nor image is provided.
    """
    if not prompt and not image:
        raise ValueError("Please provide either a text prompt, base64 image or both as input")
    
    # Initialize the Amazon Bedrock runtime client
    client = boto3.client(service_name="bedrock-runtime")
    model_id = "amazon.nova-2-multimodal-embeddings-v1:0"
    
    # Detect image format if image is provided
    image_format = "png"  # Default format
    if image:
        try:
            image_bytes = base64.b64decode(image)
            # Check PNG signature
            if image_bytes[:8] == b'\\x89PNG\\r\\n\\x1a\\n':
                image_format = "png"
            # Check JPEG signature
            elif image_bytes[:2] == b'\\xff\\xd8':
                image_format = "jpeg"
            # Check GIF signature
            elif image_bytes[:6] in (b'GIF87a', b'GIF89a'):
                image_format = "gif"
            # Check WebP signature
            elif image_bytes[:4] == b'RIFF' and image_bytes[8:12] == b'WEBP':
                image_format = "webp"
        except Exception:
            # Default to png if detection fails
            image_format = "png"
    
    # Build the request body based on input type
    body = {
        "taskType": "SINGLE_EMBEDDING",
        "singleEmbeddingParams": {
            "embeddingDimension": output_embedding_length,
            "embeddingPurpose": "GENERIC_INDEX"
        }
    }
    
    if prompt:
        body["singleEmbeddingParams"]["text"] = {
            "truncationMode": "END",
            "value": prompt
        }
    
    if image:
        body["singleEmbeddingParams"]["image"] = {
            "format": image_format,
            "source": {"bytes": image}
        }

    try:
        response = client.invoke_model(
            modelId=model_id,
            body=json.dumps(body)
        )

        # Process and return the response
        result = json.loads(response.get("body").read())
        return result["embeddings"][0]["embedding"]

    except ClientError as err:
        print(
            f"Couldn't invoke Nova embedding model. Here's why: {err.response['Error']['Code']}: {err.response['Error']['Message']}"
        )
        raise

Let’s use the generate_multimodal_embeddings function to generate embeddings of every item extracted from the PDF#

embedding_vector_dimension = 1024
for item in tqdm(items, "Generating embeddings"):
    if item['type'] == 'text':
        item['embedding'] = generate_multimodal_embeddings(prompt=item['text'], output_embedding_length=embedding_vector_dimension)
    else:
        item['embedding'] = generate_multimodal_embeddings(image=item['image'], output_embedding_length=embedding_vector_dimension)

Generating embeddings: 100%|██████████| 85/85 [00:22<00:00,  3.85it/s]

4. Create vector database#

In this section, we will create an index using FAISS, similar to Lab 2. We will create a FlatIndex which measures the L2 (or Euclidean) distance between all given points between our query vector, and the vectors loaded into the index.

all_embeddings = np.array([item['embedding'] for item in items])

Challenge

Challenge: Types of Index#

    In this lab, we will use FlatIndexL2 as the index type for the vector database. Try using a different index and observe how the speed and the quality of results changes. You can try IndexFlatIP, IndexHNSWFlat or IndexIVFFlat

# Create FAISS Index
index = faiss.IndexFlatL2(embedding_vector_dimension)
index.reset() # Clear any pre-existing index
index.add(np.array(all_embeddings, dtype=np.float32))

5. Generate a RAG Response#

In this section, we will define the function generate_rag_response to generate a response with a retrieval-augmented prompt.

First, let’s import the invoke_claude_3_multimodal function that we used in Lab 1 to generate a response to a multimodal prompt.

The following function, generate_rag_response, generates a prompt containing the user query, retrieved items and invokes the LLM to generate a RAG response.

from mlu_utils.multimodal_utils import invoke_nova_lite_multimodal

def generate_rag_response(prompt, matched_items):
    
    # Create context
    text_context = ""
    image_context = []
    
    for item in matched_items:
        if item['type'] == 'text':
            text_context += str(item["page"]) + ". " + item['text'] + "\n"
        else:
            image_context.append(item['image'])
    
    # Only 5 images are supported by Claude3 models
    if len(image_context) > 5:
        image_context = image_context[:5]
    
    final_prompt = f"""You are a helpful assistant for question answering.
    The text context is relevant information retrieved.
    The provided image(s) are relevant information retrieved.
    
    <context>
    {text_context}
    </context>
    
    Answer the following question using the relevant context and images.
    
    <question>
    {prompt}
    </question>
    
    Answer:"""
    
    return invoke_nova_lite_multimodal(final_prompt, image_context, ['image/png' for _ in image_context])

6. Test RAG Workflow#

Now that we have our functions ready, let’s test our RAG application using a few prompts.

The steps we follow to generate a RAG response are:

  1. Generate an embedding of the user query. The embedding would represent the text and the images provided in the user query.

  2. Retrieve similar items from the vector database used a nearest neighbor search

  3. Create a prompt using the user query as well as the retrieved items.

  4. Generate a response using the retrieval-augmented prompt.

query = "How is the scaled-dot-product attention calculated?"

query_embedding = generate_multimodal_embeddings(prompt=query,output_embedding_length=embedding_vector_dimension)
distances, result = index.search(np.array(query_embedding, dtype=np.float32).reshape(1,-1), k=5)

result.flatten()

array([18, 20, 17, 29, 24])

matched_items = [items[index] for index in result.flatten()]

response = generate_rag_response(query, matched_items)

display.Markdown(response)

The scaled-dot-product attention is calculated using the following steps:

  1. Inputs: The inputs to the scaled-dot-product attention are:

    • Queries (Q) of dimension (d_k)

    • Keys (K) of dimension (d_k)

    • Values (V) of dimension (d_v)

  2. Dot Product: Compute the dot products of the query with all keys. This results in a matrix of shape ((N, T_{q}, T_{k})), where (N) is the batch size, (T_{q}) is the number of queries, and (T_{k}) is the number of keys.

  3. Scaling: Divide each of these dot products by (\sqrt{d_k}). This scaling factor helps stabilize the gradients by preventing the softmax from having extremely small gradients.

  4. Softmax: Apply the softmax function to the scaled dot products to obtain the attention weights. The softmax function ensures that the weights sum to 1.

  5. Weighted Sum: Multiply the attention weights by the values (V) to obtain the output of the attention mechanism.

Mathematically, the scaled-dot-product attention can be expressed as: [ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V ]

Where:

  • (Q) is the matrix of queries.

  • (K) is the matrix of keys.

  • (V) is the matrix of values.

  • (d_k) is the dimension of the queries and keys.

This process allows the model to focus on different parts of the input sequence when producing the output sequence.

Nice. We have seen a few example questions and answers. Let’s try asking more questions. Some example questions are given below.

Activity

Activity: Asking more questions about the document#

    So far, we have seen a few example questions. Let's use the following questions and examine the responses.

        "How long were the base and big models trained?"
        "Which optimizer was used when training the models?"
        "What is position-wise feed-forward neural network mentioned in the paper?"
        "What is the BLEU score of the model in english to french translation (EN-FR)?"
        "What is the BLEU score of the model in english to german translation (EN-DE)?"

7. Quiz Questions#

Well done on completing the lab! Now, it’s time for a brief knowledge assessment.

Challenge

Challenge: Try it Yourself!#

    Answer the following questions to test your understanding of using prompt templates for inference.

from mlu_utils.quiz_questions import lab3b_question1, lab3b_question2

lab3b_question1.display()
lab3b_question2.display()

Conclusion#

In this lab, you have:

    Learned how to process PDF documents to extract text and images
    Generated multimodal embeddings using Amazon Nova 2 Multimodal Embeddings
    Created a vector database using FAISS for efficient similarity search
    Implemented a multimodal RAG system that can answer questions about the document
    Tested the RAG workflow with various queries

Additional Resources#

    Attention Is All You Need paper
    FAISS Tutorial

Thank you!#