Combining Lexical and Embedding-Based Retrieval in RAG Systems

We are now in the fourth and final lesson of this course on Beyond Basic RAG: Improving Our Pipeline! Up to this point, we have explored ways to enhance Retrieval-Augmented Generation (RAG) systems by refining chunking strategies and leveraging advanced retrieval methods. In this lesson, you will learn how to merge a lexical-based retrieval approach (using Okapi BM25) with your existing embedding-based retrieval mechanism, creating a powerful hybrid retrieval pipeline.

By the end of this lesson, you should be able to:

1. Grasp the intuition behind Okapi BM25 for lexical retrieval.
2. Construct a BM25 index on your corpus.
3. Combine BM25 scores with embedding-based retrieval scores using a configurable weight parameter, alpha.

Let’s start by understanding what Okapi BM25 is and why it’s useful for retrieval.

Within the category of lexical-based search methods, Okapi BM25 is a popular choice. It focuses on the presence of specific keywords, rewarding relevant chunks that contain more occurrences of the query terms. At the same time, it avoids overemphasizing repeated words by incorporating a saturation effect.

A few core ideas behind BM25:

• Term Frequency (TF): More keyword matches in a chunk can signal higher relevance.
• Document Length Normalization: BM25 accounts for chunk length, ensuring that very long chunks with many repeated words are not unfairly scored.

BM25's handling of term saturation via the k1 parameter (not shown here) limits the score contribution of repeated terms, preventing documents from being unfairly ranked just because they contain many repetitions of the same keyword. Even though this parameter isn't directly configured in the example, its default behavior is important when interpreting why a document with many keyword repetitions might not rank as high as expected.

Although the underlying formula has several parameters and normalizations, the general purpose is straightforward: favor chunks containing the search terms, but don't let them dominate purely by repeating keywords.

Now that you have a sense of what BM25 does, let’s see how to build a BM25 index for your chunked corpus.

To create a BM25 index from your chunked corpus, you can use the Bm25Index struct. It precomputes sparse embeddings from your text, allowing for efficient scoring of queries against your document chunks.

pub struct Bm25Index {
    embedder: Embedder,
    doc_embeddings: Vec<Embedding>,
}

impl Bm25Index {
    pub fn new(chunks: &[Chunk]) -> Self {
        let corpus: Vec<&str> = chunks.iter().map(|c| c.text.as_str()).collect();
        let embedder = EmbedderBuilder::with_fit_to_corpus(Language::English, &corpus).build();
        let doc_embeddings = corpus.iter().map(|&text| embedder.embed(text)).collect();
        Bm25Index { embedder, doc_embeddings }
    }

    pub fn score(&self, query: &str) -> Vec<f32> {
        let q_emb = self.embedder.embed(query);
        self.doc_embeddings.iter().map(|doc_emb| dot(&q_emb, doc_emb)).collect()
    }
}

In this snippet, the Bm25Index struct holds both the embedder and the precomputed document embeddings. The new function builds the index from your chunks, and the score function computes similarity scores between a query and each chunk using a dot product over sparse vectors. This approach allows you to efficiently retrieve chunks that share keywords with the query.

With a BM25 index in place, the next step is to see how we can combine its scores with those from an embedding-based retrieval system.