RAG with Python Cookbook

Practical Recipes from Data Preprocessing to LLM Agents

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

Dominik Polzer

RAG with Python Cookbook

by Dominik Polzer

Copyright © 2026 Dominik Polzer. All rights reserved.

Printed in the United States of America.

Published by O’Reilly Media, Inc., 1005 Gravenstein Highway North, Sebastopol, CA 95472.

O’Reilly books may be purchased for educational, business, or sales promotional use. Online editions are also available for most titles (http://oreilly.com). For more information, contact our corporate/institutional sales department: 800-998-9938 or corporate@oreilly.com.

Editors: Jill Leonard and Nicole Butterfield

Production Editor: Aleeya Rahman

Interior Designer: David Futato

Cover Designer: Karen Montgomery

Illustrator: Kate Dullea

April 2026: First Edition

Revision History for the Early Release

2025-03-11: First Release

2025-05-28: Second Release

2025-10-03: Third Release

2025-11-19: Fourth Release

See http://oreilly.com/catalog/errata.csp?isbn=9798341600560 for release details.

The O’Reilly logo is a registered trademark of O’Reilly Media, Inc. RAG with Python Cookbook, the cover image, and related trade dress are trademarks of O’Reilly Media, Inc.

The views expressed in this work are those of the author and do not represent the publisher’s views. While the publisher and the author have used good faith efforts to ensure that the information and instructions contained in this work are accurate, the publisher and the author disclaim all responsibility for errors or omissions, including without limitation responsibility for damages resulting from the use of or reliance on this work. Use of the information and instructions contained in this work is at your own risk. If any code samples or other technology this work contains or describes is subject to open source licenses or the intellectual property rights of others, it is your responsibility to ensure that your use thereof complies with such licenses and/or rights.

979-8-341-60051-5

Brief Table of Contents (Not Yet Final)

Preface (unavailable)

Chapter 1: Getting Started with RAG (Unavailable)

Chapter 2: Generation and Prompt Engineering (Unavailable)

Chapter 3: Loading Data (available)

Chapter 4: Data Preparation (available)

Chapter 5: Embeddings (available)

Chapter 6: Vector Databases and Similarity Search (available)

Chapter 7: Retrieval (available)

Chapter 8: Agentic RAG (Unavailable)

Chapter 9: Graph RAG (Unavailable)

Chapter 10: Evaluation (Unavailable)

Chapter 11: RAG Chatbot (Unavailable)

Chapter 1. Loading Data

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 3rd chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the author at dmnkplzr@googlemail.com.

In traditional companies, about 80% of information is unstructured, buried within PowerPoint presentations, Word documents, Excel files, emails, and voice memos. Figure 1-1 provides an overview of common file types and data sources for structured and unstructured data. This chapter gives you a toolbox for loading all kinds of file types, whether Word, PDF, Excel, CSV, SQL, audio, or images.

Figure 1-1. A common distribution of data in companies.

To robustly process and analyze these unstructured documents, we need highly versatile models that can understand different unstructured data types like text, images, and audio. This is where so-called foundation models come into play. Foundation models are large machine learning (ML) models trained on massive amounts of data, enabling them to mimic specific human abilities, such as understanding language, generating text, recognizing images, or processing audio.

Large Language Models (LLMs), for example, excel at tasks like generating text summaries, translations, and even code. Multimodal generation models go a step further by handling various input formats, including text, audio, images, and videos. However, the real unsung heroes are embedding models. These models have learned the semantic meaning behind text and map it into a multidimensional vector space, which is crucial for many machine learning tasks such as classification, clustering, and recommendation systems.

The capabilities of these models are remarkable, but there are a number of challenges—such as hallucinations, knowledge gaps, efficiency, cost, or a lack of transparency about how answers are generated. The biggest challenge is querying up-to-date information from external knowledge bases. LLMs only know what they learned up to their training cutoff; anything beyond that is either not answered or, more likely, you get a hallucinated response. With the RAG (Retrieval-Augmented Generation) approach, we combine the strengths of LLMs and classic search engines. Answers from RAG systems are:

Traceable: The answer is based on information from a knowledge base that we manage and control ourselves.

Transparent: Users can trace the sources of the answer and evaluate it based on those sources.

Up-to-date: The knowledge base can be updated at any time to include the latest information—no retraining of foundation models required.

With the RAG approach, we retain control over the information we provide to our users and combine it with the strengths of LLMs and multimodal models. Ultimately, we make it easier for users to digest important pieces of information.

Note

Even though the capabilities of foundation models are impressive, they are sometimes overused. Introducing LLMs into an existing process only makes sense if they truly provide a benefit. Google occasionally integrates LLM responses into its search engine. Instead of the classic search results, a summary generated by Google appears at the top as the first result. For me, this use case is definitely a positive example, mainly because Google manages to implement it so seamlessly.

For a typical RAG chatbot, this generation step adds latency. In RAG systems connected to a single knowledge source, the benefit is often limited. For me, the next stage—with significantly greater value—is agentic-like searches, where we break down questions and solve even complex problems step by step using tools and different knowledge sources. Almost every model provider now offers so-called “Deep Research” capabilities. Behind these deep search capabilities is an agentic multi-step approach, often supported by smaller, more cost-efficient models that can answer increasingly complex questions step by step.

The “R” in RAG stands for “Retrieval,” meaning we load content from an external source. While this source is often a vector store, it can also be a graph database, an SQL database, or even a simple Excel file that we pre-filter. The following steps focus on the most popular example of a RAG system, which uses a vector store as the external knowledge base.

Here’s how the RAG flow looks step by step (see Figure 1-2):

Load: Load data from various sources (Word, PDF, Excel, SQL, etc.).

Preprocess: Split the text into smaller chunks. Since embedding models and LLMs have a maximum context size, we divide the text to stay within this limit.

Pre-Retrieval: Translate data into embedding vectors that can later be used to filter relevant content. The idea is that these vectors represent the content and its semantic meaning. Vectors that are close in the vector space are similar. Content that is similar to the original user query is often relevant for answering the question. The most common form of content is text. In this case, we use text embedding models to generate vector representations of the text. These embeddings are vector representations of the data, allowing us to calculate similarities between the text snippets and the user’s question. This enables us to find text chunks that exhibit semantic similarity and are likely to contain related information.

Once all the data is loaded into the vector store, the app can use it during runtime. When a user asks a question, the system follows these steps:

Retrieval: Convert the user’s question into embeddings and use them to search for the most similar text chunks in the vector store by calculating the distance between the vectors.

Post-Retrieval: Construct a meaningful prompt that includes general instructions for the model, such as the specific task, the context in which the user is asking the question, and the desired tone of the response (formal, informal, helpful, direct, etc.). Additionally, include the retrieved text chunks directly in the prompt.

Generation: The LLM analyzes the information from the retrieved text chunks and generates a response that aligns with the instructions in the prompt.

In essence, a RAG system combines a search engine and an LLM. The search engine finds the most relevant information in our database, and the LLM generates the answer to the user’s question.

Figure 1-2. Overview of the components of a RAG system.

While the process may sound straightforward, it gets more tricky when we deal with data from different sources in different modalities. This makes the design of the loading pipelines so essential and requires creativity.

This chapter provides tools and techniques for designing customized loading pipelines.

Warning

Most recipes in this book show you how to implement the components for your RAG system from scratch, using established and mature Python libraries.

While there are ready-to-use orchestration frameworks that combine many of the functions needed for RAG applications, they are still in an early stage and are constantly evolving. This is mainly due to the active open-source community that continues to develop them every day. While the frequent changes can cause issues when our solution is deployed in production, the approaches in these orchestration frameworks serve as a great source of inspiration.

By replicating only the functions, such as text splitter functions, you can avoid dealing with the large number of dependencies that orchestration frameworks come with, as they rely on a wide range of libraries. This can make maintaining your solution more challenging later on, since you need to update and monitor security vulnerabilities for all of them.

If there’s no good reason, I recommend avoiding introducing more libraries to your project than are actually needed. In the best case, I build a whole project using just mature libraries I trust, like pandas, numpy, sqlalchemy, pydantic and requests.

Before we dive in, here is a general note that applies to all recipes in this chapter and book. Since we are using OpenAI models, we will use the openai Python library. The openai library uses the OpenAI API endpoints to access the models. For authentication, you need an API key, which you can generate in the OpenAI console. Just search for “OpenAI API key” and you should find a link to the OpenAI console where you can generate your key. You only need to log in with your account and add some credit to use the models. If you are working with these models for the first time, don’t worry—for simple local tests with small datasets, a few dollars should be enough to get you started.

Warning

Please do not hardcode API keys directly into your code. Even if it’s just for an offline test, you never know if you might later push it to a repository. Checking afterward to ensure that you haven’t left any hardcoded keys or that all keys have been removed from the GIT history is an unnecessary manual effort. It’s better to store your keys in a .env file and load them using the python-dotenv library. Just make sure that the .env file is included in your .gitignore so it doesn’t accidentally get uploaded to a public repository.

1.1 Loading Word Files in Python

Problem

You want to load Word files (.docx, .doc) to preprocess the raw content and load it into the knowledge pool of your RAG system.

Solution

You have a few options here. The most popular library is python-docx, which is great for creating and updating Word files. If you need to differentiate between elements like headings, paragraphs, or listings, check out the unstructured python library. Figure 1-3 shows the basic process. With unstructured, you can break down the document into its structured elements, resulting in a list of raw elements with their respective types and useful metadata.

Figure 1-3. How to load Word files using python-docx or unstructured.

Let’s dive into some examples using both libraries. In Example 1-1, we’ll use the python-docx library to load the Word file. Then, in Example 1-2, we’ll achieve the same result but with the unstructured library, which also breaks down the document into a list of elements, each with its type.

Prerequisites

We will use python-docx and unstructured to load the Word file. In both scenarios, we will use Pandas to store and further process the loaded data.

Python-docx focuses on extracting plain text from the document. The unstructured library goes a few steps further and can also handle images and tables.

Steps

In Example 1-1, we are using python-docx to load the text from the sample Word file Most_Used_Feature_Engineering_Techniques.docx.

To load the document and get its text as a single string, we use the docx.Document() function. This function returns an object with a paragraphs attribute, which is a list of paragraphs. In Example 1-1, we iterate over the paragraphs and concatenate their text into one string.

Example 1-1. Load Word File using python-docx

import os
from docx import Document

file_path = "../datasets/word_files/2023_Jan_7_Feature_Engineering_Techniques.docx"

doc = Document(file_path)

text = []
for paragraph in doc.paragraphs:
    text.append(paragraph.text)

full_text = "\n".join(text)

We get a string with the whole text (see Figure 1-4).

Figure 1-4. The output of the loaded Word file in form of a string

If our RAG system needs to differentiate between elements like titles, text, and lists, we can use the partition_docx function. This function returns a list of objects, each with a text attribute for the element’s text and a type attribute for the element’s type (e.g., Title, NarrativeText, ListItem). In Example 1-2, we iterate over these elements and create a dictionary element_dict for each to store the desired metadata.

Example 1-2. How to load and partition Word files using unstructured

from unstructured.partition.docx import partition_docx
import os
import pandas as pd

elements = partition_docx(filename=file_path)

list_of_elements = []

for element in elements:
    element_dict = {
        "element_id": element.id,
        "file_path": file_path,
        "category": element.category,  # e.g. "Title", "NarrativeText", "ListItem"
        "text": element.text,
        "last_modified": element.metadata.last_modified,
    }

    list_of_elements.append(element_dict)

elements_df = pd.DataFrame(list_of_elements)

The result is a list of elements. Each element is a dictionary containing the category, text, and file path (see Figure 1-5).

Figure 1-5. The output of the loaded Word file using the unstructured library

Discussion

Word is widely used in companies to create project documentation, contracts, and books. Loading text from Word (.docx) files is straightforward. For more complex documents with listings, tables, and images, we can build customized processing pipelines for each element type. For images, we might create summaries explaining the content using LLMs, and for tables, we might summarize key insights.

Knowing the document structure can also optimize the retrieval step in an RAG system. For example, in a book with unrelated chapters, we might first find the best matching chapter title before searching for relevant information within the chapter.

See Also

While this recipe focused on Word files, the Python library unstructured is especially powerful for PDFs. When your loading pipeline becomes more complex, optimizing it for one type of document can make your life easier. Instead of having multiple pipelines for different document types, you could have one pipeline optimized for PDF files, and a step before that pipeline to convert all other document types into PDFs. See Recipe 1.10 for more on handling PDFs.

1.2 Loading PDF Files

Problem

You want to load PDF files (.pdf) into your RAG system.

Solution

Use the PyPDF2 library to load PDFs into your RAG system. This library allows you to extract the text from each page of the PDF. It also allows you to store the exact page number along with the text.

Tip

To load and process PDFs, I often use the unstructured library. Unstructured is described in the documentation as “an open-source toolkit designed to simplify the ingestion and pre-processing of diverse data formats.”

I use it in this book, for example, in Recipe 1.10 to process PDF files with various components such as images, tables, and narrative text effectively.

Alternatively, I frequently use multimodal models with vision capabilities to read PDFs. This approach is more costly but extremely flexible. You simply transform each page of your PDF into an image and let the vision model extract the data. We can instruct the model to directly read narrative text, extract tables by translating them into Markdown format, and extract information from images, for example, by letting the model describe the content.

Figure 1-6 shows the pipeline step by step. First, it loads the PDF and then goes through each page to extract the text and store it along with some useful metadata in a dictionary. After processing all the pages, it returns a list of dictionaries, each containing the text and metadata for a page.

Figure 1-6. Loading PDF files and storing the text and metadata

Prerequisites

The PyPDF2 library is an open-source library for splitting, merging, cropping, and transforming the pages of PDF files. It uses Pillow under the hood to extract images from the PDF.

Steps

Example 1-3 shows how to use PyPDF2 to read a PDF. The function PyPDF2.PdfReader() returns an object with a pages attribute, which is a list of dictionaries. Each dictionary contains the page’s text and metadata.

We can extract the metadata we need by iterating over the pages and filling a customized dictionary for each page. In this example, the dictionary contains the page text, filename, producer, page number, and images.

Example 1-3. Loading PDF Files: Using PyPDF2 to load plain text and metadata

import PyPDF2
import os
import pandas as pd

file_path = "../datasets/pdf_files/2023_Jan_7_Feature_Engineering_Techniques.pdf"

with open(file_path, "rb") as file:
    reader = PyPDF2.PdfReader(file)

    # Initialize an empty string to store the extracted text
    list_of_pages = []
    page_counter = 1

    for page in reader.pages:
        page_dict = {
            "file_name": reader.metadata.get("/Title"),
            "producer": reader.metadata.get("/Producer"),
            "page_number": page_counter,
            "text": page.extract_text(),
            "images": page.images,
        }

        list_of_pages.append(page_dict)

        page_counter += 1

The code snippet returns a list of dictionaries, each containing the text and defined metadata (file name, producer, page number, images on the page). Later in the RAG process, the text from the page is chunked into text chunks and translated into embedding vectors.

Note

We create embedding vectors using so-called embedding models, which map the meaning of words and sentences into a multidimensional vector space. With vectors, we can perform mathematical operations, allowing us to approximate the similarity between text fragments by calculating the distance between the vectors. You can find more details in [Link to Come].

The metadata extracted here will be linked to each text chunk, which allows us to give the correct and detailed references to the information used (see Figure 1-7).

Figure 1-7. The output of the loaded PDF file in form of a data frame

Discussion

PDF is the most commonly used file format for sharing information. Whether the original file was created in Word, PowerPoint, or another tool, it is often converted to PDF before being shared with a broader audience. As a result, most documents are available in PDF format.

Therefore, loading and processing PDF files is crucial for many RAG systems. If you frequently work with orchestration frameworks like LangChain, Haystack, or LlamaIndex, you’ll find that the functions and tools for loading PDFs are often the most advanced and mature.

If your RAG system needs to handle documents in various formats, you can simplify your workflow by optimizing most of your application for PDF documents and adding a processing step at the beginning of the pipeline to convert documents from Word, PowerPoint, Markdown, LaTeX, or other formats into PDFs.

See Also

Loading the pure text is pretty straightforward. However, it gets more challenging when we have PDFs with images and tables. You can find some techniques to handle multimodal content in PDFs in Recipe 1.10.

1.3 Loading and Handling Tabular Data from Excel Files

Problem

You want to load tabular data, such as Excel files, and add them to your RAG system.

For example, this could be sales data that includes sales figures, customer information, and product details. By adding this data to your RAG system, users can analyze it by asking questions.

Solution

We have three fundamentally different ways to handle Excel files in RAG systems. In this recipe, you will find a brief introduction to all three options. The options mainly differ in the “R” of RAG—the retrieval step. In the first example, we prepare each row so that it can be stored as a standalone information piece in a vector database (see “Option 1 - Transforming table rows into text”). In the second example, the retrieval step simply consists of loading the entire table; we do not apply any filtering and just include the whole table in the prompt (see “Option 2 - Embedding the whole table into the prompt”). In the third example, the retrieval engine is a standard SQL database. We first load the data from the Excel file into an SQL database and then use an LLM to generate suitable SQL queries to retrieve the required data (see “Option 3 - Uploading the Excel to a SQL query and use the Text-to-SQL approach”).

The practical example in Example 1-4 focuses on the first option, transforming CSV rows into text.

Option 1 - Transforming table rows into text

You can turn each row of a Excel file into a sentence that makes sense on its own. Think of it like explaining the data to a child who has never seen it before. Take one row, look at the column headers, and describe what the data means in simple words.

Once we have these text snippets, we can treat them like regular text. Finally, we create embeddings (numerical representations) for each snippet and store them in a vector database, as shown in the image Figure 1-8.

Note

This approach is not scalable. More complex analyses are also not possible with it, since the model always only sees the k-nearest neighbors and thus only the information from a few rows. Semantic similarity search will also have problems because the semantic meaning behind the rows is often very similar. In my experience, however, this method is sufficient if, for example, you have a contact list in Excel format with names, addresses, phone numbers, and email addresses. If you search for a name that is explicitly used in the cells, semantic search is good enough. You can also try to improve the search results using a hybrid search approach, which combines semantic and keyword search.

Figure 1-8. How to load Excel files by transforming table rows into text.

Option 2 - Embedding the whole table into the prompt

If the table is not too large, you can add the entire table to the prompt. Most models can interpret tables directly in the prompt, at least when they have the right format. Transform the table into markdown syntax to ensure the LLM recognizes it as a table.

In Figure 1-9, you can see a basic prompt template with an attached markdown table.

Figure 1-9. How to load Excel files by embedding the whole table into the prompt.

Option 1 is straightforward to implement but can only answer questions using information from a few rows. Option 2 can be expensive for large tables and struggles with complex queries. For a more detailed analysis, use the text-to-SQL approach in option 3.

Option 3 - Uploading the Excel to a SQL query and use the Text-to-SQL approach

For more complex queries, like aggregations over the whole table, we can upload the Excel file to an SQL database and use the text-to-SQL approach. The LLM will generate an SQL query to get the needed information from the database based on the user’s question.

Once the data is retrieved, we create a prompt with the user’s question and the data (see Figure 1-10).

Figure 1-10. How to load Excel files by uploading the Excel file to an SQL database and using the Text-to-SQL approach to answer user questions about the inserted data.

The rest of this recipe shows a practical example for option 1.  Example 1-4 is transforming each line of the loaded rows into a text chunk.

Prerequisites

We will use the census income dataset, which is commonly used for binary classification tasks. It is available under the CC BY 4.0 license and thus can be shared and adapted for any purpose.

We will save the dataset as an Excel file containing 15 columns and 48,000 rows. To load the Excel file, we will use the Python library openpyxl. For further processing of the tabular data, we will use Pandas.

Steps

First, we load the Excel file census-income.xlsx into a Pandas data frame. If you have Once loaded, Example 1-4 applies the function create_text_description to each row of the data frame and thus creates a new column text_description with a text snippet explaining the values in the other columns. To create the text chunk, we are using a static template, which you can find in Example 1-4.

Example 1-4. How to transform table rows into text

import os
import pandas as pd

file_path = "../datasets/csv_files/census-income.xlsx"
df_excel = pd.read_excel(io=file_path)

def create_text_description_of_row(row):
    row["text_description"] = (
        f"""The candidate {row['age']} years old is working in the
            {row['workclass']} sector. The candidate was born in
            {row['native-country']}, is {row['marital-status']}
            and has a {row['relationship']} relationship.
            The candidate has a {row['education']} degree
            and is working as a {row['occupation']}.
            The income of the candidate is {row['income']}."""
    )

    return row

# Apply the function create_text_description_of_row to each row of the data frame
df_extended = df_excel.apply(create_text_description_of_row, axis=1)

Figure 1-11 show the content of the newly created column text_description. Each row is a text snippet that describes the values in all the columns.

Figure 1-11. How the list of text snippets looks like after transforming the rows into text

Discussion

LLMs and embedding models are specifically designed to process and analyze text. This is a significant step forward, as it allows us to analyze information from sources such as books, blog articles, emails, and even social media posts, and transform it into structured data.

But what do we do with data that is already in tabular form? We need to prepare it to leverage the strengths of LLMs and embedding models.

This recipe showed some methods to prepare tables for use with LLMs. It’s important to understand the strengths and limitations of each approach.

For small tables, convert them to Markdown format. This helps the model recognize it as a table and include it in your prompt.

For large Excel files, break down the information and describe each row as a separate piece of data. This method is useful for simple questions like Is candidate 25 earning more than 50k? or What is the Education Level of candidate 64?.

For complex questions, like What percentage of candidates with a Bachelor’s degree earn more than 50k?, use a method that can handle aggregations on the entire dataset, such as the Text-to-SQL approach. Libraries like Vanna AI make this straightforward to implement.

See also

In this recipe, we converted each row of the dataset into a text snippet to use with text embedding models. We can use the same approach also for classification tasks. For more details, refer to the paper TabLLM: Few-shot Classification of Tabular Data with Large Language Models.

1.4 Loading structured data from a PostgreSQL Database

Problem

You want to add data from a PostgreSQL database by connecting to the database and executing SQL queries to retrieve the data.

Solution

In Recipe 1.3, I have summarized a few ways to process tabular data. These methods also apply to tables loaded from SQL databases.

The only difference is the way we load the tables. This recipe shows you how to connect to your PostgreSQL database using Python and execute SQL queries to load the data.

Prerequisites

First, install the PostgreSQL server and pgAdmin on your machine. PgAdmin is a popular tool for managing PostgreSQL databases. While you can run all SQL queries directly from Python, pgAdmin makes it easier to browse data and execute queries manually. You can find installation instructions for PostgreSQL at postgresql.org/download/. PgAdmin will be installed by default when you install PostgreSQL.

Note

I rarely work with local database servers anymore; most of the time, I use cloud services (e.g., PostgreSQL via Amazon Aurora). This makes my databases highly scalable from the start, offers high availability, and is relatively cost-efficient. I say “relatively” because it might be cost-efficient compared to other enterprise-grade alternatives, but it can be overkill for small projects or experiments. For small (side) projects, I recommend AWS Lightsail, which offers a fixed monthly price for the services.

If you work with multiple database types, consider using DBeaver or DataGrip. These tools support various databases like MySQL, PostgreSQL, SQLite, Oracle, DB2, and SQL Server.

To connect to and load data from our databases, use SQLalchemy. SQLalchemy is the most used Python SQL toolkit. It gives us the power and flexibility of SQL in Python.

Steps

When you launch pgAdmin for the first time, you can either connect to an existing PostgreSQL server or create a new local server. Follow the steps below to populate and query the database:

Replicate Sample Database: To have some data to work with, we will replicate the sample database from w3schools.com. Follow the tutorial on w3schools.com to fill the database with sample data.

Install Required Libraries: Next, install the Python libraries SQLalchemy and psycopg2-binary to connect to the database.

Create User and Grant Permissions: Create a user and grant the necessary permissions by following the PostgreSQL documentation on creating a user and granting permissions. Store the user credentials in a .env file.

Set Up Connection and Execute Query: With the user created, set up a connection. In Example 1-5, we use this connection to execute an SQL query that retrieves all data from the table categories.

Example 1-5. Loading Data from a PostgreSQL: Connect to the database and execute queries using SQLAlchemy

import os
import pandas as pd
from sqlalchemy import create_engine

connection_string = (
    f"postgresql+psycopg2://{username}:{password}@{host}:{port}/{database}"
)
engine = create_engine(connection_string)

with engine.connect() as connection:
    query = """SELECT * FROM categories ORDER BY category_id ASC """
    result = pd.read_sql(query, connection)
    print(result)

After running the query, the retrieved data will be stored in the data frame results_df. The code snippet fetches all records from the table categories and sorts them by the column category_id.

Discussion

Most of our data is stored in relational databases, and PostgreSQL is one of the most popular ones. Websites, tools, and machines collect and save data in these databases.

When creating a web app for our RAG workflows, we need to load data from or save data to our database. PostgreSQL is a great choice because it can store text data and, with the pgvector extension, it can also store vector embeddings and perform similarity searches. To learn how to generate embeddings, store them in PostgreSQL, and calculate distance metrics, see [Link to Come].

See also

In this recipe, we focused on loading data from PostgreSQL. In other recipes, you will find different parts in our RAG system using PostgreSQL:

If you want to know how to store embeddings in PostgreSQL and perform similarity searches on top of it, you can find it in [Link to Come]

[Link to Come] will use PostgreSQL as an app database for RAG web apps

1.5 Loading Audio Files by Using Speech-to-Text Models

Problem

You want to load audio files to your RAG system.

Solution

To load audio files to our RAG system, we first need to convert them into text. We achieve this using speech-to-text models. These advanced models are trained on thousands of hours of audio in multiple languages and have significantly improved in recent years. They can automatically detect languages and handle loud background noise effectively.

Prerequisites

There are several speech-to-text services available out there, both open-source and commercial. You find some popular ones in the discussion section of this recipe.

I will use OpenAI’s Whisper model for this task. You can either install Whisper locally or use it through OpenAI’s API. The choice depends on your use case:

For quick setup and scalability, use a cloud service like AWS, Azure, GCP, or another platform that offers text-to-speech models.

If you need to handle sensitive data, you can host an open-source model like Whisper on your own servers. This ensures that only you have access to the data.

In Example 1-6, we will use OpenAI’s cloud services. We will load the OpenAI key from a .env file using the python-dotenv library.

Steps

Example 1-6 uses the official OpenAI Python library. Just provide the file_path of the audio file and the model you want to use (in this case, whisper-1).

That’s all it takes. One API call and your product can be labeled as AI-powered.

Example 1-6. How to load and transcribe audio files using Whisper

import os
import openai

audio_file_path = "../datasets/audio_files/harvard.wav"

# initialize the OpenAI client with your API key
client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

with open(audio_file_path, "rb") as audio_file:
    transcription = client.audio.transcriptions.create(
        model="whisper-1", file=audio_file
    )

Discussion

I’m a big fan of OpenAI’s Whisper. Compared to many traditional speech-to-text (STT) models, Whisper is multilingual and remarkably robust against background noise. On top of that, it is fully open source and can theoretically be run on edge devices. This can be helpful, for example, if you want to use speech-to-text models in sensitive environments and need to process data locally. Imagine an ML app on the shop floor that helps identify defective parts—here, it’s extremely useful to have your hands free to log defects while ensuring that highly sensitive data is not sent to third-party services.

Alternatively, there are a number of speech-to-text services available. Of course, every major cloud provider offers a suitable service, whether it’s Google, AWS, or Azure. In addition, there are specialized providers focused on speech-to-text, such as AssemblyAI or Deepgram. Table 1-1 shows a summary of the most popular alternatives to OpenAI’s Whisper.

Table 1-1. Summary of Popular ASR Alternatives

Provider
Key Features
Link

Google Cloud Speech-to-Text

Uses Google’s advanced AI (Chirp model) for accurate speech-to-text, with features like automatic punctuation, speaker separation, and filtering of inappropriate content. It can adapt to specific terms and handle multiple audio channels.

https://cloud.google.com/speech-to-text

Amazon Transcribe

A fully managed service offering accurate speech-to-text with automatic punctuation, custom vocabularies, language detection, speaker separation, and redaction of sensitive information.

https://aws.amazon.com/transcribe/

Microsoft Azure Speech Services

Provides speech-to-text, text-to-speech, and translation. It supports custom models, voice AI agents, real-time translation, and post-call analytics, with flexible deployment options including on-device.

https://azure.microsoft.com/en-us/products/ai-services/ai-speech/

Deepgram

A Voice AI platform offering speech-to-text, text-to-speech, and voice agent APIs. It uses proprietary models for high accuracy and speed, supporting both real-time and batch processing, along with speaker separation and punctuation.

https://deepgram.com/

AssemblyAI

Offers speech-to-text for both recorded and live audio, plus advanced speech understanding. Features include speaker identification, automatic text formatting, and language detection.

https://www.assemblyai.com/

See also

If you need to transcribe not just audio files but also video files, refer to Recipe 1.11. This involves breaking the video into parts and processing each part separately, with the audio track being a key element.

For a detailed tutorial on using Whisper, check out the guide from OpenAI on Hugging Face.

1.6 Extracting Text from Images and PDFs Using Tesseract OCR

Problem

You want to extract text from images or image-like documents using open-source OCR engines you can host on your own infrastructure so you can process sensitive data without sending it to third-party services.

Solution

This recipe explains how to extract text using Optical Character Recognition (OCR) engines. You can also use multimodal foundation models for this task (see Recipe 1.7). Each method has its pros and cons. Multimodal models are flexible and can be easily adjusted by changing the prompt. OCR engines are fast, cost-effective, and can be run on your own servers.

Figure 1-12 illustrates how OCR engines extract text from images, including scanned PDFs. For scanned PDFs, the process involves loading the PDF, saving each page as an image, and then using the OCR engine to extract text from each image. The end result is a list of dictionaries, each containing the image path and metadata for a page.

Figure 1-12. OCR engines can extract text from images, e.g., from photographs, technical drawings, or scanned contracts.

Note

Traditional OCR engines, such as Tesseract OCR, primarily rely on pattern recognition and statistical models to extract text from images. Historically, they struggled with complex images containing tables, symbols, drawings, or diverse fonts. You can find some alternatives in the discussion section of this recipe.

Prerequisites

The most popular OCR engine is Tesseract. You can either install Tesseract locally or use one of the existing cloud services—some of which still use Tesseract as their OCR engine under the hood.

For Example 1-7, we will install Tesseract locally. You can find installation guides for different operating systems in the Tesseract documentation. If you are on Windows, you can use the Windows installer provided by the University of Mannheim.

Once Tesseract is installed, we will install the required Python packages: pdf2image, pytesseract, and pillow.

Steps

Let’s go through two examples. In Example 1-7, you’ll see how to extract text from a sample finance PowerPoint slide, which you can find in the book’s repo.

Example 1-7. Example 1: Loading text from a .png image using Tesseract

import os
from pdf2image import convert_from_path
from PIL import Image
import pytesseract

# Load the sample .png file
image = Image.open(fp="../datasets/images/example_finance_reporting_slide.png")

# Use Tesseracst to do OCR on the image
text = pytesseract.image_to_string(image)

Example 1-8 shows another example, this time applying the same concept to PDFs. The code snippet loads the PDF Most_Used_Feature_Engineering_Techniques.pdf, converts each page into an image, and then uses Tesseract to extract the text from each PDF page.

Example 1-8. Example 2: Loading text from PDF files using Tesseract

import os
from pdf2image import convert_from_path
from PIL import Image
import pytesseract

file_path = "../datasets/images/2023_Jan_7_Feature_Engineering_Techniques.pdf"

# Convert PDF to a list of images
images = convert_from_path(pdf_path=file_path)

text = []
for i, image in enumerate(images):
    page_text = pytesseract.image_to_string(image)
    text.append(page_text)

Discussion

OCR engines are effective for images that primarily contain text, such as scanned documents. They offer a reliable and cost-effective method for extracting text from these images. When converting PDF pages into images, make sure to save the file paths for both the PDF and the images, as well as the page numbers and other relevant metadata. This way, you can reference the original document and page when using the extracted text in your RAG system.

While Tesseract remains a robust and free option for clear, high-quality printed text, it often requires extensive preprocessing and struggles with more complex scenarios like handwritten text, varied image quality, and intricate document layouts. Beside Tesseract there is also a wide range of open-source alternatives. The alternatives listed in Table 1-2 represent a newer generation of open source OCR tools that predominantly leverage deep learning techniques.

Table 1-2. Open-source alternatives to Tesseract

Tool
Key Features
Official Website

EasyOCR

80+ languages, GPU support, Python-based

https://github.com/JaidedAI/EasyOCR

PaddleOCR

High accuracy, layout analysis, handwriting support

https://www.paddleocr.ai/

docTR

Deep learning, modular, PyTorch & TensorFlow support

https://www.mindee.com/platform/doctr

TrOCR

Transformer-based, excels at handwritten text

https://www.microsoft.com/en-us/research/publication/trocr

Tip

If you are dealing with highly sensitive data, you should think about using open-source solutions that you can deploy on your own infrastructure, isolated from external access. In companies, we speak about highly sensitive data when we mean financial data, personal HR data, or product design details.

I rarely use traditional OCR models like Tesseract these days because multimodal models make text extraction so straightforward and are highly flexible. By multimodal models, I mean large transformer-based models with vision capabilities, such as GPT-4.1, Gemini 2.5 Pro, or Claude 3 Sonnet. Even the smaller models in these model families, like GPT-4.1 mini, Gemini 2.5 flash or Claude 3 Haiku, often do a solid job. The smaller models are often up to 10 times more cost effective.

In a single pass, the model can decide whether the image is a scan of continuous text or a photo of people or landscapes, and handle it accordingly. For continuous text, simple extraction is sufficient, while for photographs, a description of the objects in the image is more helpful. The only time I would choose OCR models is when I know the documents are mostly text, have a consistent structure, and need to be processed quickly and cost-effectively.

Note

While multimodal models have impressive OCR capabilities, they can struggle with complex documents like large tables with poor scan quality. They may hallucinate numbers (e.g., reading 4.00 as 4000) or skip table rows and columns entirely.

One solution is using an agentic approach, where multiple models work together iteratively to extract information, each checking and refining the others’ work.

Additionally, it’s often better to ask the model specific questions about the content rather than having it extract all text first. For example, with a product table, directly ask for the name, price and weight of each product and attach the image of the document page.

See also

With multimodal models, we can take it a step further. Depending on the specific images, we can either ask the model to extract just the text (see Recipe 1.7) or to describe every detail of the image (see Recipe 1.8). The models can even decide on their own what makes the most sense for each image.

1.7 Extracting Text from Images using Multimodal Models

Problem

You want to extract text from images or image-like documents in a robust and highly flexible way.

Solution

Recipe 1.6 used OCR engines to extract text from images. In this recipe, we use multimodal models.

The concept is simple. We create a basic prompt and ask the model to describe the image in detail. In Figure 1-13, you can see an example of a prompt and the model’s output. For each image, we aim to get a dictionary containing the extracted text and relevant metadata, such as the image path or creation date.

Figure 1-13. How to extract text from images using multimodal models

Requirements

All we need is a multimodal model. In this example, I’m using OpenAI’s multimodal model, GPT-4o.

Steps

Example 1-9 shows a basic example of a prompt. If you have specific knowledge about the types of images you are working with, use that information to improve the prompt.

Example 1-9. How to extract text from images using multimodal models

import os
from PIL import Image
import base64
import openai

png_file_path = "../datasets/images/example_finance_reporting_slide.png"

with open(png_file_path, "rb") as image_file:
    base64_image = base64.b64encode(image_file.read()).decode("utf-8")

    prompt = (
        "Extract the text from the image attached. Make sure to only "
        "extract only the text. If there is no text in the image, "
        "please return with the sentence 'No text found in the image."
    )

    response = openai.chat.completions.create(
        model="gpt-4o",  # define the model to use
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    },
                ],
            }
        ],
        max_tokens=500,
    )

    content = response.choices[0].message.content

Discussion

OCR with multimodal models works well for small-scale use, but becomes challenging when processing thousands of documents. While powerful, these models can be expensive and take a while to process. At current prices, processing a 100-page document can cost anywhere from a few cents to a dollar depending on the model used. This can become problematic when trying to digitize legacy data pools containing hundreds of thousands of documents.

Nevertheless, I still prefer using multimodal models today despite the higher costs. Just keep in mind that even flagship models have limitations. Scans with tables containing 20+ rows and many columns are often not fully processed. The model may hallucinate numbers when they are difficult to read. These limitations should be considered and either mitigated or at least made transparent to users - they need to understand that the results are AI-generated and not always 100% accurate.

See also

In Example 1-9, we used a multimodal model to extract text from images. What makes these models so powerful is their flexibility; by simply adjusting the prompt, we can use them for various preprocessing tasks, for example:

Recipe 1.8 shows how to use multimodal models to extract text and create detailed summaries for images.

Recipe 1.9 demonstrates how to use multimodal models to extract key insights from tables and generate text summaries.

1.8 Generating Text Description for Images Using Multimodal Models

Problem

You want to extract insights from images by generating detailed image descriptions. This allows us to use text embedding models to create embeddings that effectively represent the image content.

Solution

Figure 1-14 shows the process. The concept is straightforward: We create a prompt with specific instructions for the model and attach the image we want to analyze and describe. With the right prompt instructions, the model generates a detailed description of the image and the information within it. Now that we have a text description of the image, we can create embeddings using the same text embedding model that we use for other text content in our RAG system.

Figure 1-14. Generating text summaries for images using multimodal models.

Note

There are different ways to handle images and incorporate them into the retrieval step. The method described in this recipe, where we first use multimodal models to translate the content of the image into text and then generate embeddings using a text embedding model, is just one approach.

Another popular option is to use embedding models that can directly handle multimodal content, such as OpenAI’s CLIP (Contrastive Language-Image Pre-training). These models can directly translate images into embeddings without the need to extract text first.

Prerequisites

We need a multimodal model that can understand and analyze images. In Example 1-10, we use GPT-4o from OpenAI. Alternatively we can use Anthropic’s 3.5 Sonnet or Google’s Gemini 1.5.

Steps

Example 1-10 shows how to load and analyze a photo of Vietnam’s capital city. We then give the model simple instructions on how to describe the image in detail.

Example 1-10. Generate Text Summaries for Images: Load the image and generate text summaries

import base64
import openai

image_path = "../datasets/images/vietnam.png"

with open(image_path, "rb") as image_file:
    base64_image = base64.b64encode(image_file.read()).decode("utf-8")

    prompt = (
        "You are an assistant for visually impaired users. "
        "Describe the image in detail."
    )

    response = openai.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    },
                ],
            }
        ],
        max_tokens=150,
    )

    content = response.choices[0].message.content

Discussion

With the rise of multimodal models, many traditional technologies are becoming outdated. Tasks that used to require specialized NLP or computer vision models can now be handled by multimodal models with the right prompt.

However, there are trade-offs:

Smaller, specialized models like OCR engines are fast and cost-effective. They can be hosted locally, which is great for sensitive data.

Multimodal models are powerful and versatile. They can handle complex images, extract text, or describe details. These models are flexible and can adapt their output based on the image type. For example, they might focus on facial expressions in a photo of a person or on process steps in a technical chart. Figure 1-15 shows different types of PowerPoint slides. Depending on the type of slide, different details are important to understand the chart.

Figure 1-15. Extracting text from different image types

What’s best for your use case depends mainly on the types of images and the amount of data you need to process. I recommend using a multimodal model when handling only a few hundred images.

See Also

Multimodal models are versatile because they are not designed for a specific use case but aim to replicate the human ability to generate new text. They have developed this skill by training on a large text corpus, providing them with a vast range of knowledge.

By asking the right questions, we can use those capabilities for a wide range of data processing tasks, including simple classification and clustering.

Here are some examples from the book:

Recipe 1.9: The model generates text summaries for tables.

Recipe 2.3: The model reads a text snippet and creates hypothetical questions based on the information.

The most advanced use is the agentic RAG concept, which mimics human problem-solving by iteratively refining the approach (see [Link to Come]).

Tip

In the recipe, we generated summaries for images. Once generated, we can create embeddings for them using text embedding models. This process converts the image information into a vector format.

Alternatively, we can directly create embeddings for the images using models that bridge the gap between image and text content, such as OpenAI’s CLIP (Contrastive Language-Image Pre-training). For more details, see Recipe 3.5, which explains how to use CLIP for a multimodal RAG application.

1.9 Generating Text Summaries for Embedded Tables Using Multimodal Models

Problem

You want to extract insights from tables and translate them into text summaries. This allows us to generate embeddings using a text embedding model.

Solution

In our RAG system, we store content as standalone pieces of information. Most RAG systems store these pieces in a vector store and analyze the semantic similarity between the user’s question and the stored content. Since tables do not contain complete sentences, semantic search can struggle to identify the content in the table as relevant.

Imagine you have a table with a company’s quarterly results (e.g., “Revenue: $100M, Profit: $10M”). A sample question might be: “What strategic challenges does XYZ AG face due to recent product developments, and how do these affect its long-term outlook?” If the business results show that the numbers worsened after the release of the latest product iteration, the revenue and profit figures may be important for answering the question. We can increase the likelihood that our semantic search identifies the table’s content as relevant by comparing the user’s query to table summaries instead of the raw table data. For example, the summary could be: “The revenue of company XYZ decreased by 20%, and profit decreased by 30% compared to the previous quarter.”

Figure 1-16 illustrates the process of using LLMs to automatically generate these table summaries.

Figure 1-16. How to generate text summaries for tables using multimodal models.

Prerequisites

We use the unstructured library to break down the PDF into text, tables, and images. This library relies on PDFMiner, a tool for extracting text from PDF files. For processing, we use OpenAI’s GPT-4o model through the OpenAI Python library.

Steps

Example 1-11 uses unstructured to partition the PDF and extract the text and tables. We will save the extracted tables in a Python list called tables.

Example 1-11. How to generate text summaries for tables: Partition the PDF into text and tables

import os
from unstructured.partition.pdf import partition_pdf

pdf_file_path = "../datasets/pdf_files/adult_data_article.pdf"

tables = []
texts = []

# partition the PDF file into its elements
raw_pdf_elements = partition_pdf(
    filename=pdf_file_path,
    strategy="hi_res",
)

for element in raw_pdf_elements:
    if "unstructured.documents.elements.Table" in str(type(element)):
        tables.append(str(element))

This extracts all tables in the PDF

The next step is to summarize the table using a suitable LLM. In Example 1-12, we will use OpenAI’s GPT-4o. The prompt is a simple question asking the LLM to highlight the key points of the table. You can customize this prompt to better fit your specific needs. Give it a try.

Example 1-12. How to generate text summaries for tables

from openai import OpenAI
import pandas as pd

def summarize_tables(row):
    summary_prompt = f"""You are an assistant tasked with summarizing tables. \
                    Give a concise summary of the table. Table chunk: {row.table}"""

    # Initialize the OpenAI API client and generate the table summary
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": summary_prompt}],
        temperature=0.7,
        max_tokens=150,
    )

    row["table_summary"] = response.choices[0].message.content

    return row

# create a pandas dataframe from the tables
tables_df = pd.DataFrame(tables, columns=["table"])

# add a column to the dataframe to store the summaries
tables_df = tables_df.apply(summarize_tables, axis=1)

Discussion

The LLM describes the table in a way that makes it interpretable as a standalone piece of information, making it suitable for our text embedding models.

When storing the summary and creating embeddings, we need to store the link to the original table. When the retrieval step of our RAG system identifies the key insights from the table as relevant to answering a user’s question, it makes sense to include the whole table instead of just the summary in the final prompt. The LLM should be capable of analyzing the raw table. This way, we don’t lose any information along the process.

See also

Today, we use LLMs and multimodal models for many data processing tasks. These models can interpret data, rephrase text, summarize images or text, and even make decisions. We can use those models in countless use cases, often just requiring a different prompt. Here are some examples from this book:

Recipe 1.8 shows how to use multimodal models to create text summaries for images.

Recipe 2.3 demonstrates how to show the model a text snippet and ask it to generate hypothetical questions based on the information.

The most advanced application is the agentic RAG concept, which mimics human problem-solving by iteratively refining the approach (see [Link to Come]).

1.10 Parsing PDFs with multiple media content using Unstructured and multimodal models

Problem

You want to load a PDF with multimodal content, such as images and embedded tables, into your RAG system.

Solution

Multimodal models open new doors to process and analyze documents with images and tables.

Earlier in the book, we saw how to generate text summaries for images in Recipe 1.8 and how to generate text summaries for tables in Recipe 1.9. In this recipe, we use both concepts to interpret the different modalities of a PDF document.

Figure 1-17 shows the full process of loading a PDF with multimodal content. It requires a multi-step pipeline with the following steps:

Partition Data: Load raw data and partition it into text, images, and tables

Generate Summaries: Use a multimodal model (such as OpenAI’s GPT-4V, Anthropic’s 3.5 Sonnet or Google’s Gemini 1.5) to generate text summaries for the images and tables

Create Embeddings: Generate text embeddings for the text chunks, image summaries, and table summaries

Store Embeddings: Add the embeddings to the vector store. Store the metadata, like timestamps, file paths, page numbers, etc.

Figure 1-17. How to load and process multimodal PDFs.

Prerequisites

We need the unstructured library to break the PDF into text, images, and tables.

The unstructured library relies on the Poppler library to process PDF files, especially for tasks like extracting images or determining page counts.

How you install Poppler depends on your operating system. For Windows, I recommend the following steps:

Go to the Poppler for Windows website.

Download the latest release ZIP file (e.g., poppler-23.05.0.zip).

Extract the contents of the ZIP file to a directory on your computer, such as C:\Program Files\Poppler.

On Ubuntu or Debian, it is even easier. You can install Poppler with the following command:

sudo apt-get install poppler-utils

Note

To handle image-based documents like scanned PDFs, unstructured uses an OCR engine. In Example 1-13, we define tesseract as the OCR engine, so you need to install it first. You can find more details on how to install Tesseract in Recipe 1.6.

Steps

In Example 1-13, we are using the unstructured library to partition the PDF into its elements. Each element has a particular category: narrative text, titles, images, or tables.

Example 1-13. How to partition the PDF into text, images, and tables

from unstructured.partition.pdf import partition_pdf
import os

# set the OCR agent to tesseract
os.environ["OCR_AGENT"] = "tesseract"

pdf_file_path = "../datasets/pdf_files/adult_data_article.pdf"
image_output_dir = "../datasets/extracted_content_from_pdfs/images"

# get elements using the function extract_pdf_elements
raw_pdf_elements = partition_pdf(
    filename=pdf_file_path,
    extract_images_in_pdf=True,
    extract_image_block_types=["Image", "Table"],
    extract_image_block_to_payload=False,
    extract_image_block_output_dir=image_output_dir,
)

# categorize elements by type
tables = []
texts = []
titles = []

# fill the just created lists with the elements
for element in raw_pdf_elements:
    if "unstructured.documents.elements.Table" in str(type(element)):
        tables.append(str(element))
    elif "unstructured.documents.elements.NarrativeText" in str(type(element)):
        texts.append(str(element))
    elif "unstructured.documents.elements.Title" in str(type(element)):
        titles.append(str(element))

We get separate lists for each type of element. Here are some ideas for processing each type further:

Text and Titles: Split them into smaller text chunks suitable for embedding models. This process involves techniques like recursive, semantic, and agentic chunking, which are detailed in Chapter 2.

Images: Extract text using OCR engines (see Recipe 1.6) or multimodal models (see Recipe 1.7). Alternatively, generate text summaries for the images using multimodal models (see Recipe 1.8).

Tables: Either load the plain text from the table or, preferably, extract key insights using multimodal models (see Recipe 1.9).

Discussion

Many PDFs contain a mix of text, images, and tables. While humans can easily understand these elements, RAG systems often rely on text embedding models that excel at processing text but struggle with visual content. To bridge this gap, we can use multimodal models to analyze and summarize images and tables into text. This way, the text embeddings can capture the essential information from all elements.

See also

For more advanced loading pipelines for RAG systems, check out the blogs from major RAG frameworks and model providers like LangChain, LlamaIndex, OpenAI, or Anthropic.

Here are two great blog posts to get you started:

How to parse PDF docs for RAG

Multi-modal RAG Notebook Tutorial from LangChain and Semi-Structured RAG Notebook Tutorial from LangChain.

1.11 Loading Videos Using Speech-to-Text and Multimodal Models

Problem

You want to load videos to your RAG system to be able to answer questions about the video content.

Solution

We don’t directly convert the video into vector data. Instead, we break it down into parts and turn each part into meaningful text chunks. These text chunks are then used to create embeddings with text embedding models.

In Example 1-14, I’ll show you how to load a video tutorial “Learn Data Science”. This is similar to many video meetings where a presenter explains something while showing PowerPoint slides.

Figure 1-18 illustrates the steps in the loading pipeline:

Partitioning: Define timestamps for when the screen content changes significantly. Image processing libraries can help detect these changes. Commonly used techniques include pixel-based methods (using scikit-image), perceptual hashing (using imagehash), and deep learning-based approaches (using tensorflow).

Save frames and create text summaries: At each timestamp, save the video frame as an image and create a text summary (see Recipe 1.8).

Transcribe audio using Speech-to-Text models: Extract the audio segment between timestamps and save it as an mp3 file. Use the “Whisper” Speech-to-Text model to transcribe these audio files into text (see Recipe 1.5).

Generate embeddings for frames and audio: Use text embedding models to create embeddings for the text summaries of images and transcribed audio.

Store text, images, and audio in the vector database: Save the embeddings and metadata in the vector database.

Figure 1-18. How to load videos into a RAG system by breaking them into their elements and processing them using speech-to-text and multimodal models.

Prerequisites

We will use the moviepy library to load the video, extract frames, and save audio clips. Moviepy is a Python library for video editing.

To process the video further, we will use two types of models:

Speech-to-Text Model: OpenAI’s Whisper model will convert the audio clips into text.

Multimodal Model: OpenAI’s GPT-4o model will create text summaries for the images.

Steps

Since we want to extract the content from the audio track and the frames, we first need to partition the video in a way that allows us to extract the most important information from the frames while also associating the audio track with the corresponding frame. To achieve this, we need to define the timestamps at which the content on the screen changes significantly.

We can, of course, do this manually—a human goes through the video step by step and defines the timestamps at which frames should be saved. However, it is more common to design a pipeline that analyzes the video and automatically defines the timestamps. For this, we need to consider the characteristics of the video.

Is it a single speaker talking, or a discussion between multiple people? Are there static slides on the screen or moving animations? How often do they change? If it’s a simple slideshow, how frequently do the slides change?

The video we are loading in Example 1-14 is a beginner Data Science tutorial. The speaker explains basic concepts using a slide presentation. The content on the screen changes infrequently, about once every 30 seconds.

Steps 1 and 2: Partition the video into sequences and save the frames as images

It makes sense to analyze the video (frames and audio track) to find the best timestamps for splitting the video and saving frames. For Example 1-14, we simplify the pipeline by defining a random list of timestamps. At each of these timestamps, we will use moviepy to save a frame from the video as an image.

Figure 1-19. How to extract frames from the loaded video and save them as images

Example 1-14. Loading video files, extracting frames and saving them as images

import os
import pandas as pd

from moviepy import VideoFileClip, TextClip, CompositeVideoClip

video_file_path = "../datasets/videos/learn-data-science-tutorial.mp4"
image_output_folder = "../datasets/videos/video_extracted_images"

clip = VideoFileClip(video_file_path)

# create a list of timestamps from which we want to extract a frame
time_step = 10  # time in seconds
timestamps = list(range(0, int(clip.duration) - time_step, time_step))

# for each timestamp extract a frame
for timestamp in timestamps:
    frame_image_path = os.path.join(image_output_folder, f"frame_{timestamp}.png")
    clip.save_frame(frame_image_path, t=timestamp)

We are getting a list of images, each representing a frame from the video (see Figure 1-20).

Figure 1-20. Screenshot: Frames extracted from the video

After running the code, you’ll have a collection of images in the specified directory, each representing a frame from the video.

Now, let’s decide how to process the saved images. Typically, we use multimodal models to describe the content. Here are some methods covered in earlier recipes:

Recipe 1.6 explains how to extract text from images using OCR.

Recipe 1.7 shows how to extract text from images using a multimodal model.

Recipe 1.8 demonstrates how to generate text summaries for images.

Step 3: Extract the audio between each of the timestamps and save it as an mp3 file

Next, we’ll do the same for the audio track. Example 1-15 iterates over the list of defined timestamps and saves the audio segments between two timestamps as mp3 files (Figure 1-21).

Figure 1-21. How to save audio sequences from the video to mp3 files.

Example 1-15. Extracting the audio between two timestamps and saving it as an mp3 file

# for each timestamp extract the audio sequence and save it to a .mp3 file
audio_output_folder = "../datasets/videos/video_extracted_audio"

for timestamp in timestamps:
    audio_clip = clip.subclip(timestamp, timestamp + time_step).audio
    output_audio_path = os.path.join(audio_output_folder, f"audio_{timestamp}.mp3")
    audio_clip.write_audiofile(output_audio_path)

After splitting and saving the audio segments, we transcribe each audio file using the Whisper speech-to-text model. For more details, see Recipe 1.5.

Figure 1-22. Transcribe the audio track between two saved frames from the “Data Science Tutorial” video.

The remaining tasks in Steps 4 and 5 are generating the embeddings and storing them along with the metadata in the vector database.

Discussion

I’ve learned half of what I know today from YouTube videos. I also spend several hours a day in video calls, many of which are recorded. But honestly, I rarely go back to watch old meeting recordings or YouTube videos unless I have a specific reason.

RAG systems can help extract useful insights from videos and turn them into a knowledge base that you can interact with. By breaking videos into sequences and frames, we can link directly to the original content and provide exact timestamps, making it much easier to revisit old videos.

Many companies use video tutorials to upskill their organization on new tools or onboard new employees. I frequently use this concept to add the content from those tutorials to my RAG systems.

Chapter 2. Data Preparation

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 4th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the author at dmnkplzr@googlemail.com.

In RAG systems, we break down longer texts into smaller chunks to make them manageable for text embedding models. These embeddings capture the meaning of the text as vectors in a multi-dimensional space (see Recipe 3.1 for how to create these embeddings). Later, we use these vectors to measure how similar two text chunks are by calculating the distance between them (see Recipe 3.3 for details).

This way, the RAG app surfaces pieces of information that are likely relevant to answer the user’s question.

To ensure effective retrieval, we create text chunks that capture single pieces of information. The art is preparing chunks so they stand alone and remain clear without surrounding context. A strong data processing pipeline cleans the raw text and splits it at the right points to produce meaningful chunks. Figure 2-1 shows some techniques for these steps.

Here are the steps in the pipeline:

Clean the raw text: Remove noise (unnecessary characters, typos, formatting issues) so chunks are readable by humans and models. One specific point worth calling out is handling abbreviations; see Recipe 2.2.

Collect and store metadata: Gather metadata, like the author’s name or the document’s creation date, to help filter out irrelevant information during retrieval. See Recipe 2.1 for how to collect metadata fields from the document and the text itself.

Split the text into smaller chunks: The simplest method is using a character splitter (see Recipe 2.4), which divides the text into fixed-size chunks. However, it’s better to use a method that considers the document’s structure, like recursive chunking (see Recipe 2.5), which splits the text at paragraph or heading breaks. For unstructured documents, like audio transcriptions or chat histories, try semantic chunking (see Recipe 2.7) or agentic chunking (see Recipe 2.8). These methods analyze the content to decide where to split the text, though they are more complex and costly than recursive chunking. The goal is to produce chunks that each state one clear idea or fact.

Figure 2-1. The simplified RAG indexing pipeline including potential data processing techniques

2.1 Adding Metadata to Enable Metadata Filtering

Problem

You want to store useful metadata with each text chunk in your vector store. This helps you apply filters when searching for relevant content, making the retrieval step in your RAG system faster and more accurate.

Solution

We can extract and generate metadata in two ways: from the document’s stored metadata or by using LLMs to scan the document text. Figure 2-2 shows the three steps to build a detailed metadata dictionary:

Extract metadata: Extract the existing metadata from the document. See “Step 1: Extract the metadata stored in the document”.

Add metadata: Add extra metadata like file location, size, or page count. Refer to “Step 2: Add extra metadata like extraction time, file size, and more”.

Generate new metadata: Use LLMs to find metadata in the document’s text, such as the author’s name and contact details. For contracts, look for specific clauses. See “Step 3: Generate new metadata by analyzing the document text” for details. The figure below shows the full process.

Start by checking the existing metadata in the document and then fill in the gaps with information from the text. PDFs store metadata in XML format, which often includes details like the author, creation date, and last modification date. However, they usually lack a description of the document type. For example, in a RAG system for HR analyzing resumes, it would be helpful to label each resume by job type, such as software developer, graphic designer, or marketing specialist. “Step 3: Generate new metadata by analyzing the document text” shows how to use LLMs to fill this gap by analyzing the file content. This way, if a user asks, "What skills are commonly mentioned in software developer resumes?“, the system can filter out non-developer resumes.

First and foremost, retrieval looks for semantically matching text chunks; metadata filtering simply narrows the search space.

Figure 2-2. How to extract and generate metadata fields

Let’s see how this looks step by step.

Prerequisites

This recipe uses the PyPDF2 library to extract metadata from PDFs and the OpenAI API to extract metadata from the document text.

pip install PyPDF2 openai pydantic

Steps

Following the process shown in Figure 2-2, we have three steps to build a comprehensive metadata dictionary:

Step 1: Extract the metadata stored in the document

In Example 2-1, we use the PyPDF2 library to load the PDF and extract the metadata stored within it. Common metadata fields include the author, creation date, last modification date, creator (software used to create the document), producer (software used to convert the document to PDF), title, and subject. Not all documents will have all these fields, but this is a good starting point to check what’s available.

Example 2-1. How to extract metadata from a PDF

    import PyPDF2
    import os

    file_path = "../datasets/pdf_files/attention_is_all_you_need_paper.pdf"

    with open(file_path, "rb") as file:
        reader = PyPDF2.PdfReader(file)
        metadata = reader.metadata

        text = ""
        for page in reader.pages:
            text += page.extract_text()

Screenshot.

Step 2: Add extra metadata like extraction time, file size, and more

Besides the standard metadata fields, you can include any other useful information. Example 2-2 demonstrates how to add values for the PDF’s page count, file size, file name, file path, and the length of the extracted text.

Example 2-2. How to add customized metadata fields to the metadata dictionary

    metadata = dict(metadata)
    metadata["page_count"] = len(reader.pages)
    metadata["file_size"] = os.path.getsize(file_path)
    metadata["file_name"] = os.path.basename(file_path)
    metadata["file_path"] = file_path
    metadata["text_length"] = len(text)

Screenshot.

Step 3: Generate new metadata by analyzing the document text

When you run the code in Example 2-2, you’ll notice that the author’s name fields are empty. However, these names are usually present somewhere in the document, especially in scientific papers. To extract this information, we need to analyze the document text directly.

Example 2-3 demonstrates how to use LLMs to find and fill in this metadata. We load the entire text and let the LLM search for the relevant information. To ensure the LLM provides only the required metadata and avoids any irrelevant details, we use OpenAI’s structured output functionality. This involves defining a Pydantic model that specifies the expected data fields and their formats. Pydantic is a popular tool for data validation, helping us clearly outline the expected values and formats.

Additionally, we need to instruct the model on what to do. These instructions, known as the system message, are placed at the beginning of the prompt.

Figure 2-3 illustrates the steps to generate a structured output using a Pydantic model.

Figure 2-3. Generate metadata from the document text by using LLMs

Example 2-3 processes a scientific paper where the author names are not in the PDF’s metadata. We use a Pydantic model called AuthorContact to specify the data fields we want, like author names, email addresses, and company names. Since there can be multiple authors, we use another model called Contacts to hold a list of AuthorContact objects. The LLM scans the text and extracts the specified fields.

Note

OpenAI’s “Structured Outputs” feature allows LLMs to generate responses that follow a predefined JSON schema or Pydantic model. This is extremely useful for extracting structured data from text documents. Beyond that, I frequently use it for decision-making tasks with LLMs. When such a decision step is part of a larger workflow, it’s important that the LLM’s response is unambiguous. For example, if you need a clear yes-or-no decision, you can define a boolean data type so that the output is strictly True or False, and nothing else.

Example 2-3. Generate metadata from the document text using LLMs

    from pydantic import BaseModel
    from openai import OpenAI

    client = OpenAI()

    class AuthorContact(BaseModel):
        name: str
        company: str
        email: list[str]

    class Contacts(BaseModel):
        entries: list[AuthorContact]

    system_message = """Extract the contact information of all authors."""

    completion = client.beta.chat.completions.parse(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": system_message,
            },
            {
                "role": "user",
                "content": text,
            },
        ],
        response_format=Contacts,
    )

    author_contacts = completion.choices[0].message.parsed

    metadata["author_contacts"] = author_contacts

The code snippet uses the LLM to extract author information directly from the text itself. Since there can be multiple authors, we receive a list. This list is added as a new field in the metadata dictionary. This gives us a comprehensive metadata dictionary that combines fields extracted from the file itself and fields extracted from the document text.

In Figure 2-4, you can see the complete metadata dictionary, including the new field author_contacts.

Figure 2-4. Screenshot of generating new metadata fields using LLMs

Discussion

Metadata helps us narrow down the search space in our vector database before performing a vector search. This can speed up the search and filter out irrelevant information. Not every text chunk that shows high semantic similarity to the user’s question is actually relevant. Here’s what I mean.

For the user query “What are the benefits of antioxidants?”, Figure 2-5 shows two statements that seem relevant but point to different contexts:

If you are a doctor, you are interested in the benefits of antioxidants for the body.

If you are an engineer, you are interested in the positive impact of antioxidants on the characteristics of raw materials.

You will rarely have a user ask for other scenarios. Thus, information that seems correct and relevant at first glance can be irrelevant or even wrong when we know the full context behind the user’s question.

Figure 2-5. How two semantically similar information sets can be contradictory

The example in Figure 2-5 is extreme, but the issue is common when building RAG apps. We often deal with documents that seem related at first glance, but have nothing to do with each other when you look closer.

Other filter criteria could be the period the documents were published, the authors’ names, whether certain names were mentioned in the text, or whether contracts contain specific clauses. What filter criteria could be helpful really depends on your specific use case.

2.2 Enhancing Data Quality by Replacing Abbreviations and Technical Terms

Problem

You want to make text chunks more self-explanatory by replacing domain-specific abbreviations. This ensures that our embedding model and LLMs clearly understand the meaning behind the text snippets.

Solution

Text chunks are segments of a larger text that can be difficult to understand when viewed in isolation.

Figure 2-6 illustrates some challenges that the RAG system may face without proper context.

Challenge 1: Your document contains many domain-specific abbreviations, making it hard to understand when only parts of the text are visible. For example, in the first column of Figure 2-6, purchase order is abbreviated as PO in a supply chain document. However, depending on the context, this abbreviation could also mean Product Owner or Post Office. Similarly, in the sports example in the middle of the figure, MJ stands for Michael Jordan. If you are not so much into basketball, you might not recognize who MJ is—it’s not always obvious. Learn how to address this in Recipe 2.2.

Challenge 2: You are dealing with documents that contain many images, such as tutorials with screenshots. This can be challenging because LLMs and most RAG apps primarily process text. For instance, a tutorial might refer to a “settings” button without showing its location or appearance. Learn how to address this issue using multimodal models in “Challenge 2: Handling documents with many images”.

Figure 2-6. The way that documents are structured influences the performance of our RAG system

These challenges are so nuanced, that there is not one solution that fits all. We need to tackle them step by step. The rest of this recipe shows a few ways to handle abbreviations and image-heavy content.

Prerequisites

pip install openai python-dotenv requests

Steps

The following sections demonstrate different approaches to handle common challenges when preparing text data for RAG systems.

Challenge 1: Handling documents with many domain-specific abbreviations

Example 2-4 demonstrates how to use a predefined abbreviation dictionary to replace abbreviations in text chunks.

For example, consider a blog post about transformer models filled with abbreviations like LSTM (Long Short-Term Memory) and FFN (Feed-Forward Network). Even experienced data scientists might not know all these terms. The process in Example 2-4 replaces abbreviations with their full terms and adds the abbreviation in brackets, e.g., LSTM becomes Long Short-Term Memory (LSTM). This ensures that our RAG system can understand user queries regardless of whether they use abbreviations or full terms: What is LSTM? vs. What is Long Short-Term Memory?.

Example 2-4. Enhancing text chunks by replacing abbreviations

    import re

    abbreviations_dict = {
        "NLP": "Natural Language Processing",
        "RNN": "Recurrent Neural Network",
        "LSTM": "Long Short-Term Memory",
        "GRU": "Gated Recurrent Unit",
        "TF": "Transformer",
        "MHA": "Multi-Head Attention",
        "FFN": "Feed-Forward Network",
    }

    # Load the sample text file
    file_path = "../datasets/text_files/blog_post_transformers.txt"
    with open(file_path, "r") as file:
        text = file.read()

    # Replace abbreviations in the text
    for abbr, full_form in abbreviations_dict.items():
        text = re.sub(rf"\b{abbr}\b", f"{full_form} ({abbr})", text)

The code snippet searches for a list of clearly defined abbreviations and replaces them with their full forms. However, sometimes we encounter abbreviations that appear in multiple variations. For example, the term “versus” often shows up as:

vs. (Common, with a period).

vs (Common, without a period).

v. (Often used in legal contexts, with a period).

v (Less common, but seen).

Example 2-5 shows an example of how you can use regular expressions (regex) to find these different forms of abbreviations in a single pass.

Example 2-5. Enhancing text chunks by replacing abbreviations using regular expressions

import re

text = """In the match vs. Bayern, the team played well. The legal case was Müller v.
Schmidt. We also saw some ads saying Product A vs Product B. Sometimes people just write v or vs."""

matches = re.findall(r'\b(?:v\.?|vs\.?)\b', text, re.IGNORECASE)

print(matches)

Figure 2-7 shows the regex matches found in the text.

Figure 2-7. Screenshot of regex matches for “v” / “vs” variations

Challenge 2: Handling documents with many images

For more complex tasks, we need to use LLMs and multimodal models. LLMs and most basic RAG apps primarily process text. Example 2-4 loads the text from a PowerPoint slide about financial figures. You can view the original slide here: finance slide image. The content is filled with numbers and finance-specific abbreviations. Without a finance background and the accompanying graphs, it’s challenging to fully grasp the slide’s information.

We use OpenAI’s multimodal model GPT-4o to interpret the content, add context, and replace abbreviations and technical terms to make it easier to understand.

Example 2-6. Enhancing text chunks using LLMs

    import os
    from openai import OpenAI

    file_path = "../datasets/text_files/EMEA_drives_revenue.txt"

    with open(file_path, "r") as file:
        text = file.read()

    prompt = f"""
        The text below contains a financial report including a lot of abbreviations and
        technical terms from the finance domain. Please replace the abbreviations with
        their full forms and provide a brief explanation of the technical terms, so the
        whole text get's easier to read and understandable for everyone.

        Make sure it's easy enough, that a 10 years old school kid could understand it.

        Often used abbreviations are:
        - EMEA: Europe, Middle East, and Africa
        - BD: Business Development
        - YoY: Year-over-Year
        - APAC: Asia-Pacific

        Text:
        {text}
        """.strip()

    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="gpt-4o",
    )

    enhanced_text = chat_completion.choices[0].message.content

Figure 2-8 shows how the text chunks look after processing. On the left side is the raw extracted text from the slide; on the right side are the enhanced text chunks after replacing abbreviations and technical terms.

Figure 2-8. Screenshot of enhanced text chunks after replacing abbreviations and technical terms

Discussion

In an ideal world, text chunks would be easy to understand on their own. This would allow anyone to grasp the content just by reading the snippet. However, in reality, we often deal with complex texts filled with domain-specific abbreviations and references to other parts of the document.

Additionally, most documents require some background knowledge to fully understand the information. For example, reading an advanced Python book assumes you already know basic Python concepts. If you encounter something you don’t understand, you might refer to a beginner book to fill in the gaps. Gradually, you build up the knowledge needed to fully comprehend the document. Ideally, text chunks would be self-explanatory without requiring too much background knowledge.

In a basic RAG system, the LLM only sees the information retrieved from the vector search. This recipe showed techniques to replace technical terms, making text chunks easier to understand. A few practical tips:

Adding contextual sentences: Instead of “It improved efficiency by 40%.”, use “The new indexing algorithm improved search efficiency by 40%.”

Avoiding pronouns and implicit references: Instead of “It was implemented last year.”, use “The customer feedback system was implemented last year.”

Preserving full names of entities: Instead of “Google launched it in 2019.”, write “Google launched the BERT language model in 2019.”

Avoiding ambiguous words: Instead of “The system works better now.”, write “The system’s response time improved from 500ms to 200ms after optimization.”

See also

In the discussion section of this recipe, I have described several issues that make text chunks difficult to comprehend.

Instead of addressing each challenge separately, we can use LLMs to handle most of these issues by scanning the entire document and forming text chunks that are complete, meaningful, and clear. Agentic chunking (see Recipe 2.8) does exactly that by identifying similar information across documents and building propositions that contain all relevant details. These propositions serve as high-quality text chunks. The only downside is that it can become quite costly to scan, analyze, and process every single document using LLMs. You need to decide if the increased quality of the text chunks is worth the additional cost.

2.3 Improving Search Accuracy by Embedding Hypothetical Questions

Problem

Improve search accuracy by bringing relevant documents closer to user questions in the embedding space using generated hypothetical questions.

Solution

We have previously used LLMs to pre-process raw text for our RAG system. For example, in Recipe 1.8, we used a multimodal model to create text summaries from images. In Recipe 1.9, we used an LLM to extract key insights from tables.

In this recipe, we use a similar approach. Instead of summarizing content, we show the LLM a final text chunk and ask it to generate potential user questions that can be answered using that chunk. Figure 2-9 illustrates this process step by step.

Load data: Load the data and split it into chunks. Recursive chunking is a good starting point; see Recipe 2.5.

Generate hypothetical questions: Iterate over the chunks and generate hypothetical questions for each one. In Example 2-7, we create five questions per chunk.

Generate embeddings: Generate embeddings for each hypothetical question (see Recipe 3.1).

Push to database: Store the hypothetical questions and their embeddings in the database. Each question links to the original text chunk in the metadata. This is crucial because the original text chunk, not the hypothetical questions, will be provided to the LLM later. The hypothetical questions are used solely to enhance the retrieval step.

Figure 2-9. How to generate hypothetical questions using LLMs

The purpose of this process is to enhance the retrieval step. When a user asks a question, follow these steps:

Generate embeddings: Generate embeddings for the user’s question to compare them with the hypothetical questions.

Retrieval: Retrieve a list of hypothetical questions that are semantically similar to the user’s question. Use the original text chunks associated with these hypothetical questions to build an appropriate prompt.

Prerequisites

To generate effective hypothetical questions, we need a capable LLM. If budget allows, using a powerful LLM is recommended to ensure high-quality questions that enhance retrieval accuracy. In Example 2-7, we will use OpenAI’s GPT-4o.

pip install openai pydantic requests

Steps

In Example 2-7, we load a sample chat history between two students discussing the usefulness of AI in manufacturing. As shown in Figure 2-9, we go through the text chunks and ask the LLM to generate hypothetical questions for each chunk. To help the LLM generate relevant questions, we can provide additional context, such as information about the document, its type, or its domain.

To ensure we receive a list of hypothetical questions formatted as a Python list of strings, we use the structured output functionality. In Example 2-7, we define the expected structure using the Python class HypotheticalQuestions.

Example 2-7. How to generate hypothetical questions using an LLM

    import os
    from openai import OpenAI
    from pydantic import BaseModel
    import textwrap

    file_path = "../datasets/text_files/AI_in_factories_chat.txt"

    with open(file_path, "r") as file:
        text = file.read()

    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

    prompt = textwrap.dedent(
        f"""
        Below you can find a chat history between two students.

        Please generate 5 hypothetical questions that could be
        answered using the information from the discussion.
        The questions should focus on key details, definitions, and
        information present in the text.

        Chat History:
        {text}
        """
    )

    class HypotheticalQuestions(BaseModel):
        questions: list[str]

    # generate hypothetical questions using the GPT-4 model
    completion = client.beta.chat.completions.parse(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="gpt-4o",
        response_format=HypotheticalQuestions,
    )

    hypothetical_questions = completion.choices[0].message.parsed.questions

As defined, we receive a list of five hypothetical questions in return (see Figure 2-10).

Figure 2-10. Screenshot of generating hypothetical questions using LLMs

Now, we generate embeddings for each hypothetical question and store them in our vector database. It’s important to also save the original text chunk and useful metadata, such as the link to the original document and the page number where the text chunk is found.

Discussion

To improve the accuracy of our RAG system, enhancing the retrieval step is crucial.

In a basic RAG setup, we directly compare user questions with text chunks. The challenge is that these text chunks can vary widely, including code snippets, text paragraphs, or tables. Because of the inherent differences between a question and, for example, a table or chat history, the vector embeddings may not align perfectly, even if the underlying information matches well. This issue is often referred to as a semantic alignment problem (see the left side of Figure 2-11).

Figure 2-11. Hypothetical questions are closer to the location of potential user questions in the embedding space

By creating hypothetical questions, we bridge the gap between user queries and document content. These questions highlight the specific information available in the documents, making it easier to retrieve the relevant ones. During the generation step in our RAG pipeline, we use the original text chunks linked to these hypothetical questions.

2.4 Splitting Documents Using Character Splitting

Problem

You need to break down lengthy documents into smaller sections to create embeddings and analyze similarities.

Solution

The Character Text Splitter is the most straightforward method for dividing text. It uses a fixed-size window to create chunks of equal length (see Figure 2-12).

While this method is quite basic and often not the best choice, it can be useful in specific scenarios. For instance, when dealing with raw log files or real-time data streams that lack clear sentence or paragraph breaks, character splitting ensures consistent chunk sizes. This makes processing simpler and helps keep the text within the token limits of foundation models. For most other cases, consider using recursive chunking (see Recipe 2.5) or document-aware chunking (see Recipe 2.6).

Figure 2-12. Splitting text using a fixed character or token size

Prerequisites

We are using LangChain’s CharacterTextSplitter. I am using LangChain here just to explain the parameters it uses, but you don’t actually need a library for this—just count the characters and split the text. There’s no need to make it more complicated than it is.

pip install langchain langchain-text-splitters

Steps

This snippet uses the text splitter module from LangChain. The Example 2-8 counts characters, sets a chunk size of 100 with zero overlap, and uses a space (' ') as the separator. This ensures that words aren’t split in the middle while keeping chunks close to the 100-character limit.

The goal is to create text chunks that are long enough to give the LLM sufficient context but still small enough to generate specific embeddings that capture individual pieces of information. Here, I’ve used 100 characters to demonstrate the concept, but in real RAG use cases, chunk sizes between 1000-2000 tokens or around 4000-8000 characters are typically chosen (since one token is roughly 4 characters in English). It’s important to stay within the token limits of the embedding models. For example, the text-embedding-3-small model from OpenAI has a context length of 8191 tokens, so your chunks can be up to 8191 tokens.

Note

Keep in mind that larger text chunks result in larger prompts sent to the LLM, which slow down response times and increase the running costs of your RAG apps. Additionally, it can negatively affect the performance of the retrieval step, since it becomes harder for the embedding model to clearly capture the meaning behind the text chunk and map it to the vector space.

Example 2-8. How to split text using a fixed text chunk size

    from langchain.text_splitter import CharacterTextSplitter

    file_path = "../datasets/text_files/blog_post_transformers.txt"

    # Load example document
    with open(file_path, "r") as file:
        text = file.read()

    text_splitter = CharacterTextSplitter(
        chunk_size=100,
        chunk_overlap=0,
        separator="",
        length_function=len,
    )

    text_chunks = text_splitter.create_documents([text])

Discussion

The character splitter is a simple way to split the text into chunks of a specific size. However, this method also has many constraints. It doesn’t take the structure of the document into account at all. Instead, it splits paragraphs and sentences somewhere in the middle. While doing that, we are ripping not only the sentences and paragraphs apart but also the text’s whole meaning. Every effort the author invested in properly structuring the document is lost. There are better ways; you can find some in “See also”.

See also

Visualizing the splitting output helps one understand the different chunking methods. A simple web app for the LangChain splitters does just that: chunkviz.up.railway.app.

You will likely find that the character splitter is not the best choice for your needs. While it may seem simple or efficient, its limitations often outweigh its benefits. Better alternatives include:

For classic text documents with headings, chapters, and paragraphs, use recursive chunking, which respects the document’s structure (see Recipe 2.5).

For documents created with markup languages like Markdown, HTML, or LaTeX, use document-aware chunking, which understands the syntax of these languages (see Recipe 2.6).

For documents without a clear structure, such as audio transcriptions or chat histories, use semantic chunking, which analyzes the content to decide where to split the text (see Recipe 2.7).

For complex documents that frequently reference other parts, use agentic chunking. This method generates standalone statements called propositions and merges similar information into one text chunk. It is very powerful but also costly (see Recipe 2.8).

2.5 Splitting Documents Using Recursive Text Splitters

Problem

You need to load and split traditional text documents that have a clear structure with headings, chapters, and paragraphs.

Solution

The recursive text splitter identifies characters that indicate the document’s structure, such as new lines, spaces, and punctuation marks. By default, it uses two new lines (“\n\n”), one new line (“\n”), spaces (” “), and punctuation marks (“.”). Figure 2-13 illustrates how a simple text snippet is divided into chunks using recursive chunking.

Figure 2-13. Splitting text using recursive chunking

Prerequisites

We are using the RecursiveCharacterTextSplitter from LangChain.

pip install langchain langchain-text-splitters

Steps

Example 2-9 processes text from a sample news page covering politics, sports, and healthcare. To ensure each chunk is of a similar size, we set a minimum chunk size. The splitter counts characters or tokens, and once it reaches the defined size, it looks for the next separator to split the text. In this example, we use a small chunk size of 200 characters. For most RAG applications, a chunk size of 1000-2000 tokens (approximately 4000-8000 characters) is more suitable. The maximum chunk size depends on the token limit of the text embedding model and LLM you are using.

Example 2-9. Recursive character text splitter

    from langchain_text_splitters import RecursiveCharacterTextSplitter
    import PyPDF2

    file_path = "../datasets/pdf_files/daily_insights.pdf"

    with open(file_path, "rb") as file:
        reader = PyPDF2.PdfReader(file)

        text = ""
        for page in reader.pages:
            text += page.extract_text()

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=200,
        chunk_overlap=0,
        length_function=len,
        is_separator_regex=False,
    )

    chunks = text_splitter.split_text(text)

The text snippet generates text chunks that are roughly 200 characters each. However, the actual length can vary. If the text contains many long sentences and the recursive chunker splits after each sentence, some chunks might exceed 200 characters. In our example text, this results in 8 chunks. All chunks, except for the last one, have a length of approximately and at most 200 characters (see Figure 2-14).

Figure 2-14. Screenshot of recursive text splitting

Note

The default list of separators works well for most texts and languages. However, for languages like Chinese, Japanese, and Thai, which lack word boundaries, these separators might split words incorrectly. To avoid this, you can customize the list of separators. Best practices can be found in the LangChain documentation.

Discussion

Ideally, each text chunk should cover a single idea. When the author finishes one thought and starts another, we should end one chunk and begin a new one. Fortunately, humans often structure texts to reflect their meaning, grouping related information into paragraphs and chapters with clear headings.

Figure 2-15 illustrates the process of loading and splitting a sample news page, where each section covers the latest news on a specific topic. Randomly creating text chunks can mix information from different sections. This makes it harder for embedding models to capture the semantic meaning accurately. Even if the retrieval step works well, we might load irrelevant information into our prompt, increasing costs, slowing response time, and making it harder for the LLM to generate a good answer. For example, randomly splitting the news page from Example 2-9 could mix information from politics and sports, adding noise to our prompt.

Figure 2-15. Each text chunk should represent one particular piece of information as clearly as possible.

See also

For most documents, the recursive chunker works well. However, for other types of documents, consider these methods:

For documents created with markup languages like Markdown, HTML, or LaTeX, use document-aware chunking, which understands their syntax (see Recipe 2.6).

For unstructured documents like audio transcriptions or chat histories, use semantic chunking, which analyzes the content to determine where to split the text (see Recipe 2.7).

For complex documents, agentic chunking can combine related information from across the document to create more meaningful text chunks (see Recipe 2.8).

2.6 Document Aware Splitting

Problem

You need to load documents with specific syntax and split them into chunks based on their structural elements.

Solution

In Markdown files, headings start with a hashtag (#), and bullet list items begin with an asterisk (*). To make text bold, use two asterisks (**).

Document-aware splitters understand this syntax and use it to divide the text. While you could create your own regex rules, it’s simpler to use existing libraries that already handle these syntax rules. LangChain provides specific splitters for different file types, such as Markdown, Python, HTML, and LaTeX. Figure 2-16 shows examples of this syntax.

Figure 2-16. Sample code example for HTML, Python, Markdown, and LaTeX

Prerequisites

We are using the text splitters from LangChain.

pip install langchain langchain-text-splitters

Steps

Example 2-10 determines the file type based on its extension and selects the appropriate loader. The code snippet handles Markdown, LaTeX, and Python files, using the corresponding LangChain splitters to divide the text.

Example 2-10. Document Aware Splitting

    import os

    file_path = "../datasets/markdown_files/random_md_code.md"
    file_extension = os.path.splitext(file_path)[1]

    with open(file_path, "r") as file:
        file_text = file.read()

    from langchain_text_splitters import (
        PythonCodeTextSplitter,
        LatexTextSplitter,
        MarkdownHeaderTextSplitter,
    )

    if file_extension == ".py":
        splitter = PythonCodeTextSplitter(chunk_size=500, chunk_overlap=50)
    elif file_extension == ".tex":
        splitter = LatexTextSplitter(chunk_size=500, chunk_overlap=50)
    elif file_extension == ".md":
        headers_to_split_on = [
            ("#", "Header 1"),
            ("##", "Header 2"),
            ("###", "Header 3"),
        ]

        splitter = MarkdownHeaderTextSplitter(headers_to_split_on)

    chunks = splitter.split_text(file_text)

The chunking result is shown in Figure 2-17. The text is split into chunks based on the document’s structure. Since we used a Markdown file for this example, the text is split at header 1, header 2, and header 3.

Figure 2-17. Screenshot of document-aware splitting for a Markdown file

Discussion

Markup languages use specific syntax to define the appearance of the final document.

To effectively chunk these documents, we need to understand their syntax rules. Fortunately, LangChain provides various splitters tailored for different file types, making this process easier.

See also

For a brief overview of all these splitters, refer to the LangChain documentation.

Document-aware splitting works only when dealing with document types where the syntax of the raw documents can be used for splitting.

However, sometimes we encounter documents that lack a clear structure. A classic example is when transcribing audio files. Using standard speech-to-text tools, we often end up with plain continuous text. To chunk this text meaningfully, we need to understand the actual meaning of the text and its individual sentences. In this chapter, you will find two approaches to address this:

Semantic chunking (see Recipe 2.7): This method analyzes the semantics of smaller chunks using embedding models and then decides whether and how many smaller text chunks should be merged into larger ones. The idea is to create larger text chunks that are semantically cohesive.

Agentic chunking (see Recipe 2.8): This approach aims for the same goal but uses language models to extract the key insights from text segments and then merges segments with similar information.

2.7 Splitting the Text Using Semantic Aware Chunkers

Problem

You want to chunk documents without any structure elements like headings and paragraphs.

Solution

The goal of the splitting process is to create text chunks that each cover a single idea. When writing a blog post or book, you naturally structure your text into paragraphs, making it easier to read. This structure is ideal for recursive chunking. But what if you’re working with unstructured documents?

Semantic chunking aims to keep related text together, regardless of the document’s structure. It analyzes the content itself to determine which sentences belong together. Using embedding models, we can measure the similarity between consecutive sentences and decide if they should be in the same chunk.

Existing libraries, like the Semantic Chunker, are still experimental. Alternatively, you can create your own semantic chunking pipeline; the process is straightforward.

Figure 2-18 shows the steps involved:

Break text: Break the text into small pieces, like individual sentences.

Generate embeddings: Generate embeddings for each small text piece using embedding models (see Recipe 3.1).

Measure similarity: Measure the semantic similarity between consecutive text pieces by calculating the distance between their embeddings (see Recipe 3.3).

Merge text chunks: Compare the similarity of consecutive text pieces. If they are similar, merge them into a larger text chunk. Continue adding sentences to this chunk until the similarity between two consecutive pieces drops below a certain threshold. This threshold can be set using methods like percentile, standard deviation, or gradient.

Figure 2-18. Define the breakpoint using semantic similarity

For the last step, we need to decide what counts as “similar enough.” In this recipe, I’m using the percentile method. We collect all the calculated distances into a list and then determine where to set the breakpoints. Here, we split the text and create a new chunk when the distance exceeds the 90th percentile.

Tip

The distance in the vector space measures how similar two text chunks are. A smaller distance means higher similarity, while a larger distance indicates lower similarity. See Recipe 3.3 for details on calculating these distances.

Prerequisites

We need an embedding model. In Example 2-11, we will use OpenAI’s Text Embedding model (Ada-003).

To perform the steps mentioned above, we will use LangChain’s SemanticChunker from the langchain_experimental package.

pip install langchain langchain-experimental langchain-openai openai

Steps

For the Example 2-11, we use a text sample that lacks structure; it has no headings, paragraphs, or bullet points. The threshold for splitting the text is set to the 90th percentile.

For testing, we will load a sample news page. This is the same text we used in Recipe 2.5 about recursive chunking, but this time all headings and paragraphs have been removed. As humans, we would need to read the text carefully to find meaningful breakpoints. The semantic chunker does this automatically.

Example 2-11. How to use LangChain’s semantic chunker

    from langchain_experimental.text_splitter import SemanticChunker
    from langchain_openai.embeddings import OpenAIEmbeddings

    # Define the Semantic Text Splitter
    text_splitter = SemanticChunker(
        OpenAIEmbeddings(),
        breakpoint_threshold_type="percentile",
        breakpoint_threshold_amount=90,
    )

    # Split textchunks
    chunks = text_splitter.split_text(text)
    return

When you print the chunks, you will see that it correctly chose the breakpoints at the end of the sports and politics sections (see Figure 2-19).

Figure 2-19. Screenshot of the split text chunks using the semantic chunker

We can adjust the breakpoints in various ways. In Example 2-11, we used the percentile method with a breakpoint_threshold_type and set the breakpoint_threshold_amount to 90.

The provided text is a simple example. You will need to experiment to find what works best for your documents.

Discussion

Ideally, each text chunk should capture one clear idea, with the next chunk covering a different concept. As mentioned in Recipe 2.5, recursive chunking usually works well for most cases since it uses the document’s natural structure, like paragraphs and headings.

However, recursive chunking can be inflexible. It uses separators but does not adjust the chunk size. If we set the chunk size to 500 tokens, it uses the separators to break the chunk but still tries to create chunks with 500 tokens, regardless of whether it makes sense semantically. And what if a proper document structure is missing?

Semantic chunking can also work for texts without any structure. It compares the semantic similarity between two consecutive text chunks. If the similarity is very high, it indicates that the two text chunks belong together.

Although this approach can be costly, it can improve the quality of the text chunks. High-quality text chunks make the retrieval step faster, more cost-efficient, and more reliable.

While the concept is useful, it may be too complex for simple RAG chatbots. However, having a variety of tools, like the semantic chunker, can be beneficial for different tasks. Machine learning is about creativity, and having an extensive toolkit helps.

See also

For a detailed tutorial on chunking, check out Greg Kamradt’s guide on GitHub, which covers 5 levels of chunking, including semantic chunking.

If you prefer an alternative to LangChain, you can explore the LlamaIndex semantic chunker.

For another approach, consider agentic chunking, which uses LLMs to analyze content. Learn more in Recipe 2.8.

2.8 Splitting Text Using Agentic Chunkers

Problem

You are dealing with long, complex documents and want to bundle related information into meaningful text chunks.

Solution

There isn’t a mature library for this yet, but the concept is straightforward. I saw the concept for the first time in Greg Kamradt’s blog post “5 Levels of Text Splitting”. I highly recommend reading it.

The idea is to create standalone statements called propositions. These propositions are clear and understandable without needing additional context. With the right prompt, the LLM can break down text chunks into clear statements and even add implied ones.

Here’s an example:

Original text chunk:

Sarah bought a new book. She enjoys reading fantasy novels.

Created propositions:

Sarah bought a new book.

Sarah enjoys reading fantasy novels.

Sarah reads fantasy novels in her free time. (implied)

The main idea of agentic chunking is to use these clear statements as text chunks, which can improve the quality and performance of our retrieval step. We can further enhance this by reviewing and merging similar statements.

To do this, we iterate over the propositions. We start with the first proposition to create the initial chunk. For each following proposition, we decide whether to add it to the current chunk or start a new one. This process continues for all propositions. Figure 2-20 illustrates this step-by-step process.

Figure 2-20. How agentic chunkers work

Prerequisites

We need an LLM to generate and compare the propositions. In this example, we are using OpenAI’s GPT-4o.

Since we need to process all text chunks multiple times, you might also consider using smaller, more cost-effective models to check if they perform adequately.

pip install openai pydantic requests

Steps

Let’s start with the main functions needed to build your own agentic chunking pipeline.

The key step is creating the propositions. The prompt template we use is from the paper Dense X Retrieval: What Retrieval Granularity should we use. The prompt is shown in Example 2-12; I have shortened the example here to save space in the book. You can find the full prompt in the LangChain Hub.

Example 2-12. How the prompt could look like that breaks down the text into propositions

Decompose the "Content" into clear and simple propositions, ensuring they are interpretable out of context.

1. Split compound sentence into simple sentences. Maintain the original phrasing from the input whenever possible.
2. For any named entity that is accompanied by additional descriptive information, separate this information into its own distinct proposition.
3. Decontextualize the proposition by adding necessary modifier to nouns or entire sentences and replacing pronouns (e.g., "it", "he", "she", "they", "this", "that") with the full name of the entities they refer to.
4. Present the results as a list of strings, formatted in JSON.

Example:

Input: Title: ¯Eostre. Section: Theories and interpretations, Connection to Easter Hares. Content:
The earliest evidence for the Easter Hare (Osterhase) [...] America where it evolved into the
Easter Bunny.
Output: [ "The earliest evidence for the Easter Hare was recorded in south-west Germany in 1678 by Georg Franck von Franckenau.", "Georg Franck von Franckenau ..", "..."]

The prompt instructs the LLM to convert the content into clear and simple propositions. Let’s try it out. In Example 2-13, we load the prompt directly from the LangChain Hub.

Example 2-13. How to create propositions

    from langchain import hub
    from langchain_openai import ChatOpenAI  # Import ChatOpenAI
    from pydantic import BaseModel
    from typing import List

    # pull the prompt template from the langchain hub
    obj = hub.pull("wfh/proposal-indexing")

    # define the llm you want to use
    llm = ChatOpenAI(model="gpt-4o")

    # A Pydantic model to extract sentences from the passage
    class Sentences(BaseModel):
        sentences: List[str]

    extraction_llm = llm.with_structured_output(Sentences)

    # Create the sentence extraction chain
    extraction_chain = obj | extraction_llm

    # Test it out
    propositions = extraction_chain.invoke(
        """
        On July 20, 1969, astronaut Neil Armstrong walked on the moon .
        He was leading the NASA's Apollo 11 mission.
        Armstrong famously said, "That's one small step for man, one
        giant leap for mankind" as he stepped onto the lunar surface.
        """
    )

    print(propositions)

Figure 2-21 shows how the generated propositions could look like. The LLM has broken down the text into clear propositions.

Figure 2-21. How the generated propositions could look like

The agentic chunker builds a reusable function around the snippet in Example 2-13. You can read the entire document into the model to generate propositions, but this approach lacks control over whether the LLM extracts all relevant statements. To improve performance and maintain control, first chunk the data into relatively large sections.

There is no one-size-fits-all solution; you need to test what works best for you. For example, you can use the recursive splitter from Recipe 2.5 and set the separator to two new lines (\n\n). This way, each chunk ends at a paragraph or chapter break.

Send these paragraphs to the LLM to create propositions. You’ll get a list of statements for each paragraph. To avoid duplicates and merge similar statements, review all propositions and decide whether to add them to an existing chunk or start a new one.

Discussion

In books, contracts, or any other lengthy documents, information is often scattered throughout, making it tricky to chunk into meaningful segments. Different chapters and clauses are interconnected. As humans, we can piece things together by flipping between chapters or using Google and Wikipedia to clear up any confusion.

In RAG systems, our model only sees content from a single vector search: a collection of text chunks that might help answer a specific question. When we just chunk the data, we lose the context of how different chapters, paragraphs, and snippets are connected.

Agentic chunkers tackle this issue by using LLMs to create text chunks that merge similar information.

See also

There is no single right or wrong approach—multiple methods exist.

The steps described above are a good starting point. In the first step, by generating propositions from the raw text, you can already significantly improve the quality of your text chunks. You can stop the chunking process at this stage and end up with a long list of statements. Whether you need the next step, which merges similar information, depends on the amount of your data and how often you encounter redundant information in your text chunks.

You can use the standard implementations available in LangChain or LlamaIndex, or you can develop your own agentic chunking process tailored to your specific use case. Either way, the documentation and implementations in orchestration frameworks like LangChain and LlamaIndex are excellent sources of inspiration.

There are some great resources available to learn more about agentic chunking. I highly recommend checking out the following:

Greg Kamradt’s GitHub repository on chunking strategies

Thuwarakesh Murallie’s blog post on agentic chunking

Chapter 3. Embeddings

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 5th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the author at dmnkplzr@googlemail.com.

What You’ll Learn

In this chapter, you’ll learn how to:

Create embeddings for text and images using state-of-the-art models like OpenAI’s text-embedding-3-small and CLIP (see Recipe 3.1, Recipe 3.5)

Understand the difference between traditional encoding methods (label encoding, one-hot encoding) and modern embedding techniques (see Figure 3-1, Figure 3-2)

Perform semantic similarity search by calculating distances between embedding vectors using cosine similarity and other metrics (see Recipe 3.3)

Visualize high-dimensional embeddings in 2D/3D space using dimensionality reduction techniques like PCA (see Recipe 3.2)

Choose the right embedding model for your specific use case based on cost, accuracy, and performance requirements (see Recipe 3.4)

Build text classification models using embeddings as feature vectors (see Recipe 3.6)

Implement hybrid search combining keyword-based (BM25) and semantic search for improved retrieval (see Recipe 3.7)

Work with multimodal embeddings to incorporate images into your RAG systems (see Recipe 3.5)

Apply best practices for embedding model selection and evaluation in production RAG applications (see Recipe 3.4)

Embeddings are a fundamental part of many classic ML use cases, since most algorithms only work with numerical data such as integers or vectors. This means any text in our dataset needs to be converted into numbers so the algorithm can process the meaning behind it. The goal of text embeddings is to capture as much information as possible from text labels, giving other ML algorithms a chance to identify patterns in text data.

Tip

There are also multimodal embeddings that can handle images or mixed content (images and text). However, most RAG systems focus on text embeddings. When we need to work with images or videos, we usually convert them to text descriptions first, then use text embedding models for indexing.

In RAG systems, we are usually interested in the semantic similarity between the user’s query and the text chunks in our vector store, because we assume that semantically similar content likely contains information that is helpful for answering the user’s question. Therefore, we need an embedding model that can robustly capture this semantic meaning of text and map this meaning to a vector space. This allows us to perform calculations based on semantic meaning.

Embedding models vs. LLMs: What’s the difference?

Embedding models transform words or phrases into numerical vectors that capture their meanings and relationships, enabling algorithms to assess semantic similarity by measuring the distance between them. In contrast, Large Language Models (LLMs) generate new text based on given prompts or questions. This is the primary distinction between these two types of models.

You may be familiar with basic methods like label encoding and one-hot encoding that convert text to numbers. These assign a number or vector to each label without considering meaning. While these work for simple cases, they’re not suitable for RAG systems where we need to understand what the text actually means.

Nevertheless, I believe it is important to understand the idea behind classic encoding methods to see why we need more powerful embedding models and what they actually do.

Note

Don’t be confused by the term “encoding” here. Encoding simply means converting data from one format to another. Embedding is a specific type of encoding that creates meaningful representations of data by capturing relationships and semantic meaning in the resulting vector space.

In RAG systems, we use embeddings to convert text into vectors that preserve semantic meaning. We can then calculate distances between these vectors to determine how similar two text chunks are.

Figure 3-1 shows a basic example of how label and one-hot encoding work. Imagine a dataset that contains information about education level, age, and yearly salary for a group of surveyed individuals. If we want to use this dataset to train a regression model that predicts the annual salary for each new person, we first need to preprocess the raw data so the algorithm can work with it. In classic ML use cases, this process is called Feature Engineering which focuses on preparing raw data so that the information in each column is as clear as possible and in a form that works best for the specific type of algorithm used.

Feature engineering can be challenging and requires creativity. For datasets with limited data, good feature engineering is often the most important factor for improving model performance. This chapter focuses on one specific aspect: converting text data into numbers while preserving important information so algorithms can process it effectively.

In Figure 3-1, we have an education level column with text values like “Bachelor”, “Master”, and “PhD”. We need to convert these into numbers so our algorithm can use them.

Two common methods are:

Label encoding: Assign a number to each label, like “Bachelor” = 1, “Master” = 2, and “PhD” = 3.

One-hot encoding: Convert labels into vectors, like “Bachelor” = [0, 1, 0], “Master” = [1, 0, 0], and “PhD” = [0, 0, 1].

Figure 3-1. How label and one-hot encoding work

The problem with randomly assigning numbers or vectors is that we basically lose all information from the text. For example, a bachelor’s degree means someone has spent at least three years at a university. If they continue their studies, they might obtain either a master’s degree or a PhD. Randomly assigning numbers or vectors ignores this underlying meaning completely. But exactly this underlying meaning and information is crucial for RAG systems, so that we are able to match the information with the user question.

Figure 3-2 shows how we can solve this problem by creating vectors where each dimension has a clear meaning. Instead of random numbers, we use characteristics that actually matter.

Let’s look at an example with royal titles: Queen, King, Prince, and Princess. We can represent these using just 3 dimensions:

Age: Kings are typically older than princes

Gender: Queens and princesses are female (1), kings and princes are male (0)

Royal Status: Kings and queens have higher status than princes and princesses

So “King” might become [0.8, 0, 1] (older, male, high status) while “Princess” becomes [0.3, 1, 0.5] (younger, female, medium status).

Now we can even perform calculations with these vectors. For example: King - Man + Woman = Queen

This works because the vectors capture real relationships between the words.

Figure 3-2. Calculating with embeddings

This example demonstrates how embeddings work in general, but for RAG systems, the example with dimensions like gender, age, and royalty level is too simplistic. We need a model that can be applied to any kind of text. The only requirement is that, in the end, we want a vector that somehow captures the semantic meaning of the text.

So, how do we create meaningful embeddings for any kind of text? We could theoretically have linguists manually design how to convert words, sentences, and text snippets into vectors. While this might work in theory, it would be extremely time-consuming and impractical. Plus, humans would be too subjective in their decisions.

A much better approach is how embedding models are trained today. We collect massive amounts of text data and use transformer-based algorithms to automatically learn the best dimensions for the embedding space. This allows the embedding model to convert any text snippet into a high-dimensional vector, where the position of each vector reflects its semantic meaning. Vectors that are close together in this space represent texts that are semantically similar.

Note

The most powerful embedding models today are based on the transformer architecture. Transformer models are a type of deep learning architecture that revolutionized natural language processing (NLP) when introduced in the 2017 paper “Attention Is All You Need”. These models form the foundation for many state-of-the-art embedding models. For a deeper understanding of how these models work, I recommend Andrew Karpathy’s YouTube video on building GPT.

Embedding models today generate embedding vectors with hundreds or thousands of dimensions. Recipe 3.1 shows how we can use different embedding models to generate these high-dimensional vectors. In the discussion section of this recipe, you will find a brief explanation of how these models are trained (see “Discussion”).

In RAG systems, we use these embeddings to find content that matches what the user is asking about. Here’s how it works:

Store embeddings: We save the embedding vectors along with the original text chunks in a vector database.

Process user questions: When a user asks a question, we convert their question into an embedding vector using the same embedding model.

Find similar content: We compare the user’s question embedding with all the stored embeddings by calculating the distance between them. The smaller the distance, the more similar the content.

Retrieve relevant text: We select the text chunks that have the smallest distances to the user’s question - these are the most relevant pieces of content to help answer their question.

This process allows us to find semantically related content even when the exact words don’t match (see Recipe 3.3).

We use the text chunks from the most similar embedding vectors and give them to the LLM along with the user’s question. We hope these text chunks contain useful information to answer the question. This is the basic idea behind RAG systems. This way, we can give the LLM information it doesn’t know about, like private company data or recent events.

Think of it like giving someone a reference book before asking them a question. The LLM gets both the question and the relevant background information it needs to answer accurately.

Microsoft Copilot works this way to answer questions about recent events that weren’t in the model’s training data. For example, if you ask ChatGPT, Gemini, or Claude about yesterday’s iOS update announcement at Apple’s WWDC conference, the LLM wouldn’t know about it. At best, it might give you old information or make something up. Every AI model only knows what was in its training data, which always has a cutoff date. Copilot and Google Search solve this problem by adding fresh information from Bing and Google search results to their answers.

Figure 3-3 illustrates how we provide the model with additional data it doesn’t already know. We perform a similarity search to find relevant information first, feed the found data to a single prompt, and send it to the LLM.

Figure 3-3. How to fill the LLM’s knowledge gaps using the RAG concept

Beyond calculating similarity between text chunks, we can use embeddings for many other tasks. Building classification models for text and images used to be complex and time-consuming. Now, embedding models make these tasks much simpler and faster.

In Recipe 3.5, we use CLIP, an embedding model that creates embeddings for both images and text. This enables us to build an image classifier in a few seconds. Since CLIP maps both text and images into the same vector space, we can compare the embedding vector of an image with the description “This is a dog.” By calculating the distance between the vectors, we get an idea of how likely it is that the image actually shows a dog.

In Recipe 3.6, we use text embeddings to convert text into meaningful vectors that capture semantic meaning. This allows classification models to identify patterns that simpler methods like label encoding or one-hot encoding would miss.

3.1 Mapping the Linguistic Meaning of Text Chunks to a Numerical Representation by Generating Embeddings

Problem

You want to turn words and sentences into numbers that represent their meaning.

Solution

Use embedding models trained on large text datasets. These models learn patterns to map text snippets into a high-dimensional vector space. The position of data points in this space shows how text chunks are related. Similar embeddings have similar meanings. For example, the embedding for “cat videos online” will be closer to “endless procrastination” than to “productive Monday morning” (see Figure 3-4).

Figure 3-4. An example of semantic similarity

Prerequisites

In Table 3-1, you can see some of the latest models from OpenAI, Google and Voyage AI, as well as a smaller one from the open-source community. In this recipe, I will show you how to use the text-embedding-3-small model from OpenAI using the official OpenAI library and the open-source model all-MiniLM-L6-v2 via the Hugging Face API.

Table 3-1. Embedding models

Model
Company
Platform
Max. Tokens
Embedding Dimensions

text-embedding-3-small

OpenAI

OpenAI API

8191

1536

voyage-large-2

Voyage API

Voyage AI Cloud

16000

1536

text-embedding-004

Google

Gemini API / Google Cloud

2048

768

all-MiniLM-L6-v2

Open Source

Hugging Face API

512

384

What is a token?

A token is roughly equivalent to 4 characters in English. For example, a model with an 8000-token context size can process about 32,000 characters of text. To see how different models break text into tokens, you can use OpenAI’s Tokenizer Web App, which lets you input text and shows how it is divided into tokens.

Most models handle between 4k and 128k tokens or more. For typical RAG applications, it is good practice to keep text chunks much smaller. This approach makes the application faster, more cost-effective, and often more accurate by focusing on essential information and filtering out noise. For more details on chunking data, refer to Chapter 2.

To learn more about tokens, visit the OpenAI Help Page.

In the rest of the recipie we will use the text-embedding-3-small model from OpenAI and the all-MiniLM-L6-v2 model via the Hugging Face API using the sentence-transformers library.

pip install openai sentence-transformers

Solution

In the first step, we use OpenAI’s SDK and one of OpenAI’s embedding models to generate embedding vectors for text. Specifically, we use the text-embedding-3-small model. To start, install the OpenAI Python library and set up your API key.

Example 3-1 shows how to use the library now to generate embeddings for two sample sentences. To convert a sentence into a vector, we only need two commands. With client = OpenAI(), we create a client that has the function client.embeddings.create() which we use to create embeddings. For this, we define the string for which we want to generate an embedding vector as input and specify the model we want to use. The result is a list of two vectors, each with 1536 dimensions.

Example 3-1. Text to embeddings using OpenAI’s embedding models

from openai import OpenAI

text_chunks = ["The sky is blue.", "The grass is green."]

client = OpenAI()  # uses the environment variable OPENAI_API_KEY

embeddings_list = []

for text_chunk in text_chunks:
    response = client.embeddings.create(
        input=[text_chunk], model="text-embedding-3-small"
    )
    embedding = response.data[0].embedding
    embeddings_list.append(embedding)

Figure 3-5 shows what the list of embedding vectors looks like.

Figure 3-5. How embedding vectors look in practice

As an alternative to OpenAI, we can also use smaller, open source models. In this example, we use all-MiniLM-L6-v2, which is one of the most popular sentence transformer models. We use the model via the Hugging Face API. To get started, get your token at huggingface.co/settings/tokens. With the Python library sentence-transformers, we can easily load and use the models.

Note

Sentence Transformers is a Python library that builds on Hugging Face Transformers and is specifically designed for generating sentence embeddings.

We get a list with five vectors. Each of the vectors has 384 dimensions.

Example 3-2. Text to embeddings using the Hugging Face API

from sentence_transformers import SentenceTransformer

# Load the model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
text_chunks = ["The sky is blue.", "The grass is green."]

# Generate embeddings
embeddings = model.encode(text_chunks)

Discussion

Translating natural language robustly into vectors opens up new possibilities for a wide range of ML applications. RAG apps and classification use cases are just two application areas, but actually you have many more possibilities every time you have text as input data. Transformer-based models do this remarkably well. But how do they achieve this, and how are they trained?

Embedding models are neural networks trained on large text datasets. Each unique word in the dataset is assigned to a neuron in the input layer. Figure 3-6 shows a simple example with a hidden layer containing only two neurons. During training, the neural network predicts the next word for each input word. For example, when the input neuron for “Google” is activated, the network is trained to predict the word “is”. Once we have fully trained the model, we get an embedding model that maps every word or sentence into a two-dimensional space. The weights of the hidden layer neurons determine the position in this space.

In this example, we use the sentence “Google is a tech company”. Each word in the sentence is an input for the neural network (NN), so our network has five input neurons, one for each word.

Figure 3-6. How vector embedding models work

Embedding models are not only useful in RAG systems but have been applied in Data Science problems for years. Here are two examples:

In “Fruits Search System”, we create a simple search engine for fruits and vegetables.

In “Delivery time prediction model”, we convert text into vectors to train a regression model for predicting delivery times.

Fruits Search System

Figure 3-7 shows how embedding models work with fruits and vegetables. When you scan a fruit at a self-checkout, the system can suggest similar items. The embedding model understands that “pears and apples” are more related than “apples and onions,” so it can recommend fruits that are commonly grouped together.

Figure 3-7. Similarity score between different types of fruit and vegetables

Delivery time prediction model

Embedding models work well for many ML tasks with text data. Most ML algorithms need numbers, not text. Figure 3-8 shows a regression example where we convert supplier names and material names into embeddings. Simple methods like label encoding assign numbers to text, but they lose important information.

For example, suppliers named “DHL Global Logistics” and “UPS International Shipping” are both logistics companies. A good embedding model will recognize this similarity and capture it in the embedding vectors.

Figure 3-8. How to translate text data to embeddings

See also

I recommend Joshua Starmer’s StatQuest YouTube channel for visual and clear explanations of every kind of ML concepts. Joshua has the talent to break down even the most complicated ML algorithms so that everyone, no matter their background, has a chance to understand them. I mention this here because he also produced a helpful tutorial on Word Embedding and Word2Vec.

3.2 Visualizing the Semantic Relationship between Text Chunks by Reducing the Dimensionality of Embedding Vectors

Problem

You want to create a visual map showing how similar different pieces of text are to each other.

Solution

Embedding models convert text into vectors with hundreds or thousands of dimensions to robustly capture their meaning. For example, Word2Vec creates vectors with 285 dimensions, OpenAI’s text-embedding-3-small with 1536 dimensions, and OpenAI’s text-embedding-3-large with 3072 dimensions.

We can only visualize two or three dimensions, so we need to reduce the number of dimensions to be able to plot them. While doing this, we will inevitably lose a lot of information, since a two-dimensional vector cannot impossibly hold as much information as a 1536-dimensional vector, but we can use techniques to minimize information loss.

This is where dimensionality reduction techniques in ML come into play (Figure 3-9). Popular methods include PCA, t-SNE, UMAP, LDA, and MDS. In this recipe, we will focus on using PCA for dimensionality reduction.

Figure 3-9. Reducing the dimensionality of embedding vectors using PCA

Prerequisites

First, we need an embedding model to convert text into vectors. We’ll use OpenAI’s text-embedding-3-small model, which generates vectors with 1536 dimensions.

Next, we’ll use the PCA (Principal Component Analysis) method from the scikit-learn library to reduce the number of dimensions of these vectors. This will allow us to plot the vectors in a 2D plot using the Matplotlib library.

pip install openai scikit-learn matplotlib

Steps

Example 3-3 generates embeddings for six sample text snippets. Then, in Example 3-3, we will use Matplotlib to reduce the dimensions of these embeddings and plot them.

Example 3-3. Visualize embeddings using PCA

    import os
    import pandas as pd
    import matplotlib.pyplot as plt
    from sklearn.decomposition import PCA
    from openai import OpenAI

    # Define text chunks
    text_chunks = [
        "The sky is blue.",
        "The sun is shining.",
        "I love chocolate.",
        "Ice cream is delicious.",
        "Roses are red.",
        "Violets are blue.",
    ]

    # Initialize OpenAI client
    client = OpenAI()  # uses the environment variable OPENAI_API_KEY
    embeddings = []

    # Generate embeddings for text chunks
    for text_chunk in text_chunks:
        response = client.embeddings.create(
            input=[text_chunk], model="text-embedding-3-small"
        )
        embedding = response.data[0].embedding
        embeddings.append(embedding)

    # Convert embeddings to a DataFrame
    embeddings_df = pd.DataFrame(
        embeddings, columns=[f"dim_{i}" for i in range(len(embeddings[0]))]
    )

To visualize the embeddings we just created, we use the PCA function to reduce their dimensions. Since we want to display them in a 2D plot, we set the variable n_components to two. At the end of Example 3-4, I save the plot as an image. PCA transforms the 1536-dimensional vector into a 2-dimensional vector by finding the directions (principal components) that capture the most variance in the data. The first principal component captures the most variance, the second captures the second most, and so on. This allows us to represent most of the important relationships between the embedding vectors in just two dimensions, though some information is inevitably lost in the dimensionality reduction process.

Example 3-4. Visualize embeddings using PCA

    # Perform PCA with 2 components
    pca = PCA(n_components=2)
    df_reduced = pca.fit_transform(embeddings_df)

    # Create a new DataFrame with reduced dimensions
    df_reduced = pd.DataFrame(df_reduced, columns=["PC1", "PC2"])
    df_reduced["text"] = text_chunks  # Add original text chunks for labeling

    # Function to create a scatter plot
    plt.figure(figsize=(8, 6))
    plt.scatter(df_reduced["PC1"], df_reduced["PC2"])

    # Add labels to each point
    for i, label in enumerate(df_reduced["text"]):
        plt.text(df_reduced["PC1"][i], df_reduced["PC2"][i], label, fontsize=9)

    # Add labels and title
    plt.xlabel("Principal Component 1")
    plt.ylabel("Principal Component 2")
    plt.title("PCA Scatter Plot")

    # Save the plot
    plt.savefig("../results/principal_component_plot.svg", format="svg")

The plot in Figure 3-10 shows the results. Sentences about food and drinks are grouped on the right, weather-related sentences are in the lower left, and flower sentences are in the upper left. This shows the model successfully captures the meaning of different text types.

The example here uses very different topics, which makes patterns easy to see. In real RAG applications, text chunks usually come from the same domain and are more similar to each other. This makes finding the most relevant chunks harder, but modern embedding models handle this well.

Figure 3-10. Reducing the dimensionality of embedding vectors using PCA

While the 2D plot doesn’t give exact measurements of similarity in the high-dimensional space, it shows whether the model is capturing the intended meaning. This gives us a quick sense of whether the model is working as intended.

Discussion

Interpretability and transparency are crucial for building trust in your ML application. Visualizing data helps users understand what the model is doing. Even a simple 2D plot can provide a sense of whether your app is functioning correctly. Making these visualizations accessible to users can enhance trust, as users gain insight into the inner workings of your app.

Note

I work in a highly regulated industry with users from traditional engineering backgrounds. In these environments, people need to trust new technology before adopting it. While performance benchmarks and test results matter, showing users how the system actually works often builds trust more effectively. This recipe demonstrates how to make your retriever componentent in your RAG system transparent and understandable.

See also

In a 2D space, it’s easy to see which data points are close together. But this doesn’t give us the complete picture because we lose important information when we compress high-dimensional vectors into just two dimensions. For more precise similarity measurements, we need to calculate distances using the original high-dimensional vectors. Recipe 3.3 explains how to do this.

3.3 Calculating the Distance Between Embeddings

Problem

You want to measure how similar two pieces of text are to each other.

Solution

There are many ways to measure the distance between vectors, such as cosine distance, Euclidean distance, Manhattan distance, Hamming distance, Minkowski distance, and Haversine distance. Some examples are shown in Figure 3-11.

For RAG systems, cosine similarity is the most commonly used. It is effective with transformer-based embedding models because it measures the angle between vectors, focusing on their direction rather than their size.

Figure 3-11. Three of the most used distance metrics in RAG systems

Prerequisites

We use the text-embedding-3-small model to generate embeddings and NumPy to calculate the similarity between them. Alternatively, you can also use the built-in function for calculating cosine similarity in the sentence-transformer library.

pip install openai numpy sentence-transformers

Steps

Let’s start with OpenAI’s model. In Example 3-5, we create embeddings for three historical facts using the OpenAI Python library.

We have three historical statements and one user question: “Tell me something interesting about diseases in history.” Our goal is to find which historical fact best matches the user’s question and could help answer it.

Example 3-5 shows how to create embeddings for both the historical statements and the user’s query using the OpenAI library.

Example 3-5. Create embeddings using OpenAI’s text-embedding-3-small model

    import numpy as np
    from numpy.linalg import norm
    import pandas as pd
    import os
    from openai import OpenAI

    text_chunks = [
        "The Great Fire of London in 1666 destroyed over 13,000 houses.",
        "Julius Caesar was assassinated on the Ides of March (March 15) in 44 BCE.",
        "The Black Death, is estimated to have killed nearly one-third of the \
            population",
    ]

    users_question = "Tell me something interesting about diseases in history"

    embeddings_df = pd.DataFrame(text_chunks, columns=["text_chunk"])

    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    embeddings = []

    def create_embeddings(text_chunk, client):
        embedding = (
            client.embeddings.create(
                input=[text_chunk], model="text-embedding-3-small")
            .data[0]
            .embedding
        )
        return embedding

    # Apply function create_embeddings to the correct column
    embeddings_df["embedding"] = embeddings_df["text_chunk"].apply(
        create_embeddings, client=client
    )

    users_question_embedding = create_embeddings(
        text_chunk=users_question, client=client
    )

Now that we have embeddings for all historical facts and the user’s question, we can find out how similar the user’s question is to each historical fact. We’ll use NumPy’s cosine similarity function to measure this similarity.

Example 3-6. Calculate the cosine similarity between two vectors using NumPy

    # create a list to store the calculated cosine similarity
    cos_sim = []

    def calculate_cosine_similarity(text_chunk_embedding, users_question_embedding):
        A = text_chunk_embedding
        B = users_question_embedding

        # calculate the cosine similarity
        cosine = np.dot(A, B) / (norm(A) * norm(B))
        return cosine

    # Apply function calculate_cosine_distance to the correct column
    embeddings_df["similarity"] = embeddings_df["embedding"].apply(
        calculate_cosine_similarity, users_question_embedding=users_question_embedding
    )

Figure 3-12 shows the results. The third historical fact, “The Black Death killed approximately 75-200 million people,” has the highest similarity score to our question about diseases in history.

Figure 3-12. Cosine similarity calculation screenshot

We can also use the sentence transformer library for convenience, which has a built-in function to calculate cosine similarity (see Example 3-7).

Example 3-7. Calculate the cosine similarity between two vectors using the sentence transformer library

    from sentence_transformers.util import cos_sim

    for document_embedding in document_embeddings:
        # Compute cosine_similarity between documents and query
        scores = cos_sim(document_embedding, query_embedding)

Figure 3-13 shows what the similarity calculations look like when visualized in a plot. The cosine similarity is calculated from the cosine distance (cosine similarity = 1 - cosine distance). This measures how similar two vectors are by looking at the angle between them, not their actual size or length. Since we’re using OpenAI’s text-embedding-3-small model, we’re working with vectors that have 1536 dimensions. Cosine similarity works well with text embedding models because it focuses on the direction of the vectors rather than their magnitude, which matches how these transformer models are trained to work.

To make the results easy to understand, I visualized the vectors again in a 2D plot. To plot the 1536-dimensional vectors, I used the dimensionality reduction techniques from Example 3-4.

Figure 3-13. How to calculate the distance between two text chunks

Discussion

In summary, modern embedding models based on transformer technology make it much easier to build search systems for documents in RAG applications. When combined with the distance calculations we’ve covered in this recipe, these tools provide the foundation for finding relevant information in a RAG system.

See also

I highly recommend checking out these helpful resources:

Josh Starmer’s StatQuest videos are great for learning. Watch his cosine similarity video to understand how it works.

James Briggs’ YouTube channel explains vector similarity search clearly. Check out his introduction to similarity search and his video on other distance methods like Jaccard and Levenshtein distance.

Maarten Grootendorst wrote a clear article about nine different ways to measure distance between data points.

3.4 Choosing the right Embedding Model

Problem

You need to pick the right embedding model for your project.

Solution

When choosing an embedding model, consider three main factors:

Cost: How much does the model cost to use?

Accuracy: How accurate and reliable is the model?

Response Time: How quickly can the model generate embeddings?

Cost and response time are closely linked. Smaller models are cheaper and faster but may be less accurate.

Accuracy is measured using benchmarks like the Massive Text Embedding Benchmark (MTEB). You can find the leaderboard on Hugging Face. Higher accuracy usually comes with more dimensions, which means you need more storage space and computing power to calculate distances between vectors. Choose a model that’s powerful enough for your needs without being overkill. For basic RAG applications, you don’t need the most powerful models available.

Note

You can explore a wide range of models on Hugging Face, where the Massive Text Embedding Benchmark (MTEB) leaderboard provides benchmarks to evaluate their performance. The MTEB leaderboard offers insights into how different models perform on various embedding tasks, including key details like MTEB scores, the number of embedding dimensions, and the maximum token length.

Figure 3-14 shows how cost and accuracy relate to each other. OpenAI’s larger model, text-embedding-3-large, costs six times more than their smaller model, text-embedding-3-small. For most RAG applications, the smaller and cheaper models work well enough. Only choose a more powerful model if getting very accurate search results is critical for your specific use case.

Figure 3-14. Cost and accuracy comparison of smaller and larger embedding models

If your dataset doesn’t change often (like once a year), you can use a more accurate model since cost isn’t a big concern. But if your dataset changes daily and contains lots of entries, cost becomes more important. In these situations, a smaller, cheaper model often works just fine.

Test different models with your actual data to see how they perform. Also, look at your entire RAG system’s performance. Sometimes there are better ways to improve results than simply using more powerful models. For more about RAG evaluation, see [Link to Come].

Here are the key points to consider when choosing an embedding model:

Cost: More powerful models are more expensive. However, for most use cases, the cost and response time of the embedding model are not the main concerns. Instead, the token costs of the LLM and the response time are much higher. The costs of the embedding model become relevant if you frequently need to reload the entire dataset into your vector store.

Dimensions: Higher dimensions generally capture and store more semantic details of the sentence. However, they require more storage space and are more compute-intensive for calculating the distance between vectors. For most use cases, this is not a critical factor.

Evaluation: An evaluation framework is essential. As with any other ML use case, we can only optimize what we can measure. Ideally, you should have a dataset with search terms and the expected best matches. This allows you to calculate an accuracy score by comparing the retrieved contexts with those from the reference dataset. You can learn more about this in the chapter on RAG evaluation.

We can choose between two main types of embedding models: static embeddings and contextual embeddings.

Static embeddings, like Word2Vec, are smaller and can be run on local or edge devices, making them suitable for offline use. However, they are less powerful compared to transformer-based models.

Contextual embeddings, such as BERT or OpenAI’s Text-Embedding-3-small, are more powerful and capture the semantic meaning of text more effectively. However, they are more computationally intensive and relatively expensive.

When performance is crucial, and you need a model that can run offline, smaller models like Word2Vec might still be a good option. Those models are small enough to run on edge devices.

Discussion

Embedding models are reasonably priced compared to LLMs. However, if you need to reload your entire vector store daily and handle massive amounts of data, a more cost-efficient model might be the better choice. On the other hand, if a more powerful embedding model actually leads to better results for your retriever, it can help reduce the overall response time of your RAG app, since higher-quality contexts will allow you to reduce the overall size of the prompt.

In any case, it’s a good idea to keep an eye on the latest embedding models. The field is evolving quickly, with new models being released at least monthly. These new models are becoming more powerful and cost-efficient.

There are now models of all sizes, response times, and costs. Depending on your use case, you can choose the right one:

If you are training classification models with long text chunks, a more powerful model may be worthwhile to capture all the details.

If you’re building a simple RAG chatbot, a smaller, more cost-effective model is often sufficient. Focusing on the quality of the text chunks stored in your vector store will have a greater impact than squeezing out the last few accuracy points from your embeddings.

Everything we’ve discussed here also applies to selecting the LLM for your RAG system. You can find more about this in Chapter 6 on the Generation Step.

See also

A great resource for exploring available models is the MTEB leaderboard.

Once you’ve picked a few models you want to test, check out the model providers’ docs for more details. For example, you can find an Embedding Guide in the official OpenAI library, with an extensive FAQ list. The current pricing can be found on the OpenAI website.

3.5 Generating Embeddings for Images and Text using CLIP

Problem

You want to work with both images and text in your RAG system.

Solution

Contrastive Language-Image Pre-training, or CLIP. CLIP is a transformer model from OpenAI that connects images and text. It is trained to match text snippets with the correct images and vice versa. This allows CLIP to create embeddings for both images and text.

Figure 3-15 illustrates how to use CLIP for retrieval tasks in a RAG system. The process works in four main steps:

Create embeddings using CLIP: First, we process images and convert them into embedding vectors using the CLIP model. These embeddings capture the visual content and meaning of each image.

Create embeddings for user’s question: When a user asks a question like “Show me all the pictures of buildings. Can you guess where they were taken?”, we use CLIP to convert this text query into an embedding vector in the same vector space.

Find relevant content (images): We search through our image and vector database to find images whose embeddings are most similar to the user’s question embedding. This identifies the most relevant images based on semantic similarity.

Formulate answer: Finally, we use a multimodal model that can process both the retrieved images and the user’s question to generate a comprehensive response that answers the query based on the visual content.

This approach allows us to search through image collections using natural language queries and retrieve visually relevant content for our RAG system, making it possible to answer questions about visual information stored in our database.

Figure 3-15. How to use the CLIP model for retrieval tasks

We can also use CLIP to search for images. If you have a database of images and need to answer questions based on those pictures, CLIP can help find the relevant images to answer the user’s question.

Prerequisites

To use CLIP on your local machine, you need to install the Hugging Face Transformers library. The transformers library provides the CLIP model implementation, torch is required for running the neural network computations, and pillow handles image loading and preprocessing.

pip install transformers torch pillow

Steps

In Example 3-8, we create a simple classifier for images of dogs and cats. The idea is to provide an image of a dog and compare it to predefined classes. The phrase “image of a dog” should have similar embeddings to the actual dog image (see Figure 3-16).

Figure 3-16. How to classify images using CLIP

First, we load the model and the processor to create embeddings for the images and text chunks. The processor prepares the raw text and images by tokenizing the text and pre-processing the images, converting them into the format required by the CLIP model.

We calculate two values:

logits_per_text: This contains the similarity scores between each image embedding and each text embedding.

probabilities: This contains the probabilities derived from the similarity scores using the softmax function, ensuring they sum to 1.

Example 3-8 loads the CLIP model and processor to create embeddings for images and text. The processor tokenizes the text and pre-processes the images. We then calculate similarity scores between the image and text embeddings, resulting in a list of probabilities for each image-text pair.

Example 3-8. Create embeddings using CLIP

    import torch
    from PIL import Image
    from transformers import CLIPProcessor, CLIPModel

    # Load the model and processor
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    # Text and image inputs
    descriptions = ["A photo of a cat", "A photo of a dog"]
    images = [
        Image.open("../datasets/images/cat.jpg"),
        Image.open("../datasets/images/dog.jpg"),
    ]

    # Inference
    with torch.no_grad():
        inputs = processor(
            text=descriptions,
            images=images,
            return_tensors="pt",
            padding=True,
            truncation=True,
        )
        outputs = model(**inputs)

    dot_products_per_text = outputs.logits_per_text

    # Calculate probabilities
    probabilities = dot_products_per_text.softmax(dim=1)

This gives us a probability matrix with the results. The first column shows how likely each image is to contain a cat, and the second column shows how likely it is to contain a dog. We’re not making a direct prediction about whether it’s a cat or dog. Instead, we’re measuring how similar each image is to the descriptions “A photo of a cat” and “A photo of a dog.” Figure 3-17 shows these results. In this case, the probability that the first image shows a cat is 0.9999, which means the probability that it shows a dog is extremely low.

Figure 3-17. Probability matrix for CLIP image classification

Discussion

CLIP is mainly used for image search, recommendation systems, and image classification. In Figure 3-15, I showed how to use CLIP in a RAG system to match a user’s query with a database of images. In this case, we type in keywords and search for matching images. This also works in the opposite direction - we can send an image to the model and search for matching text chunks or other images.

We can use this same concept for image classification. In the past, it took days or weeks to train image classification models. First, we needed a training dataset with labeled images. To ensure accuracy, this was often done manually - human labelers would sit in front of images and assign labels like “dog” or “cat.” This was extremely time-consuming.

These simple classification models have been largely replaced by LLMs, multimodal models, and multimodal embedding models like CLIP. Distinguishing between a dog and a cat, or a train and a car, is no longer challenging for modern transformer models. The latest flagship models can do much more - they can for example automatically classify images of technical components and match them to parts lists, analyze medical imagery to identify potential anomalies and suggest differential diagnoses, or process satellite imagery to detect changes in land use patterns over time.

Tip

More common than using CLIP directly in your RAG system is combining multimodal or vision models with text embedding models. You can use LLMs to create text descriptions of images (like “A dog playing in a garden”) and then apply regular text embedding models to these descriptions. Even the classification example we showed above can be solved in the same way. We can either let the multimodal models make the decision themselves, or we can have multimodal models generate a summary of the image and then compare the text embeddings to the text embeddings of the predefined sentences, like “a photo of a dog”, to figure out if there is a dog in the image.

See also

If you want to dive deeper into CLIP, I recommend to check out the following sources:

Check out the articles on the OpenAI website. They have a detailed CLIP research page explaining the model and its functionality.

For a comprehensive guide on CLIP, read the post by Szymon Palucha.

Additionally, the blog post from LlamaIndex is great for learning how to use OpenAI’s CLIP to build a multimodal RAG system for better video content analysis.

3.6 Performing Text Classification Using Embeddings

Problem

You want to automatically sort text into different categories and use it to make decisions.

Solution

Imagine you’re building a chatbot for a newspaper website. You have all the newspaper’s articles chunked and stored in a vector store, covering topics like politics, sports, and economy. Without proper categorization, information between these sections can easily get mixed up. Individual sentences from different topics might be semantically similar even when they come from articles that have nothing to do with the user’s question.

Most users ask questions that can be answered with information from just one section. If we can identify which section to search before querying the database, we can significantly improve both accuracy and speed of our RAG solution.

For example, the question “What new policies did the government announce this week?” should clearly search only the “Politics” section. While this is obvious to humans, how does our RAG system automatically decide which section is relevant? We need a classification model that can categorize the user’s question and route it to the appropriate content section.

Figure 3-18 illustrates the concept we can use. We will train a classification model using a diverse set of text snippets from various articles and topics. The goal is to predict the topic of the text. For this example, we will use two articles:

One about the history of the Premier League (the top football league in the UK)

One about the history of Deep Learning

We are building a binary classifier that determines whether the user is asking a question about football or Deep Learning. Alternatively, we could also have LLMs perform the classification, but with a classification model we are definitely faster and more cost-effective.

Figure 3-18. Building a text classification model using text embeddings

In this recipe, we use scikit-learn, the most well-known Python library for classical machine learning tasks in Python, and train a Random Forest classifier. We generate an embedding for each text chunk and then split the embedding vectors. For example, if our embedding model has 1536 dimensions, we split the embedding vectors into 1536 individual features, so our input matrix X has 1536 columns, with the Y values being the topic labels.

Tip

For me, scikit-learn has been the bible of the last decade, as it provides mature implementations for almost all classical ML algorithms that you can apply to your data with a single line of code. So I was excited when I found out there was a library called Scikit-LLM that bridges scikit-learn with the new transformer models world and thus enables enhanced text analysis tasks. It provides a range of pre-built functions, including one-shot prompting classifications and chain-of-thought text classifications, which performs a series of reasoning steps before arriving at a final classification.

We are not using Scikit-LLM in this recipe, but you can find out more in the docs or the GitHub repository.

Prerequisites

We’ll use the text-embedding-3-small model from OpenAI to generate embeddings for text classification. The Random Forest classifier from scikit-learn will serve as our classification algorithm, and we’ll use langchain-text-splitters for chunking text documents into smaller pieces for processing.

pip install openai scikit-learn pandas langchain-text-splitters PyPDF2

Steps

In Example 3-9, we load two example files: one about Premier League history and another about Deep Learning history. These topics are completely different from each other. We split the text into smaller pieces using a method called recursive chunking, then create embeddings for each piece of text and the user’s question. Recursive chunking breaks text apart at natural points like periods and other punctuation marks to create manageable chunks (see Recipe 2.5 for more details).

Example 3-9. Text Classification using Embeddings

Next, we generate embeddings for each text chunk and save them along with the text in the data frame chunks_df.

Example 3-10. Text Classification using Embeddings

Now we can train our classifier using the data in chunks_df. We’ll use the embeddings as input data and the topic labels (like Deep_Learning or Premier_League) as the answers we want the model to learn.

Before we can train the model, we need to set up our data properly. The Random Forest classifier needs the data in a specific format - it wants each dimension of our embedding vector to be in its own column. Since our embeddings have 1536 numbers in each vector, we’ll create a table with 1536 columns where each column represents one of those numbers.

Example 3-11. Text Classification using Embeddings

    from sklearn.ensemble import RandomForestClassifier
    from sklearn.datasets import make_classification

    y = chunks_df["label"]
    X = chunks_df["embedding"].apply(
        lambda x: pd.Series(eval(x)) if isinstance(x, str) else pd.Series(x)
    )

    # Train a random forest classifier
    clf = RandomForestClassifier()
    clf.fit(X, y)

Now we can use our trained model to make predictions. In Example 3-12, we test the model with the question “What is the name of the top football league in England?”. The Random Forest classifier tells us which category the question belongs to and how confident it is about that prediction. In this example, the model is 94% confident that this question should be answered using information from the Premier League article rather than the Deep Learning article.

Example 3-12. Text Classification using Embeddings

    # Predict probabilities
    test_embedding = create_embeddings(
        text_chunk="What is the name of the top football league in England?",
        client=client,
    )

    X_test = [test_embedding]

    # Predict the most likely class
    predicted_classes = clf.predict(X_test)  # e.g. ['Premier_League']
    probabilities = clf.predict_proba(X_test)  # e.g. array([[0.06, 0.94]])

Discussion

The idea is to use embeddings to help make automatic decisions. In RAG systems, we often need to answer questions about many different topics. This creates a challenge: sometimes text chunks that seem similar in meaning actually come from completely different subject areas and contain different types of information.

For example, imagine searching for information about “antioxidants.” You might find text chunks from both medical research and environmental engineering that discuss antioxidants. While both topics use similar words and concepts, the medical information talks about health benefits, while the environmental information discusses pollution control. The semantic meaning appears similar, but the actual facts and context are completely different.

By using a classification model, we can automatically sort the user’s question into the right category before searching, which helps us find more relevant and accurate information.

See also

I found Shaw Talebi’s blog post on text embeddings, classification, and semantic search helpful and easy to follow.

In this chapter, we used Random Forest as our classification algorithm, but you can use any other classification algorithm in the same way.

For experimenting with different classification algorithms, scikit-learn offers the most comprehensive collection of ML algorithm implementations for classification, regression, and clustering tasks. It’s an excellent resource for exploring alternatives to Random Forest that might work better for your specific use case.

3.7 Improving Search Results Using a Hybrid Search Approach

Problem

You want to get better search results by combining different search methods.

Solution

We use Hybrid Search to make our search results better and more reliable. Think of it like asking two different experts for help and then combining their answers to get the best possible result.

This approach is similar to how ensemble methods work in Machine Learning. Just like Random Forests combine many decision trees to make better predictions, hybrid search combines different search methods to reduce mistakes and improve accuracy.

Here’s how it works in four steps:

Keyword search (using BM25): This is like traditional search - it looks for exact word matches in your documents and ranks them based on how often those words appear and how important they are.

Create embeddings: We convert the user’s question into numbers (vectors) that capture what the question actually means, not just the specific words used.

Semantic similarity search: We use those number vectors to find content that has similar meaning to the user’s question, even if it uses completely different words.

Combine the results: We merge the rankings from both keyword and semantic search to create a final list that takes advantage of what each method does best.

Figure 3-19. How Hybrid Search works

Prerequisites

To combine keyword and semantic search effectively, we need to run both methods on the same data and merge their results. Here’s what we need:

Keyword search: BM25 algorithm finds documents based on exact word matches.

Semantic search: The text-embedding-3-small model creates embeddings and uses cosine similarity to find meaning-based matches.

Combine results: Reciprocal Rank Fusion (RRF) merges the rankings from both search methods.

pip install openai rank-bm25 sentence-transformers

Solution

In Example 3-13, we start with three sample text chunks and a user question. We use the BM25 algorithm to find which text chunks best match the user’s question based on the actual words used. This creates a list called doc_scores that ranks each text chunk by how well it matches the question - the higher the score, the better the match.

Example 3-13. How to rank text chunks using BM25

from rank_bm25 import BM25Okapi

text_chunks = [
    "The Great Fire of London in 1666 destroyed over 13,000 houses.",
    "Julius Caesar was assassinated on the Ides of March (March 15) in 44 BCE.",
    """The Black Death, is estimated to have killed nearly one-third of the European
    population""",
]
tokenized_chunks = [chunk.split(" ") for chunk in text_chunks]

bm25 = BM25Okapi(tokenized_chunks)

users_query = "Tell me something interesting about diseases in history"
tokenized_query = users_query.split(" ")

doc_scores = bm25.get_scores(tokenized_query)

ranking_keyword_search = (
    pd.DataFrame(doc_scores, columns=["scores"])
    .sort_values(by="scores", ascending=False)
    .index.to_list()
)
from rank_bm25 import BM25Okapi
import pandas as pd
import os

text_chunks = [
    "The Great Fire of London in 1666 destroyed over 13,000 houses.",
    "Julius Caesar was assassinated on the Ides of March (March 15) in 44 BCE.",
    """The Black Death, is estimated to have killed nearly one-third of the European
    population""",
]
tokenized_chunks = [chunk.split(" ") for chunk in text_chunks]

bm25 = BM25Okapi(tokenized_chunks)

users_query = "Tell me something interesting about diseases in history"
tokenized_query = users_query.split(" ")

doc_scores = bm25.get_scores(tokenized_query)

ranking_keyword_search = (
    pd.DataFrame(doc_scores, columns=["scores"])
    .sort_values(by="scores", ascending=False)
    .index.to_list()
)

Next, we create embeddings for both the text chunks and the user’s question. Then we calculate how similar the question is to each text chunk using cosine similarity. This creates a list called doc_scores where higher numbers mean the text chunk is a better match for the user’s question.

Example 3-14. Text classification using embeddings

from sentence_transformers.util import cos_sim

embeddings_df = pd.DataFrame(text_chunks, columns=["text_chunk"])

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
embeddings = []

def create_embeddings(text_chunk, client):
    embedding = (
        client.embeddings.create(
            input=[text_chunk], model="text-embedding-3-small")
        .data[0]
        .embedding
    )
    return embedding

embeddings_df["embedding"] = embeddings_df["text_chunk"].apply(
    create_embeddings, client=client
)

users_question_embedding = create_embeddings(
    text_chunk=users_query, client=client)

# Compute cosine_similarity between documents and query
embeddings_df["similarity"] = cos_sim(
    embeddings_df["embedding"], users_question_embedding
)

ranking_semantic_search = embeddings_df.sort_values(
    by=["similarity"], ascending=False
).index
from sentence_transformers.util import cos_sim

embeddings_df = pd.DataFrame(text_chunks, columns=["text_chunk"])

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
embeddings = []

def create_embeddings(text_chunk, client):
    embedding = (
        client.embeddings.create(
            input=[text_chunk], model="text-embedding-3-small")
        .data[0]
        .embedding
    )
    return embedding

embeddings_df["embedding"] = embeddings_df["text_chunk"].apply(
    create_embeddings, client=client
)

users_question_embedding = create_embeddings(
    text_chunk=users_query, client=client)

# Compute cosine_similarity between documents and query
embeddings_df["similarity"] = cos_sim(
    embeddings_df["embedding"], users_question_embedding
)

ranking_semantic_search = embeddings_df.sort_values(
    by=["similarity"], ascending=False
).index

Once we have the keyword and semantic search scores, we merge them into one. In Example 3-15, we use Reciprocal Rank Fusion (RRF) to combine the two lists. RRF takes the rankings from both search methods and combines them by giving higher scores to documents that rank well in both lists. It works by taking the reciprocal of each document’s rank (1/rank) and adding these values together, which means documents that appear near the top of both lists get the highest combined scores.

Example 3-15. How to calculate a combined rating

#  calculate a combined similarity score using Reciprocal Rank Fusion (RRF)
combined_score = []

for i in range(0, len(ranking_semantic_search), 1):
    k = 60
    rrf_score = 1 / (k + ranking_keyword_search[i]) + 1 / (
        k + ranking_semantic_search[i]
    )
    combined_score.append(rrf_score)

combined_score_df = pd.DataFrame(combined_score, columns=["combined_score"])
new_ranking = (
    combined_score_df.sort_values(
        by=["combined_score"], ascending=False).index + 1
)

#  calculate a combined similarity score using Reciprocal Rank Fusion (RRF)
combined_score = []

for i in range(0, len(ranking_semantic_search), 1):
    k = 60
    rrf_score = 1 / (k + ranking_keyword_search[i]) + 1 / (
        k + ranking_semantic_search[i]
    )
    combined_score.append(rrf_score)

combined_score_df = pd.DataFrame(combined_score, columns=["combined_score"])
new_ranking = (
    combined_score_df.sort_values(
        by=["combined_score"], ascending=False).index + 1
)

Discussion

Hybrid search blends the strengths of keyword and semantic search. Keyword search, like BM25, excels at finding exact matches and rare terms, while semantic search grasps the context and meaning of the text. By using both, we achieve more accurate and dependable search results.

From my experience, hybrid search gives you better and more reliable results. Semantic search is definitely stronger, but I’ve noticed that many users still search the way they’re used to - by typing in specific keywords they think are important. Hybrid search helps capture these keyword-focused queries while also understanding the meaning behind what users are asking for.

See also

For more information on Hybrid Search, I recommend Dr. Leon Eversberg’s blog post on using hybrid search for better RAG retrieval.

Chapter 4. Vector Databases and Similarity Searches

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 6th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the author at dmnkplzr@googlemail.com.

What You’ll Learn

In this chapter, you’ll learn how to:

Choose the right vector database solution based on your project’s scale, complexity, and performance requirements (see Recipe 4.1)

Store and search embeddings using FAISS for fast vector operations in lightweight applications and prototypes (see Recipe 4.2)

Build and deploy vector databases using ChromaDB for small to medium-scale RAG applications with persistent storage (see Recipe 4.3)

Store embeddings in PostgreSQL using the pgvector extension to combine structured and vector data in SQL environments (see Recipe 4.4)

Perform efficient similarity search in PostgreSQL using different distance metrics like cosine similarity and Euclidean distance (see Recipe 4.5)

Speed up vector searches for large datasets using indexing techniques like IVFFlat and HNSW in PostgreSQL (see Recipe 4.6)

Implement hybrid search combining keyword-based (full-text) and semantic search for improved retrieval accuracy (see Recipe 4.7)

Understand the trade-offs between different vector database options from simple libraries to production-ready solutions

Apply best practices for vector database selection, indexing strategies, and performance optimization in RAG systems

Vector databases have grown rapidly in popularity. In recent years, many new companies have started offering vector database solutions, while existing ones have gained significant attention. Figure 4-1 shows when some of the most popular vector databases were first released: FAISS came out in 2017, followed by Milvus and Weaviate in 2019, Vald in 2020, Pinecone in 2021, and ChromaDB in 2023. At the same time, some traditional databases like PostgreSQL and Elasticsearch have added vector search features to their existing platforms.

This chapter can’t cover every vector database option available, so I’ve picked the ones I use most in my own projects. We’ll start with simple vector libraries like FAISS, move to easy-to-use in-memory stores like ChromaDB, and then explore PostgreSQL with the pgvector extension that adds vector search to a regular database.

Figure 4-1. The evolution of vector stores

Before diving into vector databases, consider whether you really need one. For simple RAG workflows or apps, a Numpy array might be sufficient. Just store the embeddings and calculate cosine similarity to the user’s query (see Recipe 3.3).

Tip

While it’s tempting to explore advanced indexing techniques and powerful vector stores, remember that simplicity wins. Focus on delivering solutions quickly and effectively rather than building overly complex setups.

Each vector store has its own strengths and weaknesses. If your database contains several hundred thousand entries but you want to keep things simple, consider using FAISS. It’s a straightforward Python library that performs extremely fast vector searches and can save its index to a file. Another option is in-memory vector stores like ChromaDB, which are equally easy to use. With ChromaDB, you can process your data once, load it into memory, and handle all the vector operations you need for typical RAG applications. If you want to keep your index available after restarting your application, ChromaDB also includes a persistence feature. This allows you to save the index to a file and load it back up when you restart.

Table 4-1 provides a quick comparison of FAISS and ChromaDB, with performance data from the research paper PERFORMANCE ANALYSIS OF CHROMA, QDRANT, AND FAISS DATABASE.

Table 4-1. Comparison of FAISS and ChromaDB. Performance metrics are referenced from the paper PERFORMANCE ANALYSIS OF CHROMA, QDRANT, AND FAISS DATABASE.

Feature
FAISS
ChromaDB

Type

Library

Vector Database

Primary Focus

High-performance similarity search

Developer-friendly, AI apps (RAG)

Scalability

Billions of vectors (with GPU/engineering)

Millions of vectors (single-node)

Performance (500k vectors)

Insertion: 0.08s; Query: 0.309s

Insertion: 434.89s; Query: 0.319s

Persistence

Manual save/load of index

Built-in (local disk, automatic)

Metadata Handling

Limited built-in filtering

Robust, integrated filtering

Tip

FAISS and ChromaDB are very lightweight and extremely easy to set up—literally just a line of code. However, they don’t include advanced features like user management, role-based access control, or built-in monitoring.

However, this basic setup can handle most RAG use cases. Many applications serve fewer than 100 users, yet we often rush to complex solutions. Start simple and focus on delivering real value. Most ML projects fail to reach production—don’t over-engineer your RAG app.

When moving from a prototype to a production RAG app with thousands of potential users, you need to rethink your setup. You’ll need a more robust solution that is scalable, can handle large amounts of data, and is easy to maintain.

This chapter covers various options, from simple vector libraries to full-featured vector databases, to help you choose the right one for your project. Figure 4-2 organizes different databases into three categories based on how complex your project is, how many users you have, and how much data you need to handle.

The diagram shows three types of databases, from simple to complex: Local and Developer-Friendly Vector DBs (like ChromaDB and FAISS) work well for simple projects and experiments, Traditional and Hybrid Databases (like PostgreSQL with pgvector and Elasticsearch) let you work with both regular data and vectors together, and Scalable Cloud-Native Vector DBs (like Weaviate and Pinecone) are built for large production applications that need to handle lots of users and data.

Figure 4-2. The overview of when to use which database

This chapter walks you through different vector database options for RAG applications. We’ll start with simple solutions like FAISS and ChromaDB that work great for prototypes and small projects. Then we’ll dive into PostgreSQL with pgvector for production applications, showing you how to set it up, create tables for embeddings, and use different distance calculations for searches. We’ll also discuss when to consider Pinecone and Weaviate for large-scale RAG applications.

I often use PostgreSQL with pgvector because it’s open source, easy to set up, and integrates well with existing infrastructure. The main reason is that my team and I are already familiar with it, and we use PostgreSQL as our SQL database in most projects anyway. But it’s not just familiarity—I find it really convenient to execute all commands in SQL and easily combine vector search with standard SQL commands on top of the application data already stored in PostgreSQL. I use AWS Aurora to host a managed PostgreSQL instance, but you will also find managed PostgreSQL database services in GCP and Azure.

If you’re starting from scratch and aren’t tied to existing infrastructure, Weaviate or Pinecone could be good options. They’re available in the AWS and Google Cloud marketplaces, making it easy to set them up as managed vector databases. If you’re an Azure user, you can use Azure Cognitive Search or deploy Weaviate on an Azure VM.

Tip

I recommend building a working solution first—you don’t need a full vector database right away. Start simple with Numpy arrays or FAISS for prototypes, then optimize later. Treat your RAG components (vector store, LLM, embedding models) as interchangeable parts. In the fast-changing GenAI space, there are always new models and databases. Focus on the business logic behind your app and the problem you want to solve, all other components can be replaced at any time.

4.1 Choosing the right Vector Database

Problem

You need to choose the right vector database for your RAG application.

Solution

For RAG, you need a place to store and calculate with embedding vectors with hundreds or thousands of dimensions. Figure 4-3 illustrates the landscape of vector storage solutions, organizing them into four distinct categories based on their architecture and primary use cases.

The diagram shows Vector Libraries (like FAISS, HNSWlib, Annoy, and ScaNN) on the left—these are lightweight, file-based solutions perfect for prototyping and smaller applications. Pure Vector Databases (including Pinecone, Weaviate, Qdrant, milvius, and ChromaDB) appear on the right, designed specifically for vector operations and optimized for large-scale deployments. Traditional Text Search Databases (such as Elasticsearch and OpenSearch) are shown at the top, offering hybrid capabilities that combine full-text and vector search. Finally, Vector-Capable SQL and NoSQL Databases at the bottom (like PostgreSQL with pgvector, MongoDB, Redis, and Cassandra) extend existing database systems with vector search capabilities.

For smaller apps and prototypes, you can start with vector libraries like FAISS or lightweight databases like ChromaDB. FAISS is not really a database, but rather a Python library—similar to Numpy or Pandas—optimized for vector operations. It’s ideal for early experiments or standalone RAG workflows where you reload data on each run and do not require robust data storage. ChromaDB bridges the gap between libraries and full databases, offering an in-memory solution that’s perfect for small applications or initial prototypes while providing optional persistence.

For production applications requiring scalability and advanced features, pure vector databases like Pinecone and Weaviate are optimized for large-scale use. Alternatively, if you already have existing database infrastructure, extending SQL databases like PostgreSQL with pgvector or NoSQL solutions like MongoDB with vector capabilities can provide a unified platform for both structured and vector data.

Figure 4-3. Different kinds of vector stores organized by type and use case

In my daily work, I choose different vector storage options based on what the project needs and which stage it’s in. Here’s a simple breakdown of the different types and how they fit together. You can also check [Link to Come] for the main benefits and drawbacks of popular vector storage options.

Vector Libraries (like FAISS): Perfect for small projects. FAISS is a fast, simple library for working with vectors. It’s great for testing ideas on your computer and building basic prototypes.

In-memory and File-based Vector Stores (like ChromaDB): ChromaDB is lightweight and quick to set up. It’s excellent for building quick prototypes and small projects, but it doesn’t have the advanced features you’d need for large production systems.

Regular Databases with Vector Features (like PostgreSQL, Redis): Databases like Redis, MongoDB, and PostgreSQL weren’t originally built for vectors, but you can add extensions to make them work with vectors. This lets you stick with databases you already know while adding vector search capabilities.

Dedicated Vector Databases (like Pinecone, Weaviate): Pinecone and Weaviate are built specifically for handling vector data. They work best with large amounts of data and offer powerful vector search features. They’re designed for AI applications and come with built-in features that RAG applications often need, like hybrid search, filtering by metadata, AI agents, and more.

Here’s a summary of the pros and cons for each type of vector store:

Table 4-2. Pros and Cons of Different Vector Stores and Libraries

Database
Type
Pros
Cons

FAISS

Vector Library

Ultra-fast vector search for large datasets

GPU acceleration available

Supports exact and approximate searches

No built-in database features (persistence, metadata filtering)

ChromaDB

In-memory / in-file vector store

Easy to set up with developer-friendly API

Good for prototyping and small/medium applications

Integrates well with Python ML frameworks

Limited scalability for production workloads

No native encryption

PostgreSQL (pgvector)

Hybrid (Relational Database with Vector Extension)

Combines structured and vector data in one platform

Uses existing PostgreSQL infrastructure

Strong ACID compliance

Can be slower for large-scale vector operations

Limited advanced vector features

Weaviate

Pure Vector Store

Highly scalable and cloud-native

Combines vector and scalar search

Built-in NLP and classification modules

Complex setup and learning curve

Resource-intensive and potentially costly

Requires performance tuning

Pinecone

Pure Vector Store (Fully Managed Service)

Fully managed with no maintenance required

Excellent scalability and performance

Real-time data ingestion

Higher cost than open-source options

No self-hosting option

Limited customization and potential vendor lock-in

Discussion

Choosing the right vector database depends on what you’re building. For small projects and experiments, simple tools like Numpy work fine for basic similarity searches. If you’re working with hundreds of thousands of vectors or more, the FAISS library is a good choice since it’s built specifically for fast vector operations.

For larger applications with millions of records and thousands of users, you’ll need a full-featured vector database. When making this choice, consider these questions:

Can you use cloud services instead of managing your own servers?

Does it integrate well with your existing tools and systems?

How fast does it perform with large datasets?

Is there good community support when you need help?

Don’t overthink this decision. Start with a simple prototype and improve it over time. Most projects benefit more from getting something working quickly.

See also

There aren’t actually that many resources that directly compare a larger number of vector stores. This is likely because most teams eventually settle on a vector database and rarely migrate unless there is a compelling reason. Still, I found Plaban Nayak’s article worth reading: Which Vector Database Should You Use? Choosing the Best One for Your Needs.

4.2 Storing and Searching Embeddings using FAISS

Problem

You need a fast and simple way to store embeddings and perform similarity searches.

Solution

FAISS is a great open-source option. It was one of the first libraries to provide fast similarity search on large-scale datasets. In this recipe, I will show you how to store and retrieve embeddings using FAISS.

Prerequisites

With just a few steps, we can get FAISS up and running. We only need our package manager pip to download the FAISS library. You can use faiss-gpu or faiss-cpu. Since I am running this experiment on my laptop without an external GPU, I am using faiss-cpu. If you have a GPU, you can speed up the vector searches using faiss-gpu.

pip install faiss-cpu openai numpy

Steps

We’ll set up and use FAISS in two simple steps: creating an index to store our data, then searching through it.

Step 1: Create a FAISS index

An index is like a filing cabinet where FAISS stores your embeddings (vector representations of text) so it can quickly find similar ones later.

In Example 4-1, we use OpenAI’s embedding model to convert our text into numbers (embeddings) and store them in a NumPy array. Then we create a FAISS index using faiss.IndexFlatL2, which uses L2 distance to measure how similar embeddings are. We need to tell FAISS how many dimensions our embeddings have - we get this from document_embeddings.shape[1].

Example 4-1. How to create a FAISS index

import faiss
import numpy as np
from openai import OpenAI

# Example list of sample strings
text_chunks = [
    "The sky is blue.",
    "The sun is shining.",
    "I love chocolate.",
    "Ice cream is delicious.",
    "Roses are red.",
    "Violets are blue.",
]

# Initialize the OpenAI embeddings model
client = OpenAI()
model = "text-embedding-3-small"

# Generate embeddings for the sample strings
def get_embedding(text):
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

document_embeddings = np.array(
    [get_embedding(text) for text in text_chunks]
)

# Convert embeddings to float32 (FAISS requires float32 type)
document_embeddings = document_embeddings.astype("float32")

# Create a FAISS index (using L2 distance)
index = faiss.IndexFlatL2(document_embeddings.shape[1])

# Add embeddings to the index
index.add(document_embeddings)

Step 2: Search the index

Once we have our index set up with data, we can search for text similar to a user’s question.

Now let’s search through our stored data using FAISS. The example in Example 4-2 shows how this works step by step:

We take the user’s question “What color do Violets have?” and convert it to an embedding using the same OpenAI model

We store this embedding in a NumPy array

We search our index to find the most similar stored embeddings

FAISS returns the closest matches to our query

Warning

Always use the same embedding model for both storing and searching. In this example, we use OpenAI’s text-embeddings-small-3 for both steps.

Example 4-2. How to query the FAISS index

# generate a query embedding for the user query
query = "What color do Violets have?"
query_embedding = np.array(get_embedding(query)).astype("float32")

# Perform the search: k = number of closest documents you want to retrieve
k = 5
distances, indices = index.search(query_embedding.reshape(1, -1), k)

# Retrieve the documents corresponding to the indices
retrieved_documents = [text_chunks[i] for i in indices[0]]

Figure 4-4 shows the search results: distances, chunk locations, and the actual text chunks. The sentence “Violets are blue.” has the smallest distance, making it the best match for our query.

Figure 4-4. FAISS query result example

This completes the retrieval step in a basic RAG application. Next step would be, we take the best matching text chunks (like the top 3 results) and combine them with the user’s question in a prompt for the language model to generate an answer.

Discussion

FAISS, created by Facebook AI Research in 2017, remains an excellent choice for small RAG applications. It’s not designed for large-scale chatbots serving thousands of users, but it excels at standalone workflows and prototypes.

The most valuable RAG applications aren’t always chatbots. Instead, they’re often simple, focused workflows that solve specific business problems (see Figure 4-5). For example, you can use multimodal models to read scanned PDFs or photographs and extract key information. This transforms piles of unreadable documents into structured, actionable data.

FAISS is perfect for these temporary, task-specific workflows. You can quickly build an index from a few documents, extract what you need, then delete the index when you’re done.

Figure 4-5. Using FAISS for RAG workflow

See also

You can find the official documentation in the GitHub repository of FAISS. It’s my go-to library for small RAG workflows. For productive RAG apps, I recommend using a hybrid database (e.g., PostgreSQL, see Recipe 4.4) or a pure vector database (e.g., Weaviate or Pinecone).

Here’s my preferred order:

Numpy for simple embedding experiments (see Recipe 3.3).

FAISS for larger datasets and standalone RAG workflows (see Recipe 4.2).

ChromaDB for quick prototypes and MVPs (see Recipe 4.3).

PostgreSQL with pgvector for production apps with lots of structured data (see Recipe 4.4). Weaviate or Pinecone for large-scale production apps with advanced vector operations.

When working on RAG use cases in a corporate environment, data security and privacy play a major role. Since most companies already have contracts and managed accounts with AWS, Google, or Microsoft, it’s easy to use one of the available options on these platforms. AWS offers managed PostgreSQL instances with AWS Aurora. The same applies to Pinecone and Weaviate, both of which are available in the AWS and Google Cloud marketplaces. Since Weaviate is open source, you can also deploy it locally or on a virtual machine.

4.3 Storing and Working with Embeddings in a Chroma Vector Database

Problem

You need a lightweight vector database to store and retrieve embeddings for your RAG application.

Solution

ChromaDB is a lightweight vector database that’s perfect for prototypes and small to medium-scale applications. Unlike FAISS, which is just a library, ChromaDB provides database features like persistence and metadata handling while remaining simple to use.

You have two storage options with ChromaDB: - Ephemeral Client: Stores data in memory (lost when application stops) - Persistent Client: Saves data to disk for permanent storage

Figure 4-6 illustrates ChromaDB’s workflow. While ChromaDB can automatically handle text processing and embedding generation, I recommend creating embeddings yourself first, then storing them in ChromaDB. This approach gives you more control over the embedding model and makes it easier to switch models later.

This creates a two-step process: First, you load text documents and create vector embeddings (either using ChromaDB’s built-in model or your own). Second, ChromaDB stores everything (collections, embeddings, and metadata), either in memory or persistently to disk.

Figure 4-6. How to store and retrieve vector embeddings from ChromaDB

Prerequisites

Setting up ChromaDB is simple. Install the ChromaDB library using the pip package manager. Additionally, we install the OpenAI library to use OpenAI’s embedding models.

pip install chromadb openai

Steps

Working with ChromaDB involves three main steps: first, you’ll create a client and collection to establish your vector storage, then add your text documents to the collection where ChromaDB will handle the embedding process, and finally query the collection to find similar documents for your RAG application.

Step 1: Creating a chroma client and collection

A collection is the designated storage for your embeddings, documents, and any additional metadata. In Example 4-3 we use the persistent client, which saves the vector store to the vector_store_path. We are using you can use “persist_directory”.

Example 4-3. Chroma Vector Database

import chromadb

# vector store settings
VECTOR_STORE_PATH = r"../02_Data/00_Vector_Store"
COLLECTION_NAME = "my_collection"

# get/create a chroma client and collection
chroma_client = chromadb.PersistentClient(path=VECTOR_STORE_PATH)
collection = chroma_client.get_or_create_collection(name=COLLECTION_NAME)

Step 2: Adding text documents to the collection

Chroma handles tokenization, embedding, and indexing for you. By default, Chroma uses the Sentence Transformers “all-MiniLM-L6-v2” model for creating embeddings. If you prefer to use a different embedding model, I recommend generating the embeddings first and then inserting them into Chroma. This approach allows you to be more flexible with your choice of embedding models and easily switch between them.

Example 4-4. How to fill the ChromaDB collection with text documents

from openai import OpenAI

text_chunks = [
    "The sky is blue.",
    "The sun is shining.",
    "I love chocolate.",
    "Ice cream is delicious.",
    "Roses are red.",
    "Violets are blue.",
]

# Generate embeddings
client = OpenAI()
model = "text-embedding-3-small"
def get_embedding(text):
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

embeddings_model = client  # for compatibility with later code
embeddings_list = [get_embedding(text) for text in text_chunks]

# add data frame to collection
collection.add(
    embeddings=embeddings_list,
    documents=text_chunks,
    ids=[
        str(i) for i in range(len(text_chunks))
    ],  # create a list of strings as index
)

That’s it. Your embeddings are now stored in ChromaDB.

Step 3: Querying the collection

To proceed with our RAG workflow, we need to perform a similarity search on the data stored in the ChromaDB vector store (see [Link to Come]). ChromaDB uses cosine similarity by default to measure the similarity between the query embedding and the stored embeddings. I usually stick with cosine similarity, as it works well with most embedding models.

Example 4-5. How to query the ChromaDB collection

# query collection
query = "What is the color of the sky?"
query_embedding = get_embedding(query)
results = collection.query(query_embeddings=[query_embedding], n_results=3)

The code snippet returns the three closest matching text chunks (k=3) along with their cosine distances. Figure 4-7 shows the resulting dictionary. The most important attribute for the next steps in the RAG flow is results[documents]. It contains the text chunks that are most similar to the user’s query. You can now provide these text chunks, together with the user’s question, to the LLM, which will generate a well-formulated answer.

Figure 4-7. Working with embeddings in ChromaDB

Discussion

Chroma is an excellent open-source choice. It’s one of the first lightweight vector store libraries and remains my go-to for experiments and small prototypes.

See also

I liked the getting started guide from the ChromaDB website. It’s a good starting point to get a feeling for the library and what it can do.

4.4 Storing Embeddings in PostgreSQL Using the pgvector Extension

Problem

You want to use PostgreSQL and pgvector as a vector store for your RAG application.

Solution

PostgreSQL is one of the most popular open-source databases and has been my go-to choice for years—even before RAG applications became common. Fortunately, with the pgvector extension, I don’t have to change that. The pgvector extension allows us to store vectors and perform vector similarity searches, making PostgreSQL a great fit for RAG applications. The big advantage is that you can build on existing infrastructure and don’t need any special upskilling, since you continue working in the familiar SQL environment. In PostgreSQL, vector search is just another operator in your SQL scripts, so you can easily connect and upgrade existing tables and SQL procedures.

Requisition

First, you need to install PostgreSQL and the pgvector extension. You can either use Docker or install PostgreSQL and pgvector directly on your computer. I recommend the Docker approach because it’s easy to manage and allows you to quickly deploy the setup to a server or the cloud.

pip install psycopg2-binary openai

I will show you both methods, first the manual installation and then the Docker setup. For the manual installation, follow these steps:

Install PostgreSQL: Download and install PostgreSQL from the official website: https://www.postgresql.org/download/.

Install pgvector: The installation process depends on your operating system:

Windows: Clone the pgvector repository and follow the instructions in the repo’s readme.

Mac: Use Homebrew to install pgvector.

Alternatively, I can highly recommend the installation via Docker, which is why the rest of this recipe focuses on it. First, install Docker from the official website: https://docs.docker.com/get-docker/.

Now, let’s see how to do this with Docker. All you need is Docker Desktop. You can find everything you need to install Docker Desktop in the official Docker documentation. Once Docker is installed, you can set up PostgreSQL with pgvector using a Docker container. To do this, follow these steps:

Create a file named docker-compose.yml in your project directory.

Add the following code (see Example 4-6) to the docker-compose.yml file. This file uses the ankane/pgvector` image, including the pgvector pre-installed extension.

Example 4-6. Docker compose file for PostgreSQL with pgvector

version: '3.8'
services:
  db:
    image: ankane/pgvector
    container_name: postgres_with_pgvector
    restart: always
    environment:
      POSTGRES_USER: rag_cookbook_user
      POSTGRES_PASSWORD: rag_cookbook_user_pw
      POSTGRES_DB: rag_cookbook
    ports:
      - "5432:5432"
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:

To start the container, navigate to the directory where the docker-compose.yml file is located and run the following command:

docker-compose up -d

This command will download the ankane/pgvector image, create a container named postgres_with_pgvector, set up a PostgreSQL user rag_cookbook_user with the password rag_cookbook_user_pw, and map port 5432 from the container to port 5432 on your computer.

To verify if your container is running, open the Docker dashboard and check the status of the container (see Figure 4-8).

Figure 4-8. How to check if the Docker container is running

To connect to the database, we use the psycopg2 Python package for PostgreSQL connections. All database operations in this recipe run through Python code. If you want to browse the database directly, you can use pgAdmin (a free tool included with PostgreSQL) or DBeaver (which works with PostgreSQL and many other databases).

Steps

If you’re familiar with relational databases like PostgreSQL, MySQL, or SQLite, you already know almost everything you need. The only new thing with pgvector is a vector data type that lets you store vectors (lists of numbers) and search through them.

Let’s create a table with a vector column and add some sample data. Example 4-7 shows how to create a table with three columns, including one for vectors. Make sure your vector column size matches your embedding model. For example, OpenAI’s text-embedding-3-small model creates vectors with 1536 numbers, so I set the vector column to store exactly 1536 dimensions.

Example 4-7. How to create a table with a vectors column

CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE IF NOT EXISTS embeddings (
 id SERIAL PRIMARY KEY,
 text_chunk TEXT,
 embedding VECTOR(1536))

Before creating the table in PostgreSQL, we need to set up a new database. Let’s create a database named rag_cookbook on the local PostgreSQL server. In Example 4-8, we grant the user all privileges, including both read and write access. This allows us to create tables, insert data, and perform queries on the database.

Example 4-8. How to create a PostgreSQL user and grant privileges

-- Connect to your PostgreSQL instance as a superuser (e.g., postgres)
CREATE USER rag_cookbook_user WITH LOGIN PASSWORD 'rag_cookbook_user_pw';

-- Grant all privileges on the database
GRANT ALL PRIVILEGES ON DATABASE rag_cookbook TO rag_cookbook_user

Now let’s connect to the newly created database using Python and create a sample table. In Example 4-9, we create a table that includes a vector column. This table will store text chunks along with their embeddings.

Example 4-9. Create table with a vector column

import psycopg2
from psycopg2 import Error

conn = psycopg2.connect(
    dbname="rag_cookbook",
    user="rag_cookbook_user",
    password="rag_cookbook_user_pw",
    host="localhost",
    port="5432",
)

cur = conn.cursor()

# creates the vector extension if it does not exist
cur.execute("""CREATE EXTENSION IF NOT EXISTS vector""")

# execute the query to create the table including the vector column
cur.execute(
    """CREATE TABLE IF NOT EXISTS embeddings_table(
        id integer PRIMARY KEY,
        text_chunk TEXT,
        embedding vector(1536)
    )"""
)

conn.commit()

Next, let’s populate the table with sample data. In Example 4-10, we loop through a list of text chunks, generate embeddings using the OpenAI embedding model, and then insert the text chunks and their embeddings into the embeddings table using an appropriate SQL query.

Example 4-10. Write vectors to PostgreSQL

from openai import OpenAI

# Define text chunks
text_chunks = [
    "The sky is blue.",
    "The sun is shining.",
    "I love chocolate.",
    "Ice cream is delicious.",
    "Roses are red.",
    "Violets are blue.",
]

client = OpenAI()
model = "text-embedding-3-small"

def get_embedding(text):
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

index = 0
cur = conn.cursor()

# Insert the embeddings into the table
for text_chunk in text_chunks:
    embedding = get_embedding(text_chunk)
    cur.execute(
        "INSERT INTO embeddings_table (id, text_chunk, embedding) VALUES (%s, %s, %s)",
        (index, text_chunk, embedding),
    )
    index += 1

The function fills the PostgreSQL table with text chunks and their corresponding embeddings. Figure 4-9 shows how the table is populated: six text chunks along with their embeddings.

Figure 4-9. How the embeddings table looks in PostgreSQL

You can use this table like any other SQL table for select, create, and update operations. With the pgvector extension, you can also perform similarity searches, such as cosine similarity, Euclidean distance, or inner product. For more details, see Recipe 4.5.

Discussion

PostgreSQL with pgvector allows us to perform fast vector searches directly in our PostgreSQL database. pgvector supports all classic vector search operations like calculating cosine similarity or Euclidean distance.

What makes PostgreSQL extremely convenient for me is that I only need one database for both the classic application data, like user profiles, search history, or any other kind of relational data we need to operate our app, and the vector data I need for the retrieval part of many RAG apps. In the tables, I just use the vector column type for this kind of data, and that’s all.

This gives me the flexibility to execute SQL queries that include all other data from our application. For example, I can use any column in the table to first filter my entire dataset before performing the similarity search on a limited subset.

See also

For more details on performing similarity searches in PostgreSQL using the pgvector extension, see Recipe 4.5.

While PostgreSQL isn’t specifically designed as a vector database, it is highly versatile and practical for many use cases. For an entertaining perspective on PostgreSQL’s capabilities, I recommend the video from the Fireship YouTube channel: I replaced my entire tech stack with Postgres…​.

4.5 Performing similarity search in PostgreSQL

Problem

You want to find the text chunks most similar to a user’s query using PostgreSQL.

Solution

The pgvector extension supports common similarity metrics like cosine similarity, Euclidean distance, and inner product. While Recipe 4.4 shows how to store embeddings in PostgreSQL, this recipe focuses on performing similarity searches.

Requisition

Ensure you have a PostgreSQL server with the pgvector extension installed. In Recipe 4.4 you can find instructions on installing the extension and how to load data into the table.

We’ll use the Python library psycopg2 to connect to the PostgreSQL server and execute SQL queries. We also need the openai library to generate embeddings for our search queries using OpenAI’s embedding models.

pip install psycopg2-binary openai

Steps

In Example 4-11, we use the cosine similarity operator (<=>) to calculate similarity. Pgvector also supports other similarity measures, such as Euclidean distance and inner product. If you’re unsure which measure to use, start with cosine similarity—it aligns well with embedding models by focusing on the vector’s orientation rather than magnitude.

Here’s a summary of the supported similarity measures and their corresponding SQL operators:

Cosine Distance: (embedding <=> vector)

Euclidean Distance: (embedding <-> vector)

Inner Product: (embedding <#> vector)

Taxicab Distance: (embedding <+> vector)

In Example 4-11, we use cosine similarity to query the table created in Recipe 4.4. This table contains random text chunks about weather, flowers, and food. The test query is “Sweets are delicious.”.”

Example 4-11. Similarity search

from openai import OpenAI

# Initialize the OpenAI embeddings model
client = OpenAI()
model = "text-embedding-3-small"

# Example text chunk
text_chunk = "Sweets are delicious."

# Get the embedding
response = client.embeddings.create(input=text_chunk, model=model)
embeded_query = response.data[0].embedding

cur = conn.cursor()
cur.execute(
    f"""
    SELECT 1 - (embedding <=> '{embeded_query}') AS cosine_similarity, *
    FROM embeddings_table
    ORDER BY 1 - (embedding <=> '{embeded_query}') DESC
    LIMIT 20;
"""
)

results = cur.fetchall()

Figure 4-10 shows the closest matches to the query “Sweets are delicious.” The chunks “Ice cream is delicious.” and “I love chocolate.” are ranked highest.

Figure 4-10. How to calculate cosine similarity in PostgreSQL

Discussion

Pgvector supports important similarity measures like cosine similarity, Euclidean distance, and inner product. Since PostgreSQL is mainly a relational database, it’s also excellent for storing other application data, such as user information.

In PostgreSQL, you can easily create SQL queries that combine both structured data and vector searches. For instance, you can filter data based on user groups and then perform a vector search on the filtered results.

See also

PostgreSQL and pgvector are popular and well-supported, so you can easily find resources and guides online. The official documentation is also excellent. For more information and key functions, visit the pgvector extension’s GitHub repository.

4.6 Speeding Up Vector Searches in PostgreSQL using Indexing Techniques Supported by pgvector

Problem

You want to make vector searches faster on large datasets.

Solution

In simple terms, indexing techniques help organize our data to make searches faster. PostgreSQL with the pgvector extension supports two popular indexing methods: IVFFlat and HNSW.

Prerequisites

For the experiments in this recipe, I used a dataset with around 200,000 embeddings. While this isn’t extremely large, it’s sufficient for our tests. Generally, indexing methods become more crucial when dealing with tables containing millions of entries.

To try it yourself, you can use a random text dataset from Kaggle. I used the job description dataset, which you can find here. To increase the number of rows, I artificially created more job descriptions and added them to the table.

pip install psycopg2 openai pandas

Example 4-12 shows the table we create. How to fill it is explained in Recipe 4.4.

Example 4-12. Create a table for the similarity search using indexing techniquess

# execute the query to create the table including the vector column
cur.execute(
    """
    CREATE TABLE IF NOT EXISTS job_description_table(
        id integer PRIMARY KEY,
        text_chunk TEXT,
        embedding vector(1536)
    )
    """
)

Steps

Let’s compare the performance of three search methods: a full search, HNSW indexing, and IVFFlat indexing. We’ll start with the full search and then move on to the two indexing techniques.

In the first example, we use the IVFFlat index. In the second example, we use the HNSW index. If you want to know how those indexing techniques work, see “Discussion”.

You can find the details here:

“Performing a Full Search Using Cosine Similarity”: This gives us a baseline to compare with the indexed search.

“How to create a IVFFlat index”: This method groups data into clusters. Each cluster has a list of data points, making searches faster by only looking at relevant clusters.

“How to create a HNSW (Hierarchical Navigable Small Worlds) index”: This method uses a multi-layered graph to organize data points. The search process navigates through these layers to quickly find approximate nearest neighbors.

Performing a Full Search Using Cosine Similarity

Example 4-13 demonstrates how to perform a full search using cosine similarity. This method sequentially searches through all entries without using an index, making it the slowest option for large datasets. The strength of indexing techniques only becomes apparent with very large datasets. For small datasets, a full search can sometimes even be faster.

Example 4-13. How to perform a full search using cosine similarity (sequential search, without index)

full_search_sql = f"""
    EXPLAIN ANALYZE SELECT 1 - (embedding <=> '{str(query_embedding)}')
    AS cosine_similarity,
    * FROM job_description_table
    ORDER BY 1 - (embedding <=> '{str(query_embedding)}') DESC
    LIMIT 20;
    """

How to create a IVFFlat index

To compare the search speed of a full search with an indexed search, I first duplicate the original table and then create an index for it. Example 4-14 shows you how to do this.

To make sure the index is actually used, we turn off sequential mode. This is crucial for small datasets because PostgreSQL might otherwise opt for a full search, which can be just as fast. However, even with this setting, PostgreSQL sometimes still defaults to a full search if it doesn’t clearly see the benefit of using the index.

Example 4-14 shows how to create an IVFFlat index and perform a similarity search. First, it duplicates the job_description_table and creates an index on the new table. This allows us to compare search speeds with and without indexing. Then, it performs a similarity search on the new table and uses EXPLAIN ANALYZE to display the query plan and execution time.

Example 4-14. Create a IVFFlat (Inverted File with Flat Compression) index on the table

from openai import OpenAI

# generate embedding vector to an example query regarding jobs
query = "I am looking for a job as a data scientist in Berlin."
client = OpenAI()
model = "text-embedding-3-small"
response = client.embeddings.create(input=query, model=model)
query_embedding = response.data[0].embedding

# perform similarity search
ivfflat_sql = f"""
    DROP TABLE IF EXISTS test_embedding_table;
    CREATE TABLE test_embedding_table AS SELECT * FROM job_description_table;
    CREATE INDEX ON test_embedding_table USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 30);  -- Increase the number of lists for faster search
    SET ivfflat.probes = 3;  -- Reduce the number of probes for faster search
    EXPLAIN ANALYZE SELECT 1 - (embedding <=> '{str(query_embedding)}')
    AS cosine_similarity, *
    FROM test_embedding_table
    ORDER BY 1 - (embedding <=> '{str(query_embedding)}') DESC
    LIMIT 20;
    """

cur.execute(ivfflat_sql)
ivfflat_search = cur.fetchall()

In my test with a database of 200,000 job descriptions and embeddings, using the IVFFlat index made the search about twice as fast. This is still a speedup that one could debate the value of, but you’ll notice that when you work with datasets of several hundred thousand or several million data entries, indexing suddenly becomes crucial to keep the computation time per request low.

In the following, I’ll use HNSW instead of IVFFlat, which is also a popular indexing method.

How to create a HNSW (Hierarchical Navigable Small Worlds) index

In Example 4-15 we create an HNSW index. In the code snippet you can see that we use the vector_cosine_ops option. This optimizes the index for cosine similarity searches. If you want to use a different similarity metric, you can choose from the following options:

vector_ip_ops → Inner Product (for cosine similarity with normalized vectors).

vector_l2_ops → Euclidean Distance.

To create and query the HNSW index, the three most important parameters we set are:

m = 16 sets the maximum number of connections per layer in the HNSW graph. The m parameter determines how many connections each data point has to its neighbors. Increasing m makes the graph denser, speeding up search queries but increasing index build time and memory usage. Typical values range from 16 to 64

ef_construction = 200 controls the size of the dynamic candidate list during index construction, affecting build quality. A higher value allows the algorithm to consider more candidates, improving index quality.

SET hnsw.ef_search = 50; controls how many candidates to consider during search, balancing search accuracy vs. speed. Higher values give more accurate results but slower queries.

Example 4-15 shows how such a query looks with an HNSW index.

Example 4-15. Perform HNSW similarity search in PostgreSQL

hnsw_test_sql = f"""
    DROP TABLE IF EXISTS test_embedding_table;
    CREATE TABLE test_embedding_table AS SELECT * FROM job_description_table;
    CREATE INDEX ON test_embedding_table USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 200);
    SET hnsw.ef_search = 50;
    EXPLAIN ANALYZE SELECT 1 - (embedding <=> '{str(query_embedding)}')
    AS cosine_similarity, *
    FROM test_embedding_table
    ORDER BY 1 - (embedding <=> '{str(query_embedding)}') DESC
    LIMIT 20;
"""

cur.execute(hnsw_test_sql)
hnsw_search = cur.fetchall()

Discussion

You don’t need to be an expert in indexing techniques to build a RAG system. The main thing to remember is: if you are working with datasets well over 100,000 entries — especially in the multimillion or billion range — then indexing techniques become crucial. Remember how important indexing is at that scale and the differences between the main approaches. Below is a summary of the main differences between IVF and HNSW.

IFV is a popular method for finding similarities between different items. It works by creating a database index that stores content and connects it to its position in a table or document. We divide the whole collection of items into partitions and their centroids. Each item can only be part of one partition at a time. When we search for similarities, we use the partition centroids to quickly find the items we are looking for (see Figure 4-11).

Figure 4-11. How to query a IVF index by finding the closest centroid

When searching for nearby points, we typically look in the closest centroid. However, if points are near the edge of a neighboring centroid, we might miss them. This is why it’s called Approximate Nearest Neighbor search. We sacrifice a bit of accuracy for faster results.

Figure 4-12. How to query a IVF index by identifying the closest data point to the centroid

To avoid this issue, we search multiple partitions instead of just one. However, the underlying problem still remains. We lose some accuracy. But that is usually ok, speed is more important when dealing with datasets with millions of rows.

On the other hand, HNSW builds a multi-layered graph where higher layers feature sparser connections that cover broader regions of the data space, acting as shortcuts for efficient search. The bottom layer, in contrast, is a dense network of local connections for detailed proximity searches (see Figure 4-13)

Tip

Analogy: Think of HNSW like a multi-layered map of a country

The highest layer is like a national highway map. It has very few roads (connections) but they cover vast distances, allowing you to quickly get from one major city (data point) to another across the country. These are your “shortcuts” for efficient broad searches.

As you go down to lower layers, it’s like zooming into a state map, then a county map, and finally a local street map. Each layer has progressively more roads (connections) that cover smaller, more detailed areas.

The bottom layer is like a detailed local street map. It has many, many roads (dense local connections) that help you navigate precisely within a neighborhood to find a specific house (the exact nearest neighbor).

Figure 4-13. How the HNSW layer approach works

The question you’re probably asking yourself now is whether you really need indexing.

Indexes are great for huge datasets, but how is a huge dataset defined? Usually tables with millions to billions of rows or several gigabytes in size, so the table is to large to fit into memory.

Rule of thumb:

Query on a small table (10,000 rows) is fine without an index.

Query on a 100-million-row table without an index could take minutes instead of milliseconds.

In between, you need to experiment a bit and find out if it’s really worth it. For vector tables with up to one hundred thousand entries, probably not.

The databases use different indexing methods. Here you can find a high-level overview:

FAISS (Facebook AI): Uses PQ, IVF, HNSW

Annoy (Spotify): Uses a tree-based approach

HNSWlib: Pure HNSW implementation

Milvus: Supports multiple index types (HNSW, IVF, PQ)

Weaviate, Pinecone, Qdrant: Use graph-based and ANN indexing

See also

For a simple explanation of indexing techniques, check out the Pinecone website’s tutorial series. James Briggs’ YouTube channel also offers easy-to-understand explanations.

For more details on using these methods in pgvector, visit the official pgvector GitHub repository.

4.7 Combining Keyword and Similarity Search to Get Better Results (Hybrid Search) with PostgreSQL

Problem

You want to combine keyword search and semantic search to get more accurate results.

Solution

RAG apps with a vector database work like a search engine with an LLM that interprets search results to generate an appropriate answer. To make the search engine more robust, we can use both keyword-based and semantic search together, called hybrid search.

You can find more details on hybrid search concepts and Python implementation in Chapter 6 (see Recipe 3.7). In this recipe, we focus on how to perform hybrid search in PostgreSQL.

Prerequisites

Just a PostgreSQL server with pgvector. See Recipe 4.4 for setup instructions. To connect to our PostgreSQL database via Python and execute commands, we use the psycopg2 library.

pip install psycopg2-binary

Steps

In machine learning, it’s common to use multiple approaches to solve the same problem and then combine their results for better overall performance. We can apply this same principle to search by combining PostgreSQL’s built-in keyword search with pgvector’s semantic search capabilities.

To do this, Example 4-16 follows these steps:

Copy the job_description_table.

Add a new column called tsv with the type TSVECTOR.

Use the ts_rank function to rank documents based on their relevance to a query.

Use the plainto_tsquery function to convert the query string into a TSQUERY object.

The resulting column text_score will show the relevance score of the text.

Example 4-16. How to perform a hybrid search in PostgreSQL

# generate embedding vector to an example query regarding jobs
from openai import OpenAI

query = "I am looking for a job as a data scientist in Berlin."
client = OpenAI()
model = "text-embedding-3-small"
response = client.embeddings.create(input=query, model=model)
query_embedding = response.data[0].embedding

# perform similarity search
hybrid_search_table = f"""
    DROP TABLE IF EXISTS test_embedding_table;
    CREATE TABLE test_embedding_table AS
    SELECT *, to_tsvector(text_chunk) AS tsv
    FROM job_description_table;
    """

cur.execute(hybrid_search_table)

hybrid_search_query = f"""
    WITH ranked_docs AS (
        SELECT
            id,
            text_chunk AS text,
            ts_rank(tsv, plainto_tsquery('PostgreSQL')) AS text_score,
            1 - (embedding <=> '{str(query_embedding)}') AS vector_score,
            ts_rank(tsv, plainto_tsquery('PostgreSQL'))
            * 0.5 + (1 - (embedding <=> '{str(query_embedding)}')) * 0.5
            AS hybrid_score
        FROM test_embedding_table
        WHERE tsv @@ plainto_tsquery('PostgreSQL') -- Filter for relevant text first
        ORDER BY hybrid_score DESC
        LIMIT 10
    )
    SELECT * FROM ranked_docs;
    """

cur.execute(hybrid_search_query)
results = cur.fetchall()

Discussion

Semantic search using transformer-based embedding models is very powerful and often enough for most cases. However, when data quality varies, or users are used to keyword searches, the hybrid search can provide more reliable results and keep users satisfied.

For example, if you have a legacy system that uses keyword search and you want to upgrade it to semantic search and add a chatbot feature. Even if the overall results improve, users might be confused because they are used to keyword searches. A user searching for “dog” might also expect to find “hot dog,” even though they do not have much in common from a semantic perspective.

See also

For more details on hybrid search concepts and implementation, see Recipe 3.7. That recipe provides a step-by-step explanation of how to implement hybrid search directly in Python.

Chapter 5. Retrieval

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 7th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the author at dmnkplzr@googlemail.com.

Improving the retrieval step is often one of the most effective ways to boost the overall performance of a RAG system.

There is no single right or wrong approach here. As with all other machine learning use cases, you need a solid toolkit to be creative. Often, the simplest solutions are most effective.

Figure 5-1 shows several techniques you can try when fine-tuning the retrieval step of your RAG system. For each of them, you can find suitable recipes in this chapter:

Table 5-1. Retrieval enhancement techniques

Technique
Description
Recipe Reference

Metadata Filtering

Using metadata to filter search results based on what you know about the user

Recipe 5.1

Multi-Query Retrieval

Creating multiple versions of the same prompt to find more relevant documents

Recipe 5.3

Query Routing System

Using a query routing system to identify the best data sources or tools to use to answer the question

[Link to Come]

Auto-Merging Retriever

Retrieving larger, more meaningful text chunks by grouping related data

Recipe 5.5

Sentence Window Retrieval

Including nearby sentences to add context to the retrieved text

Recipe 5.6

Hypothetical Document Embeddings (HyDE)

Generating hypothetical documents to improve search results

Recipe 5.2

Query Decomposition

Breaking down complex queries into simpler sub-queries

Recipe 5.8

Reranking

Reviewing and evaluating the relevance of the retrieved documents using LLMs

Recipe 5.7

Figure 5-1. An overview of techniques to enhance the retrieval process

In this chapter, I explain each concept separately, but in real-world RAG workflows, we often combine them. Figure 5-2 shows an example of a complete retrieval process, combining multiple steps and techniques:

Start by breaking down a complex question into simpler ones using query decomposition (see Recipe 5.8). For example, the question “Was Google’s revenue higher than Microsoft’s in 2022?” can be split into two:

“What was Google’s revenue in 2022?”

“What was Microsoft’s revenue in 2022?”

For each question, perform a search and apply metadata filtering (see Recipe 5.1) to focus on the most relevant parts of the dataset. You can use a query routing system (see Recipe 5.4) to direct each question to the best data collection. For example, use the dataset about Microsoft’s financials for one question and Google’s financials for the other.

Combine the results and use a reasoning step to create a single, unified answer. Before this, you can add a reranking step (see Recipe 5.7) to refine the retrieved documents and ensure the most relevant ones are prioritized.

Figure 5-2. How a multi-step retrieval process could look like

Think of the recipes in this chapter as building blocks you can combine to improve your RAG workflow and achieve better search results.

5.1 Optimizing Query Results using Metadata Filtering in PostgreSQL

Problem

You want to make searches faster and more accurate by filtering your vector database using metadata before performing the semantic search. This focuses the search on the truly relevant parts of the dataset.

Solution

Storing metadata with embeddings lets you filter searches to find the most relevant information. This improves speed, accuracy, and robustness. In PostgreSQL, a simple WHERE clause is sufficient.

Example: Imagine you’re building a chatbot for a sports website covering football, tennis, and basketball.

A question like “Who is the best player of all time?” is too vague without knowing which sport the user means.

By using what we know about the user, we can make searches more specific and accurate:

User Profile Example: Sarah, 28, tennis enthusiast from New York

When Sarah asks “Who is the best player of all time?” we can focus on tennis players like Serena Williams or Roger Federer

For the question “What were yesterday’s results?” we can prioritize tennis tournament scores over football or basketball games

By using what we know about the user, we can make searches more specific and answers more relevant. Figure 5-3 shows another example, where the query comes from Jim, a 25-year-old FC Liverpool fan.

Figure 5-3. How to use metadata filtering as part of an router retriever. The router prefilters the search space based on metadata before performing the search.

Prerequisites

All you need is a PostgreSQL database with the pgvector extension installed.

pip install psycopg2 openai

Steps

Let’s start with a simple example. Imagine a dataset containing general knowledge about football and tennis. We’ll use the example from the introduction of this recipe: the question “Who is the greatest player of all time?”. To answer this question effectively, we first need to identify the sport and filter our dataset accordingly.

We’ll create a table that includes metadata, then use this metadata to refine our search.

The code in Example 5-1 demonstrates how to create a table that stores both the embedding vectors and the metadata fields. We store all metadata in a dedicated “metadata” column, which holds a JSON object. In this example, we add a “topic” field to the JSON object to capture the category of each text chunk (see Example 5-1).

Example 5-1. How to create a table using metadata filtering

    # creates the vector extension if it does not exist
    cur.execute("""CREATE EXTENSION IF NOT EXISTS vector""")

    # Create table with metadata column
    cur.execute(
        """
        CREATE TABLE IF NOT EXISTS embeddings_table_with_metadata (
            id SERIAL PRIMARY KEY,
            text_chunk TEXT,
            embedding VECTOR(1536),
            metadata JSONB
        )
    """
    )
    conn.commit()

Next, we’ll populate the newly created table with sample text chunks. The code in Example 5-2 demonstrates how to insert sample data into our metadata-enabled table. For this example, we use just four text chunks for simplicity: two snippets about tennis and two about football.

Each row contains: - The text content of the chunk - Its embedding vector - A metadata JSON object with a “topic” field indicating the sport category

You can extend the metadata structure to include additional fields like publication date, author, source, or language of the text.

Example 5-2. How to populate the table with sample data and metadata

    # Define text chunks and metadata
    text_chunks = [
        {
            "text": "Roger Federer has won 20 Grand Slam titles in tennis.",
            "topic": "tennis",
        },
        {
            "text": "The FIFA World Cup is the most prestigious football tournament.",
            "topic": "football",
        },
        {
            "text": "Serena Williams is one of the greatest tennis players of all time.",
            "topic": "tennis",
        },
        {
            "text": "Lionel Messi has won multiple Ballon d'Or awards in football.",
            "topic": "football",
        },
    ]
    # Initialize OpenAI client
    client = OpenAI()

    # Insert text chunks with embeddings and metadata
    for chunk in text_chunks:
        # Get embedding using OpenAI API
        response = client.embeddings.create(
            input=chunk["text"],
            model="text-embedding-3-small"
        )
        embedding = response.data[0].embedding

        metadata = {"topic": chunk["topic"]}
        cur.execute(
            "INSERT INTO embeddings_table_with_metadata (text_chunk, embedding, metadata) "
            "VALUES (%s, %s, %s)",
            (chunk["text"], embedding, json.dumps(metadata)),
        )

    conn.commit()

After running the code above, your PostgreSQL table will contain four rows with the following structure: - Text chunks about tennis and football - Corresponding embedding vectors for semantic search - Metadata JSON objects categorizing each chunk by sport

Figure 5-4 shows how this data appears when viewed in pgAdmin, displaying the text content, embeddings, and metadata fields in their respective columns.

Figure 5-4. The populated metadata filtering table as viewed in pgAdmin

Now, let’s use the metadata to improve our search. In Example 5-3, the user asks, “Who is the best player?”. Since we know the user is a football fan, we filter the dataset to show only football-related content. This is done using a simple WHERE clause in the SQL query.

The SQL query uses two placeholders: * One for the embedding vector from the user’s question * One for the metadata filter (in this case, filtering for “topic” = “football”)

Example 5-3. How do we generate an appropriate SQL command to restrict the search space using metadata?

    # Query the table with metadata filtering
    query = "Who is the best player?"
    response = client.embeddings.create(
        input=chunk["text"],
        model="text-embedding-3-small"
    )
    query_embedding = response.data[0].embedding
    topic_filter = "football"

    cur.execute(
        f"""
        SELECT text_chunk, 1 - (embedding <=> %s::vector) AS similarity
        FROM embeddings_table_with_metadata
        WHERE metadata->>'topic' = %s
        ORDER BY similarity DESC
        LIMIT 5
        """,
        (query_embedding, topic_filter),
    )
    results = cur.fetchall()

The query returns football-related results, with “Lionel Messi has won multiple Ballon d’Or awards …​” being the most relevant match (see Figure 5-5).

Figure 5-5. Query results showing metadata filtering in action

Tip

Use an LLM to automatically choose the right metadata filters. Add a quick LLM call before your search to decide which part of the database to look in. This makes searches more efficient and often cheaper by ensuring you only get relevant content. Use a fast, cheap model like GPT-4.1 mini for choosing filters, and save the powerful models like GPT-4.1 for the final answer generation.

Discussion

Gathering the right metadata might take some work upfront, but using it to limit the search space is extremely effective. It makes searches more accurate, cheaper, and your whole RAG system easier to test.

Tip

For any RAG implementation, start small. For example, build a chatbot specifically tailored for football enthusiasts. Once it performs well, you can gradually expand its scope to include tennis, basketball, and other sports. I recommend keeping the datasets separate by using appropriate metadata. This makes it easier to test and maintain your solution and knowledge base later on.

See also

Most vector stores support metadata filtering with straighforward commands.

PostgreSQL uses standard SQL WHERE clauses for filtering. You can store metadata in a JSONB column or separate columns - both work well.

Other vector stores like ChromaDB and Pinecone use similar operators like $eq (equals), $gt (greater than), and $lt (less than).

Note

Using metadata in RAG workflows is easy once it’s stored. The hard part is collecting the right metadata upfront. Start by gathering more metadata than you think you need, then remove what’s unnecessary later. Recipe 2.1 shows how to collect and generate useful metadata.

5.2 Enhancing Search Results by Extending the Original Query with Generated Pseudo-Documents

Problem

You want to improve search results by creating pseudo-documents that match the style of real content in your database.

Solution

The HyDE (Hypothetical Document Embeddings) approach uses an LLM to generate fake documents that help find better search results.

Here’s how it works: Instead of searching directly with the user’s question, we first ask an LLM to write a hypothetical answer. Even though this answer might be completely made up, it’s written in a style similar to real documents in your database.

We then use this fake document to search the vector database. This often finds more relevant real documents because the search is based on content that matches your data’s style and structure.

The process is simple:

User asks a question

LLM generates a hypothetical document that might answer it

Search the vector database using this fake document’s embeddings

Then use the retrieved real documents to generate the final answer.

Figure 5-6 explains the process and includes two more examples I found in the paper Precise Zero-Shot Dense Retrieval without Relevance Labels.

Figure 5-6. How does the LLM generate a hypothetical paragraph to assist in answering the question?

Another example could be:

“What is the revenue of Company X in 2022?”

The LLM might generate:

“Company X reported approximately $5 billion in revenue for 2022, driven by growth in its technology and services divisions.”

Even though this information is made up, it sounds like a real financial report. This helps our search find similar content in the database.

Prerequisites

You need an LLM that can generate hypothetical documents. For this example, I’m using GPT-4o mini because it’s fast, cost-effective, and capable of generating meaningful hypothetical content.

pip install openai pydantic

Steps

In Example 5-4, we use the LLM to generate hypothetical documents for the sample query “What is the revenue of Company X in 2024?”. We use a Pydantic model to ensure we get a list of hypothetical documents back. Each document is a string.

Example 5-4. How to generate hypothetical documents

    from pydantic import BaseModel
    from openai import OpenAI

    user_query = "What is the revenue of Company X in 2024?"

    client = OpenAI()

    class HypotheticalDocuments(BaseModel):
        documents: list[str]

    prompt = f"""
    You are an AI assistant. Based on the user query below, generate three hypothetical
    text chunks that contain relevant information to answer the query.
    """

    completion = client.responses.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": prompt},
            {"role": "user", "content": user_query},
        ],
        text_format=HypotheticalDocuments,
    )

    hypothetical_documents = completion.output_parsed.documents

Figure 5-7 shows the generated hypothetical documents. The LLM creates paragraphs that look like financial reports, which help find better results in the vector database.

Figure 5-7. Hypothetical documents generated to answer “What is the revenue of Company X in 2024?”

Now use the generated hypothetical documents for the actual vector search.

Note

This concept is similar to multi-query retrieval (Recipe 5.3). With multi-query, we create multiple versions of the question and search with each one. With HyDE, we often generate multiple hypothetical documents and search with those instead.

Discussion

HyDE improves retrieval accuracy by helping find the most relevant documents in your vector database. However, this approach has trade-offs: it’s slower and more expensive because of the extra generation step, and you’ll run multiple queries (for the original question plus all generated documents). You then need to merge these results effectively.

Rather than just hoping for better quality, you should measure it. [Link to Come] shows how to evaluate your retrieval accuracy and overall RAG system performance.

See also

HyDE was first introduced in a research paper by the University of Waterloo. You can learn more about the concept in the original paper: Precise Zero-Shot Dense Retrieval without Relevance Labels.

5.3 Improving Search Results with Multi-Query Retrieval

Problem

You want to improve search results when a single query might miss relevant information.

Solution

Multi-query retrieval creates several variations of your original question, searches with each one, then combines all the results for a more comprehensive answer.

The process is straightforward: 1. Take the original query and generate 2-3 alternative phrasings 2. Search the vector database with each query version 3. Merge the retrieved documents and send them to the LLM

Figure 5-8 illustrates this workflow. Each query variant finds different relevant documents, giving you broader coverage of your knowledge base.

Figure 5-8. How multi-query retrieval works

Example: A user asks “What are the benefits of renewable energy?”

Multi-query retrieval generates these variations:

“What are the environmental benefits of renewable energy?”

“How does renewable energy contribute to economic growth?”

“What are the drawbacks of renewable energy?”

Each query targets different aspects: environmental impact, economic effects, and potential challenges. Together, they retrieve a more comprehensive set of documents than the original question alone.

Prerequisites

You just need an LLM that can generate hypothetical documents based on a user query. For this example, I’m using GPT-4o mini because it’s fast, cost-effective, and more than capable of generating meaningful query variations. By using smaller models, we keep the overall latency of the RAG system low.

pip install openai pydantic

Steps

Example 5-5 shows you everything you need. It includes a prompt template that instructs the LLM to rephrase the original query into several subqueries. Since we have defined a response template in the form of a Pydantic model, the response is a list of strings—our subqueries.

Example 5-5. How to implement a basic multi query retriever

    from openai import OpenAI
    from pydantic import BaseModel
    import os

    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    question = "What are the benefits of renewable energy?"

    query_prompt = f"""You are an AI language model assistant. Your task is
    to create three alternative versions of the provided user query to
    enhance the retrieval of relevant documents from a vector database.
    By offering diverse variations of the query, your goal is to help
    mitigate the limitations of distance-based similarity search. Provide
    these alternative queries, each on a new line.

    Original query: {question}
    """

    # send the query prompt to OpenAI
    class QueryVariations(BaseModel):
        queries: list[str]

    completion = client.responses.parse(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": query_prompt,
            },
        ],
        text_format=QueryVariations,
    )

    queries = completion.output_parsed.queries

This gives us three alternative versions of the same question (see Figure 5-9):

Figure 5-9. Screenshot: How the generated queries look like

For all three queries (plus the original question), we perform a similarity search. The only remaining task is to decide how to merge the retrieved documents into a single set. You can simply combine all results, or use reranking methods to prioritize the most relevant ones.

For example, you could rerank all retrieved documents using an LLM. Recipe 5.7 shows how this could look. Afterwards, you can decide whether to use all results or only the top-k most relevant ones.

Discussion

Multi-query retrieval uses a simple but powerful idea: create different versions of a query, search with each one, and combine the results. This approach is similar to how ensemble models work in machine learning.

Think of it like a random forest, which combines many simple decision trees. Each tree makes predictions independently, and while individual trees might not be very accurate, combining hundreds of them creates robust and accurate predictions. This mirrors how groups of people often make better decisions than individuals alone.

With multi-query retrieval, we use slightly rephrased questions to target different areas of our knowledge base. Together, these variations provide stronger overall retrieval results than a single query alone.

See also

LangChain offers a built-in feature for multi-query retrieval called the MultiQueryRetriever. However, I recommend building your own custom multi-query retriever. The concept is relatively simple, and this allows you to better tailor the retriever to your specific datasets and requirements.

5.4 Addressing Complex Requests by Designing a Query Routing System

Problem

You want to handle complex questions by directing them to the right data sources or tools.

Solution

A query routing system directs questions to the right data sources or tools. For example, “What is 1+8?” goes to a calculator, while “What was Google’s revenue?” goes to a financial database.

In the specific example in Figure 5-10, the user is asking for a comparison between the revenue of Google and Microsoft. Let’s assume we have a vector database filled with the financial reports of various companies. To ensure that our RAG system doesn’t mix data from different companies, we can break down the question into “What was Google’s revenue?” and “What was Microsoft’s revenue?” separately, and then interpret the retrieved documents together.

Figure 5-10. Query routing with function calling models. Here, the agent can query either the vector database or the SQL database. The agent should decide which route is most appropriate.

Since knowledge pools often contain similar but unrelated information, we can improve efficiency by keeping databases separate. The agent first decides which database section is most relevant for the query.

Tip

The idea of query routing is very similar to what we do in Agentic RAG systems, only at scale. In [Link to Come] you can find a deep dive on how to define agents that make autonomous decisions and have a toolset they can use to answer questions.

Prerequisites

We’ll create our own router represented by an LLM and suitable instructions. By providing the right prompt, the model will decide which tool or function to use.

Each tool is typically represented as a function. This function could retrieve data from a specific collection, execute a database query, or even perform calculations in Python, such as mathematical operations.

Figure 5-11. How to create a custom routing system for function calling

To build a router, we need the following packages installed for this recipe:

pip install openai pydantic requests

Steps

In Example 5-6, we create a prompt that instructs the model to choose between three different tools. In this case, we simulate having three different data sources that we can access using various tools:

A knowledge pool with information about general football knowledge, such as “What is the offside rule?”

A knowledge pool with information about general tennis knowledge, such as “How many points does a player typically need to win a tennis game?”

Additionally, we simulate a connection to a SQL database containing the latest football results. Based on the user’s question, the router should assign one of the following labels: “general_football_knowledge”, “general_tennis_knowledge”, or “latest_football_results_sql”.

The function that returns the user questions as a list of dictionaries. Each dictionary has the user’s query and a selected_data_source field that gets filled with the correct label based on the router’s decision.

Example 5-6. Implementing a basic query router

    from pydantic import BaseModel
    from openai import OpenAI

    client = OpenAI()

    user_queries = [
        {
            "query": "Who is the all-time top scorer in the FIFA World Cup?",
            "selected_data_source": None,
        },
        {
            "query": "What are the four Grand Slam tennis tournaments?",
            "selected_data_source": None,
        },
        {
            "query": "Did Manchester United win their last game?",
            "selected_data_source": None,
        },
    ]

    prompt = f"""
    You are an exptert at routing a user question to the appropraite data source.
    Given a user question choose which of the data sources in list_of_data_sources
    is the best to answer the question.
    """

    from typing import Literal
    from pydantic import Field

    class RouterDecision(BaseModel):
        data_source: Literal[
            "general_football_knowledge",
            "general_tennis_knowledge",
            "latest_football_results_sql",
        ] = Field(
            ..., description="The best data source to answer the user's question."
        )

    for user_query in user_queries:
        completion = client.responses.parse(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": prompt},
                {"role": "user", "content": user_query["query"]},
            ],
            text_format=RouterDecision,
        )
        user_query["selected_data_source"] = completion.choices[
            0
        ].message.parsed.data_source

Figure 5-12 shows the results for the three sample questions.

Figure 5-12. Screenshot: Query routing system results

I’ve added comments showing how the LLM chose between functions. The LLM’s reasoning seems correct. With enough context, LLMs can reliably decide which tools are most suitable.

Discussion

In this recipe, we used an LLM to decide which tool is best suited to answer the user’s question. However, there are other methods that can be faster and more efficient. Figure 5-13 shows two examples:

Vector Search with Majority Voting: Perform a similarity search across all your data and examine the top results (e.g., the top 30). Count how often results from each knowledge pool appear and select the “winner”—the knowledge pool with the most results. This works best with balanced datasets where each topic has roughly the same number of entries.

Embedding Classifier: Train a classifier using embedding vectors as input and the knowledge pool labels as the target variable, we want to predict. When a new query arrives, generate its embedding and use the classifier model to predict the most relevant knowledge pool. See Recipe 3.6 for training details.

Figure 5-13. Alternative approaches to build a query routing system

Tip

You don’t need to build a router from scratch. Frameworks like LangChain, LlamaIndex, Amazon Bedrock Agents, and Azure AI Agent Service offer ready-made solutions. However, I recommend keeping it simple - a router is often just an advanced prompt, so avoid adding unnecessary frameworks to your project.

See also

You can also find router functions documented in LangChain and LlamaIndex. I personally don’t use them and prefer to write my own router definitions, but the packages are a great source of inspiration. Here you can find the details:

LlamaIndex: Using the Router Query Engine

LangChain: Routing with Function Calling Models

5.5 Increasing Search Efficiency by Designing an Auto-Merging Retriever (aka Parent Document Retriever)

Problem

You want to give the LLM larger, more complete text chunks instead of small pieces. This provides better context and improves answer quality while keeping search accuracy high.

Solution

To use auto-merging retrieval, organize your text into smaller chunks that link to larger parent chunks. When the system finds enough smaller chunks from the same parent during retrieval, it provides the entire parent chunk to the LLM for better context.

Figure 5-14 illustrates a three-level tree structure: a parent node, two child nodes, and six leaf nodes. The system compares the query against the leaf nodes during semantic search.

For instance, if leaf nodes 1, 2, and 5 are relevant and two of them belong to child node 1, the system retrieves the complete text from child node 1. This gives the LLM more context.

Figure 5-14. How the auto-merging retriever works: parent and child nodes

The key is to connect parent chunks to their smaller child chunks. This way, instead of getting only small pieces of information, the system can retrieve larger, more complete information pieces.

Prerequisites

You need an embedding model and a vector store to save text chunk embeddings. This example uses text-embedding-3-small and PostgreSQL with pgvector. In this recipe, we store the child node chunks and parent node chunks along with the embeddings for the child nodes in a PostgreSQL database, since we want to perform similarity search on top of it later.

You don’t necessarily need PostgreSQL; you can store everything in another vector database. However, if you want to rebuild this in PostgreSQL and haven’t installed SQL yet, you can find everything about the setup in Recipe 4.4. To connect to the PostgreSQL database from Python we use the psycopg2 package.

pip install psycopg2 openai

Steps

First, create a simple structure with parent and child nodes. The child nodes are small text chunks, and parent nodes are larger chunks that contain multiple child chunks.

The auto-merging retriever works like this:

Search for relevant child chunks using similarity search

If multiple child chunks from the same parent are found, return the entire parent chunk instead

This gives the LLM more complete context while keeping search accuracy high.

In Example 5-7, we create a simple two-level structure. Each child chunk contains 250 characters, and every 4 child chunks are grouped into one parent chunk (1000 characters total). In the code snippet, we iterate over the smaller text snippets that are as long as we defined the leaf node size. We give each leaf node an ID, and when we have 4 leaf nodes together, we define a new parent node with a parent node ID. We write the text chunk of each individual leaf node plus the leaf node ID and the corresponding parent node ID into a dictionary.

The subsequent retrieval process is straightforward: if we find enough relevant child chunks from the same parent (say, 3 out of 4), we replace them with the full parent chunk to provide better context.

Example 5-7. Test data for auto-merging retriever

    file_path = "../datasets/text_files/random-text-about-5-different-stories-with-pargraphs.txt"

    with open(file_path, "r", encoding="utf-8") as f:
        text = f.read()

    leaf_node_size = 250  # size of the leaf nodes in characters
    parent_merge = 4  # number of leaf nodes to merge into a parent node
    parent_node_size = (
        parent_merge * leaf_node_size
    )  # size of the parent nodes in characters
    text_chunks = []  # list to store the text chunk dictionaries

    for leaf_node_id in range(0, len(text) // leaf_node_size, 1):
        parent_node_id = leaf_node_id // parent_merge
        leaf_chunk_start = leaf_node_id * leaf_node_size
        parent_chunk_start = parent_node_id * parent_node_size

        text_chunk = {
            "leaf_node_id": leaf_node_id,
            "leaf_chunk": text[leaf_chunk_start : leaf_chunk_start + leaf_node_size],
            "parent_node_id": parent_node_id,
            "parent_chunk": text[
                parent_chunk_start : parent_chunk_start + parent_node_size
            ],
        }
        text_chunks.append(text_chunk)

Next, we create a PostgreSQL table to store the leaf node embeddings, their IDs, parent node IDs, and the corresponding text chunks. Example 5-8 shows the SQL command for this table, which includes columns for leaf and parent IDs, their text, and the embedding vector.

Example 5-8. Creating a suitable table for the auto-merging retriever

    # Create a new table for storing chunks and embeddings
    cur.execute("""CREATE EXTENSION IF NOT EXISTS vector""")

    cur.execute(
        """
        CREATE TABLE auto_merging_retriever_text_chunks (
            leaf_node_id SERIAL PRIMARY KEY,
            leaf_chunk TEXT,
            parent_node_id INTEGER,
            parent_chunk TEXT,
            leaf_chunk_embedding VECTOR(1536)
        )
    """
    )
    conn.commit()

Now fill the table with each leaf node’s embedding, leaf node ID, parent node ID, and parent chunk. To generate the embeddings for the leaf nodes, we again use OpenAI’s text-embedding-3-small model.

Example 5-9. Inserting data for auto-merging retriever

    from openai import OpenAI

    client = OpenAI()  # initialize OpenAI client

    # Insert chunks and their embeddings into the table
    for text_chunk in text_chunks:
        leaf_embedding = (
            client.embeddings.create(
                input=[text_chunk["leaf_chunk"]], model="text-embedding-3-small"
            )
            .data[0]
            .embedding
        )

        cur.execute(
            """
            INSERT INTO auto_merging_retriever_text_chunks
            (leaf_node_id, leaf_chunk, parent_node_id, parent_chunk, leaf_chunk_embedding)
            VALUES (%s, %s, %s, %s, %s)
            """,
            (
                text_chunk["leaf_node_id"],
                text_chunk["leaf_chunk"],
                text_chunk["parent_node_id"],
                text_chunk["parent_chunk"],
                leaf_embedding,
            ),
        )
    conn.commit()

After performing a similarity search, we check which parent nodes the retrieved leaf nodes belong to. If enough leaf nodes from the same parent are found (for example, 3 out of 4), we return the entire parent chunk instead of just the individual leaf nodes. This gives the LLM more context and helps it generate better answers.

Discussion

In practice, auto-merging retrieval can be very effective, but it has trade-offs. If irrelevant child chunks are selected, including the whole parent chunk may add unnecessary content and noise to your prompt. Table 5-2 summarizes the main pros and cons.

Table 5-2. Pros and Cons of Auto-Merging Retrieval

Pros
Cons

Enhanced Contextual Understanding: Provides the LLM with broader context, leading to more complete and coherent responses for complex queries.

Increased System Complexity: Requires managing parent-child relationships, additional storage, and more complex setup.

Improved Retrieval Precision: Initial search uses small “child” chunks, which have focused semantic meanings, resulting in highly accurate and precise retrieval of relevant information.

Higher Computational Costs: Passing larger parent documents to the LLM increases token usage, leading to longer response times and higher API costs.

See also

Almost all orchestration frameworks, such as LlamaIndex, Haystack, and others, have their own definition and implementation of an AutoMergingRetriever. Here are some resources for more information:

AutoMergingRetriever from LlamaIndex: In LlamaIndex, the Auto-merging retriever works in conjunction with the HierarchicalNodeParser, which automatically creates a multi-level hierarchy of nodes (e.g., with default chunk sizes of 2048, 512, and 128 tokens). Each child node explicitly maintains a reference to its larger parent node. The retriever then initially searches for relevant leaf nodes via a vector store and, if multiple retrieved leaf nodes belong to the same parent and meet a specified threshold, it returns the larger parent node to the LLM. For more details, visit the LlamaIndex documentation.

AutoMergingRetriever from Haystack: Haystack utilizes its HierarchicalDocumentSplitter to organize documents into parent and leaf structures, storing them in separate document stores. The AutoMergingRetriever is then integrated into the retrieval pipeline. After an initial retriever fetches relevant leaf documents, the AutoMergingRetriever decides whether to return these individual leaf documents or their larger parent document. This decision is based on a threshold parameter; for example, a threshold of 0.5 means if at least 50% of the leaf documents below a common parent match the query, the parent document will be returned. Find out more at the Haystack blog.

5.6 Increasing Search Results by Designing a Sentence Window Retriever

Problem

You want to improve retrieval by including more context around the most relevant text chunks, helping the LLM give more meaningful and coherent answers.

Solution

The Sentence Window Retriever improves search results by including a few sentences before and after the most relevant chunk. This gives extra context, making the retrieved text more complete and easier to understand.

Figure 5-15 shows an example with a question about FC Bayern Munich. Instead of returning just one sentence, the retriever adds nearby sentences, so the answer covers the club’s history and achievements more clearly.

Figure 5-15. How sentence window retrieval works

To create a sentence retriever, ensure your text chunks are numbered in the database. This lets you easily add surrounding text before and after a selected chunk (see Figure 5-16).

Figure 5-16. How the sentence window retriever works in practice

Prerequisites

You need a vector store that allows you to store text chunks with an ID. Ich nutze hier PostgreSQL und die pgvector extension. To later use the order in which the text chunks appear in the document, you need to store both a document ID and a sequential ID for each text chunk.

Um mich mit Python mit meiner PostgreSQL datenbank zu verbinden nutzt ich das psycopg2 package.

pip install psycopg2

Steps

The concept is straightforward, so I’ll just show how the table would look and skip the data insertion.

Example 5-10 shows how to create a simple embedding table with an extra column for the chunk order. In this example, we only need a chunk ID, but in real-world cases, you’ll also want a document ID to track which document each chunk comes from.

Example 5-10. How to create a table using sentence window retrieval

    cur.execute(
        """
        CREATE TABLE sentence_window_retriever_text_chunks (
            chunk_id SERIAL PRIMARY KEY,
            chunk TEXT,
            chunk_embedding VECTOR(1536)
        )
    """
    )
    conn.commit()

To fill the table, split the text, generate embeddings, and insert them as in previous recipes. The important part is that chunk IDs are sequential, showing the order of text in the document.

When a user asks a question and a relevant chunk is found, nearby chunks may also be important for context. By including surrounding chunks, you make the retrieved information more complete and coherent.

Discussion

The Sentence Window Retriever helps keep text chunks short and focused, improving retrieval quality. It also adds nearby context, so important information isn’t missed if it’s split across chunks.

However, including extra context makes prompts longer. You should test if the concept of the Sentence Window Retriever improves your system’s performance without adding too much cost or delay. You find some evaluation metrics that can help you measure this in [Link to Come].

See also

If you’re already using LangChain for your RAG setup, take a look at their SentenceWindowNodeParser.

5.7 Improving Search Accuracy with Reranking Methods

Problem

You want to keep only the most relevant documents and avoid adding unnecessary information to the prompt by reranking the retrieved results.

Solution

Vector search and embeddings generation are the steps in the RAG flow which are relatively fast, so we can run it multiple times without much delay. Techniques like HyDE (see Recipe 5.2) or multi-query retrieval (see Recipe 5.3) do this by searching with different versions of the question or with generated documents. But after multiple searches, we get a large set of results—how do we pick the best ones?

Reranking helps by sorting the retrieved documents so only the most relevant ones are sent to the LLM. Here’s how it works:

Perform a semantic search to find the top matching documents.

Use a reranking model to review and sort these documents by relevance.

The reranker sorts the documents by relevance (see Figure 5-17). You can then select just the top results to send to your generation step.

Figure 5-17. How to apply reranking to the retrieved documents

Note

Although this recipe mentions semantic vector search, the same idea applies to other retrieval methods. Querying a SQL database or REST endpoint is usually much faster than LLM generation, so you can run multiple queries quickly without slowing down the user experience.

Prerequisites

In this recipe we are using LLM-based rerankers, so the only thing we need is an LLM that has a look at the retrieved documents and assesses their relevance, and if necessary, reranks them. For this example, I am using GPT-4o mini because it is fast, cost-effective, and more than capable of reranking documents based on their relevance to the query.

pip install openai pydantic

Steps

The LLM assesses the relevance of each document. We provide the model with all retrieved documents in a prompt, along with instructions to rerank them. Example 5-11 shows a basic example of a prompt that instructs the LLM to rerank the documents.

Example 5-11. How to implement a basic reranking function

    import textwrap
    from openai import OpenAI
    from pydantic import BaseModel

    text_chunks = {
        1: "Tesla's Supercharger network and tech lead face rising competition from BYD "
        "and established automakers.",
        2: "Tesla's production grows, but price competition threatens its market share.",
        3: "The automotive industry is shifting to EVs due to climate change and "
        "regulations.",
        4: "Semiconductor shortages are disrupting automotive supply chains.",
        5: "Consumer demand for autonomous driving and advanced tech impacts EV "
        "competition.",
    }

    prompt = textwrap.dedent(
        f""""
        Query: Will Tesla remain the market leader in electric vehicles?

        Documents:
            1. {text_chunks[1]}
            2. {text_chunks[2]}
            3. {text_chunks[3]}
            4. {text_chunks[4]}
            5. {text_chunks[5]}

        Instructions:
        Please assess the relevance of each document to the query and
        provide a relevance score from 1 to 5, where 5 is the most relevant.#

        Relevance Scores:
        """
    )

    client = OpenAI()  # initialize OpenAI client

    class RankedDocuments(BaseModel):
        text_chunk_ids: list[int]
        relevance_scores: list[int]

    completion = client.responses.parse(
        model="gpt-4o",
        input=[
            {"role": "user", "content": prompt},
        ],
        text_format=RankedDocuments,
    )

    ranked_docs = completion.output_parsed

Figure 5-18 shows the new order of relevance after reranking.

Figure 5-18. How the reranking results look like

Note

We used LLM-based rerankers, but there are other options. For example, ColBERT combines vector search and cross-encoders for fast and accurate reranking. LlamaIndex provides a ColbertRerank function—see the LlamaIndex documentation for details.

Discussion

Semantic search is very fast, quickly finding similar embedding vectors. LLMs can analyze text more deeply, but are slower and more expensive. By adding an LLM-based reranking step after semantic search, you get both speed and better relevance.

Using an LLM reranker adds some extra latency and cost. However, if it helps you reduce the number of irrelevant chunks sent to the LLM, it can be worth it in terms of latency and model costs.

See also

You can find suitable implementations in LangChain, LlamaIndex, and Haystack. If you want to further refine your reranker, you will certainly find some inspiration here:

LangChain’s RankLLM Reranker: A versatile LLM-based reranker module.

LlamaIndex’s ColbertRerank: Utilizes the ColBERT model for fast and accurate reranking, acting as a Node Postprocessor.

LlamaIndex’s SentenceTransformerRerank: Improves LLM accuracy by pruning irrelevant nodes using cross-encoder models, also integrated as a Node Postprocessor.

Haystack’s SentenceTransformersSimilarityRanker: Ranks documents based on query similarity using Hugging Face cross-encoder models.

5.8 Decomposing Complex Queries into Multiple Sub-Queries

Problem

You want to answer complex questions by breaking them down into simpler, manageable parts.

Solution

Figure 5-19 shows the entire process of decomposing complex queries into simpler sub-queries in four steps:

Decomposition: Break down the complex query into simpler sub-queries.

Retrieval: For each sub-query, follow the standard RAG process: generate embeddings, perform a similarity search, and retrieve relevant documents.

Answer Generation: Generate answers for each sub-query, resulting in a collection of question-answer pairs.

Reasoning: Combine the answers to the sub-queries to form a final response to the original complex query.

Figure 5-19. How to decompose complex queries

Prerequisites

For this recipe, we need an LLM that can break down complex questions into simpler ones. Here I intentionally use the larger model GPT-4o because it’s important that the model actually understands all facets of the original question.

pip install openai pydantic

Steps

First, in Example 5-12 we define a Pydantic model that captures the question-answer pairs. In the first step, I only populate the questions, leaving the answers empty. Later, I use my RAG system to answer each individual question step by step and store them in the same dictionary.

In addition to the Pydantic model, you’ll find a suitable prompt that instructs the LLM to break down the complex original question into simpler ones. If the question is very simple, the model should keep it as a single question.

Example 5-12. How to decompose complex queries using LLMs

    from pydantic import BaseModel
    from typing import Optional
    from openai import OpenAI

    class Question(BaseModel):
        question: str
        answer: Optional[str] = None

    class Questions(BaseModel):
        questions: list[Question]

    splitter_prompt = """
    You are a helpful assistant for a RAG chatbot.

    Your job is to break down complex questions into simpler ones that are easy to
    answer. When the answers to these simpler questions are combined, they should
    fully answer the original question. If the question is already simple, leave
    it as it is. Handle one question at a time.

    Example:
        1. Query: Did Microsoft or Google make more money last year?

    Decomposed Questions:
       1. How much profit did Microsoft make last year?
       2. How much profit did Google make last year?
    """

We now use this prompt template with a sample question “What are the benefits of renewable energy compared to fossil fuels?” (see Example 5-13). The question may not seem very complex at first, but if we as humans had to make a truly educated statement, we would first research the advantages and disadvantages of renewable energy and then the pros and cons of fossil fuels, and then come to a comprehensive answer ourselves.

That’s exactly what we’re trying to imitate by instructing the model do break down the question first in sub-questions.

Example 5-13. How to decompose complex queries using LLMs

    client = OpenAI()

    query = "What are the benefits of renewable energy compared to fossil fuels?"

    completion = client.responses.parse(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": splitter_prompt},
            {"role": "user", "content": query},
        ],
        text_format=Questions,
    )

    decomposed_questions = completion.output_parsed.questions

Figure 5-20 shows the results of the decomposition. The LLM breaks down the original question “What are the benefits of renewable energy compared to fossil fuels?” into simpler sub-questions, like “What are the benefits of renewable energy?” and “What are the benefits of using fossil fuels?”.

Figure 5-20. How the decomposed sub-queries look like

The following steps are the standard RAG flow, but we execute it separately for each sub-query and add a final consolidation step for all generated answers.

We iterate over each question, convert it into embeddings using an embedding model (e.g., text-embedding-3-small), query the vector database, and generate an answer using an LLM (e.g., GPT-4o). This time, we perform this process for each individual question.

This gives us a list of question-answer pairs. These answers should make it much easier for the LLM to answer the original, complex question.

Discussion

Breaking down complex questions into smaller, simpler parts makes them easier to solve. It’s like solving a problem step by step, learning and improving as you go. This approach is highly effective, and when used correctly, you can achieve better results with smaller, more cost-efficient LLMs by running multiple prompts in sequence instead of sending a single prompt to a large, expensive LLM.

Note

This is also the core principle behind the deep research capabilities of OpenAI, Anthropic, and Google. Typically, these systems use smaller, faster models that, like humans, search the internet, interpret results, and follow a research plan—very similar to how we would browse the web and gather important information.

You can combine query decomposition with query routing (Recipe 5.4) to split complex questions and direct each part to the best data source. Use query decomposition only if it improves your specific use case. Like other techniques in this chapter, it increases system complexity, implementation effort, token usage, and response time. Table 5-3 shows the main trade-offs.

Table 5-3. Pros and Cons of Query Decomposition

Pros
Cons

Enhanced Accuracy and Relevance: Breaks complex queries into focused sub-queries, leading to more precise and relevant document retrieval, addressing ambiguity.

Higher Latency: Multiple retrieval and generation steps increase response time, impacting real-time applications.

Reduced Hallucination: Grounds each sub-output in verifiable, retrieved documents, strengthening the factual accuracy and trustworthiness of the final response.

Increased Development Effort: Requires sophisticated orchestration and complex debugging for multi-step pipelines.

See also

You can find existing implementations of query decomposition in frameworks like LangChain, LlamaIndex, and Haystack. These frameworks often combine query decomposition with other techniques like HyDE, multi-query retrieval, and query routing to improve retrieval accuracy.

For more details on using query decomposition in RAG workflows, see Tuanwa Celik’s article on the Haystack blog: Haystack blog.

Tip

There are no limits to creativity here; the techniques presented in this chapter are simply meant to provide you with a useful toolkit.