Preparing data for a RAG pipeline using Data Prep Kit (DPK)

Retrieval-augmented generation (RAG) enhances large language models (LLMs) by providing them with relevant, domain-specific information that they wouldn't otherwise know. LLMs are trained on vast amounts of public data but lack access to private or proprietary information. RAG solves this problem by retrieving relevant contextual information and feeding it to the model, thereby improving accuracy and reducing hallucinations.

RAG is widely used for question-answering systems, customer support automation, and knowledge-based AI applications. For an in-depth understanding, check out What is Retrieval-Augmented Generation?.

This tutorial will walk you through preparing PDF documents for a RAG pipeline and running queries on them effectively.

The following diagram illustrates the RAG pipeline workflow that you will implement in this tutorial.

Prerequisites

To follow along, you’ll need:

• A local Python development environment. While you can use an Anaconda Python envronment, other tutorials in the Data Prep Kit learning path have used a Python virtual environment on Python v3.11.
• The code for this tutorial. Clone the Data Prep Kit repo locally, and you'll find the Jupyter Notebook in the examples/rag-pdf-1 directory.
• A (free) account at Replicate to run LLMs. Use this invite to add some credit to your Replicate account! The free account will give you a few API calls for free. That is enough for this tutorial. Once you sign up, create a token by going to your account, selecting API tokens, and creating a new token, called rag-1 (you can use any name).

Steps

The tutorial is structured into these main steps (click on the links to jump to the section)

1. Processing PDFs – Extract, clean, and chunk documents.
2. Saving Data in a Vector Database – Store processed data efficiently.
3. Performing Vector Searches – Retrieve relevant document chunks.
4. Querying Documents Using LLMs – Interact with your documents via an LLM.

Step 1: Processing PDFs using Data Prep Kit

In this step, we will process PDF documents to prepare them for Retrieval Augmented Generation (RAG) queries. You can find the code for this step in the rag_1_dpk_process_python.ipynb Notebook.

The steps involved are as follows:

1. Reading and extracting contents from PDFs
2. Removing duplicates
3. Splitting the extracted content into chunks
4. Creating embeddings for the chunks

Step 1.1: Setting up configuration parameters

First, we define all our configuration parameters in one class, which keeps the global namespace clean. We have common configuration parameters defined in my_config.py. The following code is a snippet of these configuration parameters:

class MyConfig:
    pass 

MY_CONFIG = MyConfig ()

## Input Data - configure this to the folder we want to process
MY_CONFIG.INPUT_DATA_DIR = "input"
MY_CONFIG.OUTPUT_FOLDER = "output"
MY_CONFIG.OUTPUT_FOLDER_FINAL = os.path.join(MY_CONFIG.OUTPUT_FOLDER , "output_final")
### -------------------------------

### Milvus config
MY_CONFIG.DB_URI = './rag_1_dpk.db'  # For embedded instance
MY_CONFIG.COLLECTION_NAME = 'dpk_papers'


## Embedding model
MY_CONFIG.EMBEDDING_MODEL = 'ibm-granite/granite-embedding-30m-english'
MY_CONFIG.EMBEDDING_LENGTH = 384

## LLM Model
MY_CONFIG.LLM_MODEL = "ibm-granite/granite-3.0-8b-instruct"
MY_CONFIG.MAX_CONTEXT_WINDOW = 4096 #tokens

This configuration class centralizes all parameters, which makes it easy to add new ones:

MY_CONFIG.SAMPLE_CONFIG = "hello"

In the main code, we initialize the configuration class:

from my_config import MY_CONFIG

Step 1.2: Fetching the PDFs

Next, we download three PDF documents from arxiv.org, using the download_file function from utils.py. The function checks for local existence before downloading.

from utils import download_file

download_file (url = 'https://arxiv.org/pdf/1706.03762', \
local_file = os.path.join(MY_CONFIG.INPUT_DATA_DIR, 'attention.pdf' ))
download_file (url = 'https://arxiv.org/pdf/2405.04324', \
local_file = os.path.join(MY_CONFIG.INPUT_DATA_DIR, 'granite.pdf' ))
download_file (url = 'https://arxiv.org/pdf/2405.04324', \
local_file = os.path.join(MY_CONFIG.INPUT_DATA_DIR, 'granite2.pdf' )) # duplicate

Why the duplicate? We intentionally download the same file twice to demonstrate duplicate removal in a later step.

Step 1.3: Setting up input and output directories

The processing of the PDF files follows a multi-stage approach:

input ---(stage-1)--> output/01_stage1_out ---(stage 2)---> output/02_stage2_out

The output (output/01_stage1_out) file becomes input to the next stage. And so on.

The final PDF processing pipeline looks as follows:

pdfs → (extract text) → 01_parquet_out → (dedupe) → 02_dedupe_out → (chunking) → 03_chunk_out → (embeddings) → 04_embeddings_out → output_final

The output folder structure looks as follows:

output
├── 01_parquet_out
├── 02_dedupe_out
├── 03_chunk_out
├── 04_embeddings_out
└── output_final

Step 1.4: Extracting the contents of the PDF files

Next, we will extract the contents from the PDF files using the Pdf2Parquet transformer:

• Class: dpk_pdf2parquet.transform_python.Pdf2Parquet
• Found in package: data-prep-toolkit-transforms[pdf2parquet] (which is installed as part of data-prep-toolkit-transforms[all])

For more details about this transform, read the PDF2Parquet documentation.

Here is the code that does extracts the contents of the PDF files (abbreviated):

from dpk_pdf2parquet.transform_python import Pdf2Parquet
from dpk_pdf2parquet.transform import pdf2parquet_contents_types

result = Pdf2Parquet(input_folder= 'input',
                    output_folder= 'output/01_parquet_out',
                    data_files_to_use=['.pdf'],
                    pdf2parquet_contents_type=pdf2parquet_contents_types.MARKDOWN,   # markdown
                    ).transform()

if result != 0:
    raise Exception ('process failed')

These are the parameters that are used:

• input_folder : where to read data files from
• output_folder : destination for output folder
• data_files_to_use : specify file filters; here we are only intereseted in PDF files

• pdf2parquet_contents_type : This controls the output format for extracted content. It can be JSON or markdown.

  • For markdown: pdf2parquet_contents_type=pdf2parquet_contents_types.MARKDOWN
  • For JSON: pdf2parquet_contents_type=pdf2parquet_contents_types.JSON'

Each input PDF generates a corresponding Parquet file in output/01_parquet_out.

├── 01_parquet_out
│   ├── attention.parquet
│   ├── granite2.parquet
│   ├── granite.parquet
│   └── metadata.json

Let's inspect the output using the read_parquet_files_as_df utility function from utils.py.

from utils import read_parquet_files_as_df

output_df = read_parquet_files_as_df('01_parquet_out')
output_df.head(5)

Output explained:

• One entry per input file (so 3 total)
• filename column is the name of the input file
• contents column has the Markdown text extracted from the PDFs
• num_pages and num_tables columns indicate detected metadata
• document_id is a unique ID generated for each document
• document_hash is a hash calcuated on the entire file
• hash is calculated from contents column, which is different from document_hash
• pdf_convert_time denotes the time in seconds that it took to process each PDF

Step 1.5: Eliminating duplicate documents

Removing duplicates reduces the amount of data we have to process down the pipeline and improves the quality of RAG results. We have 2 identical copies of granite.pdf, which have identical hash values. We are going to eliminate one of the duplicates.

We will use Exect Deduplicate (Ededupe) transform for this. Ededupe computes the hash values on the contents of each file. Files with duplicate hashes are filtered out.

• Class: dpk_ededup.transform_python.Ededup
• Found in package: data-prep-toolkit-transforms[ededup] (which is installed as part of data-prep-toolkit-transforms[all])

For more details about this transform, read the Exact dedeupe filter documentation.

from dpk_ededup.transform_python import Ededup

result = Ededup(input_folder='01_parquet_out',
                output_folder='output/02_dedupe_out',
                ededup_doc_column="contents",
                ededup_doc_id_column="document_id"
                ).transform()

if result != 0:
    raise Exception ('process failed')

Parameters explained:

• input_folder: This is the output created by the previous transform (pdf2parquet)
• output_folder: The filtered out files will be in output/02_dedupe_out
• ededup_doc_column = 'contents': This is the column we will calculate hash values for
• ededup_doc_id_column = 'document_id': The column where we store the computed hash

Again, let's inspect the output using the read_parquet_files_as_df utility function from utils.py.

from utils import read_parquet_files_as_df

output_df = read_parquet_files_as_df('output/02_dedupe_out')
output_df.head(5)

After deduplication, one of the duplicate granite.pdf files is removed.

Step 1.6: Splitting documents into chunks

In this tutorial we use a chunking size of 128 tokens (~100 words) with an overlap of 30 tokens (~25 words). You should experiment with chunking settings to find one that works best for your documents. Learn more about chunking strategies, in this article, "Enhancing RAG performance with smart chunking strategies." Really dig into the details of effective chunking strategies in the IBM RAG Cookbook.

We split documents into smaller retrieval-friendly chunks using the DocChunk transform.

• Class : dpk_doc_chunk.transform_python.DocChunk
• Found in package : data-prep-toolkit-transforms[doc_chunk] (installed as part of data-prep-toolkit-transforms[all])

For more details about this transform, read the Chunking transform documentation.

from dpk_doc_chunk.transform_python import DocChunk

result = DocChunk(input_folder='output/02_dedupe_out',
                    output_folder='output/03_chunk_out',
                    doc_chunk_chunking_type= "li_markdown",
                    # doc_chunk_chunking_type= "dl_json",
                    doc_chunk_chunk_size_tokens = 128,  # default 128
                    doc_chunk_chunk_overlap_tokens=30   # default 30
                    ).transform()

Parameters explained:

• input_folder and output_folder are self explanatory.
• doc_chunk_chunking_type = "li_markdown" Used to define the markdown parser. In this tutorial, we will be using Llama-Index markdown parser. If we want to split JSON text into chunks, set this to "dl_json".
• doc_chunk_chunk_size_tokens = 128 : Size of each chunk is about 128 tokens (approx 100 words). Default value for this is 128.
• doc_chunk_chunk_overlap_tokens=30 : The overlap between chunks is about 30 tokens (~25 words). Default value for this is 30.

Again, let's inspect the output using the read_parquet_files_as_df utility function from utils.py.

from utils import read_parquet_files_as_df

input_df = read_parquet_files_as_df('output/02_dedupe_out')
output_df = read_parquet_files_as_df('output/03_chunk_out')

print (f"Files processed : {input_df.shape[0]:,}")
print (f"Chunks created : {output_df.shape[0]:,}")

output_df.sample(min(3, output_df.shape[0]))

Output:

Files processed : 2
Chunks created : 60

The 2 PDF files have been split into 60 chunks. Here is a sample output of few chunks:

Step 1.7: Create embeddings for the chunks

The final step in our processing pipeline is to generate embeddings for the extracted chunks.

Why embeddings? Embeddings play a crucial role in the RAG process, enabling fast search and retrieval of complex data like text and images. You can learn more about embeddings in the IBM RAG Cookbook.

We generate embeddings for each chunk using the TextEncoder transform.

• Class : dpk_text_encoder.transform_python.TextEncoder
• Found in package : data-prep-toolkit-transforms[text_encoder] (installed as part of data-prep-toolkit-transforms[all])

For more details about this transform, read the Text Encoder documentation

Here is the code (abbreviated) to transform chunks into embeddings:

from dpk_text_encoder.transform_python import TextEncoder

result = TextEncoder(input_folder= 'output/03_chunk_out', 
                    output_folder= 'output/04_embeddings_out', 
                    text_encoder_model_name = 'ibm-granite/granite-embedding-30m-english'
                    ).transform()

Parameters:

• input_folder and output_folder are used respectively for reading and writing.
• text_encoder_model_name = 'ibm-granite/granite-embedding-30m-english' - Here we can specify any open source embedding model. There are numerous embedding models we can use. We are using ibm-granite/granite-embedding-30m-english as it is a small model (hence quick to run) and gives decent results. Hugging Face's embedding model leaderboard is an excellent resource for finding suitable embedding models.

Again, let's inspect the output using the read_parquet_files_as_df utility function from utils.py.

from utils import read_parquet_files_as_df

output_df = read_parquet_files_as_df('output/04_embeddings_out')
output_df.sample(3)

In the produced output, we see a new column called embeddings. This column has embeddings computed for text contents column, these chunked texts:

Step 2: Storing data in a vector database

Now that we've processed our PDF data, the next step is to store it in a vector database for efficient retrieval. In this tutorial, we'll be using Milvus. Milvus is an open-source vector database that is designed for handling unstructured data like text embeddings. Milvus is built for high-performance vector similarity search and is optimized for working with embeddings generated from text, images, and audio. It allows for fast and scalable querying, making it ideal for retrieval-augmented generation (RAG) workflows like ours.

You can find the code for this step in the rag_2_load_data_into_milvus.ipynb Notebook.

The steps involved are as follows:

1. Loading the processed PDF data
2. Inserting the data into Milvus

Step 2.1: Loading the processed PDF data

Before we begin, let’s import our configuration settings:

from my_config import MY_CONFIG

Now, we'll load the data we saved in the previous step (pdf2parquet) from the output/output_final directory. This folder contains Parquet files, one Parquet file for each input PDF.

The following code does the following:

• Reads all Parquet files using pandas.read_parquet
• Combines them into a single DataFrame using pandas.concat

import pandas as pd
import glob

# Get a list of all Parquet files in the directory
parquet_files = glob.glob(f'{MY_CONFIG.OUTPUT_FOLDER_FINAL}/*.parquet')

# Create an empty list to store the DataFrames
dfs = []

# Loop through each Parquet file and read it into a DataFrame
for file in parquet_files:
    df = pd.read_parquet(file)
    dfs.append(df)

# Concatenate all DataFrames into a single DataFrame
data_df = pd.concat(dfs, ignore_index=True)

We are also renaming the columns as follows, to match the expected format for Milvus.

• embeddings --> vector
• contents --> text

data_df = data_df.rename( columns= {'embeddings' : 'vector', 'contents' : 'text'})

After this step, our processed data looks like this:

Step 2.2: Creating a Milvus collection

First, we need to connect to Milvus. We’re connecting to an embedded Milvus instance (a local database file rag_1_dpk.db). While this works for testing, in production, you’d connect to a cloud-hosted instance.

from pymilvus import MilvusClient

milvus_client = MilvusClient(MY_CONFIG.DB_URI)

Before inserting data, we clear any existing collection to ensure a fresh import. Learn more about Milvus collections in the Milvus Docs.

Then, we create a collection with these properties:

• collection_name = 'dpk_papers' - The name of our collection
• dimension=384 - Matches the output length of our embedding model
• metric_type='IP' - Uses Inner Product Distance for similarity search
• consistency_level='Strong' - Ensures consistent query results. For more details, check the Milvus documentation
• auto_id=True - Milvus assigns primary IDs automatically

# if we already have a collection, clear it first
if milvus_client.has_collection(collection_name='dpk_papers'):
    milvus_client.drop_collection(collection_name='dpk_papers')


milvus_client.create_collection(
    collection_name='dpk_papers',
    dimension=384,
    metric_type="IP",  # Inner product distance
    consistency_level="Strong",  # Strong consistency level
    auto_id=True
)

Step 2.3: Inserting data into our Milvus vector database

Now that our collection is set up, let’s insert our processed data into Milvus.

res = milvus_client.insert(collection_name='dpk_papers', 
                            data=data_df.to_dict('records'))

print('inserted # rows', res['insert_count'])
milvus_client.get_collection_stats('dpk_papers')

Example Output:

inserted # rows 60

And, finally, after the data is successfully stored, we close the Milvus connection to free up resources.

milvus_client.close()

Step 3: Performing a vector search

Now that we’ve successfully imported our data and embeddings into the Milvus vector database, it’s time to run some powerful queries!

Unlike traditional keyword searches, these queries perform semantic searches, which means they retrieve results based on meaning rather than exact word matches.

You can find the code for this step in the rag_3_vector_search.ipynb Notebook.

The steps involved are as follows:

1. Set up embeddings
2. Run a search query

Step 3.1: Set up embeddings

Our first step is to load our configuration settings:

from my_config import MY_CONFIG

Then, we need to connect to our vector database. We’ll be connecting to an embedded Milvus instance, which is a local database file (rag_1_dpk.db). While this setup is great for testing, a production environment would typically connect to a cloud-hosted Milvus instance.

from pymilvus import MilvusClient

milvus_client = MilvusClient(MY_CONFIG.DB_URI)

Now, we can convert our query strings into vector embeddings using the same embedding model we used earlier.

Here is the code:

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

from sentence_transformers import SentenceTransformer

embedding_model = SentenceTransformer(MY_CONFIG.EMBEDDING_MODEL)

def get_embeddings (str):
    embeddings = embedding_model.encode(str, normalize_embeddings=True)
    return embeddings

In this code, we first set the HF_ENDPOINT to Hugging Face’s mirror to ensure smooth model downloads. Then, the get_embeddings function converts a text query into an embedding vector.

Now, let's test the embedding conversion:

# Test Embeddings
text = 'Paris 2024 Olympics'
embeddings = get_embeddings(text)
print ('sentence transformer : embeddings len =', len(embeddings))
print ('sentence transformer : embeddings[:5] = ', embeddings[:5])

Sample Output:

sentence transformer : embeddings len = 384
sentence transformer : embeddings[:5] =  [ 0.02468892  0.10352131  0.0275264  -0.08551715 -0.01412829]

To streamline our search, we will define a helper function:

def  do_vector_search (query):
    query_vectors = [get_embeddings(query)] 

    results = milvus_client.search(
        collection_name=MY_CONFIG.COLLECTION_NAME,  # target collection
        data=query_vectors,  # query vectors
        limit=5,  # number of returned entities
        output_fields=["filename", "page_number", "text"],  
        # specifies fields to be returned
    )
    return results
## ----

This code completes these steps:

• Converts the query text into an embedding vector.
• Executes a vector search in the Milvus database.
• Returns the top 5 most relevant results.

Step 3.2: Running a search query

Now, let’s test it out!

query = "What was the training data used to train Granite models?"

results = do_vector_search (query)
print_search_results(results)

This will output the top matches, ranked by search score (higher scores indicate better matches, with a range of 0 to 1.0).

Sample Output:

num results :  5
------ result 1 --------
search score: 0.8260288238525391
filename: granite.pdf
text: ...

------ result 2 --------
search score: 0.7880659699440002
filename: granite.pdf
text: ...
...

From this output, we can see that the search has successfully retrieved relevant results, with the top match scoring 0.826.

Step 4: Querying documents using LLMs

Welcome to the final step of our RAG (Retrieval-Augmented Generation) pipeline! Here, we’ll explore how to query documents effectively using a large language model (LLM).

You can find the code for this step in the rag_4_query_replicate.ipynb Notebook.

The steps involved are as follows:

1. Set up embeddings
2. Query the LLM

Step 4.1: Load configuration

We need to do some up front configuration.

First, let’s load our configuration settings:

from my_config import MY_CONFIG

Then, we will load some settings from our .env file. Why use .env files? These .env files store sensitive configurations such as database passwords and API keys. These files should remain private and should not be checked into the codebase.

Create a file named .env (the file name must start with a dot!) in the same examples folder.

The .env file will look like this

REPLICATE_API_TOKEN=xyz

Replace xyz with the Replicate token you created from prerequisites section.

Here is the code:

from dotenv import find_dotenv, dotenv_values

config = dotenv_values(find_dotenv())

MY_CONFIG.REPLICATE_API_TOKEN = config.get('REPLICATE_API_TOKEN')

The final configuration step is connecting to our vector database, Milvus. We'll establish a connection to an embedded Milvus instance, which stores our document embeddings. For production environments, a cloud-hosted Milvus instance is recommended.

from pymilvus import MilvusClient

milvus_client = MilvusClient(MY_CONFIG.DB_URI)

Now we are ready to set up the embeddings. An embedding model converts text into vector representations for efficient searching.

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

from sentence_transformers import SentenceTransformer

embedding_model = SentenceTransformer(MY_CONFIG.EMBEDDING_MODEL)

def get_embeddings (str):
    embeddings = embedding_model.encode(str, normalize_embeddings=True)
    return embeddings

In this code, we first set the HF_ENDPOINT to Hugging Face’s mirror to ensure smooth model downloads. Then, the get_embeddings function converts a text query into an embedding vector.

Step 4.2: Querying the LLM

Before querying the LLM, we need to retrieve relevant documents using vector search.

def fetch_relevant_documents (query : str) :
    search_res = milvus_client.search(
        collection_name=MY_CONFIG.COLLECTION_NAME,
        data = [get_embeddings(query)], 
        limit=3,  # Return top 3 results
        search_params={"metric_type": "IP", "params": {}},  # Inner product distance
        output_fields=["text"],  # Return the text field
    )
    retrieved_docs_with_distances = [
        {'text': res["entity"]["text"], 'distance' : res["distance"]} for res in search_res[0]
    ]
    return retrieved_docs_with_distances

In this code, milvus_client.search performs a semantic search using Inner Product (IP) distance. Then, it retrieves top 3 most relevant documents.

Next, we need to initialize the LLM. We will use the replicate service to run LLMs. Replicate allows calling LLMs using easy to use APIs.

We’ll use an open-source model from the available options. In this tutorial, we will use ibm-granite/granite-3.0-8b-instruct.

The ask__LLM query function prepares the query for LLM. Here is the abbreviated code:

def ask_LLM (question, relevant_docs):
    context = "\n".join(
        [doc['text'] for doc in relevant_docs]
    )

    system_prompt = """
    You are an AI assistant. 
    You are able to find answers to the questions \
    from the contextual passage snippets provided.
    """

    user_prompt = f"""
    Use the following pieces of information enclosed in <context> tags \
    to provide an answer to the question enclosed in <question> tags.
    <context>
    {context}
    </context>
    <question>
    {question}
    </question>
    """
    ...
    # parameters to the LLM
    params = {
            "top_k": 1,
            "top_p": 0.95,
            "temperature": 0.1,
            "prompt": user_prompt,
            "system_prompt": system_prompt,

            ...
    }

This code includes:

• system prompt: guides the models behavior and directing it to generate responses only from the provided context.
• User Prompt: Supplies the retrieved documents alongside the user’s question.

• Key Parameters:

  • top_k=1: picks the first result from possible results
  • temperature: controls the 'randomness' or 'creativity' of a model.Higher temperatures makes the model more 'creative'. Smaller values makes the model more deterministic and focused, providing concise and precise answers.

Finally, here is the code to query the LLM:

question = "What was the training data used to train Granite models?"
relevant_docs = fetch_relevant_documents(question)
ask_LLM(question=question, relevant_docs=relevant_docs)

Expected Output:

Granite Code models were trained on 3.5T to 4.5T tokens of code data
and natural language datasets related to code. The code data comprised 
116 languages, and the natural language datasets included high-quality 
data from various domains such as technical, mathematics, and web documents. 
During phase 1, the models were trained solely on code data, while in phase 2, 
additional high-quality language data was included to improve the model's 
performance in reasoning and problem-solving skills.

The model provides a precise response based on the available documents.

Now, let’s test if our model adheres to instructions when asked about something outside our dataset:

question = "When was the moon landing?"
relevant_docs = fetch_relevant_documents(question)
ask_LLM(question=question, relevant_docs=relevant_docs)

Expected Output:

I'm sorry, the provided context does not contain information about the moon landing.

Success! The model correctly limits responses to retrieved documents.

Summary

By completing this tutorial, you’ve learned how to:

• Process and prepare PDFs for a RAG pipeline.
• Store data in a vector database.
• Perform vector searches.
• Query documents using an LLM.

This workflow is a foundation for building AI-powered applications that integrate private knowledge into LLMs efficiently.

Happy coding!