Hybrid Retrieval in Retrieval-Augmented Generation Systems

Introduction

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.

Understanding the Okapi BM25 Algorithm

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.

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.

Building a BM25 Index

Here is a simple function that builds a BM25 index from your chunked corpus. We assume you already have a collection of text chunks ready.

import java.util.List;
import java.util.stream.Collectors;

public class BM25Index {
    private List<String[]> corpus;
    private float avgDocLength;
    private float k1 = 1.5f;
    private float b = 0.75f;

    public BM25Index(List<String[]> corpus) {
        this.corpus = corpus;
        calculateAvgDocLength();
    }

    private void calculateAvgDocLength() {
        float totalLength = 0;
        for (String[] doc : corpus) {
            totalLength += doc.length;
        }
        this.avgDocLength = totalLength / corpus.size();
    }

    public float[] getScores(String[] query) {
        float[] scores = new float[corpus.size()];
        for (int i = 0; i < corpus.size(); i++) {
            float score = 0;
            String[] doc = corpus.get(i);
            for (String term : query) {
                int termFreq = countOccurrences(doc, term);
                if (termFreq > 0) {
                    float idf = calculateIDF(term);
                    float numerator = termFreq * (k1 + 1);
                    float denominator = termFreq + k1 * (1 - b + b * doc.length / avgDocLength);
                    score += idf * numerator / denominator;
                }
            }
            scores[i] = score;
        }
        return scores;
    }

    private int countOccurrences(String[] doc, String term) {
        int count = 0;
        for (String word : doc) {
            if (word.equals(term)) {
                count++;
            }
        }
        return count;
    }

    private float calculateIDF(String term) {
        int docsWithTerm = 0;
        for (String[] doc : corpus) {
            if (Arrays.asList(doc).contains(term)) {
                docsWithTerm++;
            }
        }
        return (float) Math.log(1 + (corpus.size() - docsWithTerm + 0.5) / (docsWithTerm + 0.5));
    }
}

public static BM25Index buildBM25Index(List<Map<String, Object>> chunks) {
    List<String[]> corpus = chunks.stream()
            .map(chunk -> chunk.get("text").toString().toLowerCase().split("\\s+"))
            .collect(Collectors.toList());
    return new BM25Index(corpus);
}

In this snippet:

  • We split chunks into tokens (words) by lowercasing and splitting their text.
  • We use a custom BM25Index class to create our lexical index.
  • Later, we'll score new queries on this index to get relevance.
Sign up

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal