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: Grasp the intuition behind Okapi BM25 for lexical retrieval. Construct a BM25 index on your corpus. 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. Javaimport 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.

Join the 1M+ learners on CodeSignal

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

Merging BM25 and Embedding-Based Retrieval: BM25 Scoring & Similarity Calculation

Below is the first segment of a hybridRetrieval function that computes BM25 scores for each chunk and converts distances from an embedding-based search to similarities:

Java

import java.util.HashMap;
import java.util.Map;
import tech.amikos.chromadb.Collection;
import tech.amikos.chromadb.Collection.QueryResponse;
import tech.amikos.chromadb.handler.ApiException;

public static List<Map<String, Object>> hybridRetrieval(
        String query, List<Map<String, Object>> chunks, BM25Index bm25, Collection collection, 
        int topK, float alpha) throws ApiException {
    
    String[] tokenizedQuery = query.toLowerCase().split("\\s+");
    float[] bm25Scores = bm25.getScores(tokenizedQuery);

    QueryResponse embedResults = collection.query(Arrays.asList(query), Math.min(topK * 5, chunks.size()), null, null, null);
    Map<String, Float> embedScoresDict = new HashMap<>();
    
    List<String> ids = embedResults.getIds().get(0);
    List<Float> distances = embedResults.getDistances().get(0);
    
    for (int i = 0; i < ids.size(); i++) {
        String idx = ids.get(i);
        float distance = distances.get(i);
        float similarity = (float) (1 / (1 + distance));
        embedScoresDict.put(idx, similarity);
    }

import java.util.HashMap;
import java.util.Map;
import tech.amikos.chromadb.Collection;
import tech.amikos.chromadb.Collection.QueryResponse;
import tech.amikos.chromadb.handler.ApiException;

public static List<Map<String, Object>> hybridRetrieval(
        String query, List<Map<String, Object>> chunks, BM25Index bm25, Collection collection, 
        int topK, float alpha) throws ApiException {
    
    String[] tokenizedQuery = query.toLowerCase().split("\\s+");
    float[] bm25Scores = bm25.getScores(tokenizedQuery);

    QueryResponse embedResults = collection.query(Arrays.asList(query), Math.min(topK * 5, chunks.size()), null, null, null);
    Map<String, Float> embedScoresDict = new HashMap<>();
    
    List<String> ids = embedResults.getIds().get(0);
    List<Float> distances = embedResults.getDistances().get(0);
    
    for (int i = 0; i < ids.size(); i++) {
        String idx = ids.get(i);
        float distance = distances.get(i);
        float similarity = (float) (1 / (1 + distance));
        embedScoresDict.put(idx, similarity);
    }

In this code:

The function signature includes the parameters we need: the user query, the chunk corpus, the BM25 index, the embedding-based collection, how many top results to return (topK), and the alpha weight.
We compute BM25 scores by tokenizing the query (lowercasing and splitting), then calling bm25.getScores().
Next, we query the embedding-based store to retrieve potential matches.
We loop through these matches to transform distances into similarities (using a simple 1 / (1 + distance) formula) and store them in a dictionary keyed by chunk index.

Merging BM25 and Embedding-Based Retrieval: Normalizing Scores & Final Ranking

Below is the second segment of the same function. It normalizes the BM25 scores, merges them with the similarity values, sorts by the final combined score, and returns the top results:

Java

    List<Map<String, Object>> merged = new ArrayList<>();
    
    float bm25Min = Float.MAX_VALUE;
    float bm25Max = Float.MIN_VALUE;
    
    if (bm25Scores.length > 0) {
        for (float score : bm25Scores) {
            bm25Min = Math.min(bm25Min, score);
            bm25Max = Math.max(bm25Max, score);
        }
    } else {
        bm25Min = 0;
        bm25Max = 1;
    }
    
    for (int i = 0; i < chunks.size(); i++) {
        float bm25Raw = bm25Scores[i];
        float bm25Norm;
        
        if (bm25Max != bm25Min) {
            bm25Norm = (bm25Raw - bm25Min) / (bm25Max - bm25Min);
        } else {
            bm25Norm = 0.0f;
        }
        
        float embedSim = embedScoresDict.getOrDefault(String.valueOf(i), 0.0f);
        float finalScore = alpha * bm25Norm + (1 - alpha) * embedSim;
        
        merged.add(Map.of("id", i, "chunk", chunks.get(i), "score", finalScore));
    }
    
    merged.sort(Comparator.comparingDouble(
        (Map<String, Object> m) -> ((Number) m.get("score")).doubleValue()
    ).reversed());
    
    List<Map<String, Object>> topResults = merged.subList(0, Math.min(topK, merged.size()));
    
    return topResults;
}

    List<Map<String, Object>> merged = new ArrayList<>();
    
    float bm25Min = Float.MAX_VALUE;
    float bm25Max = Float.MIN_VALUE;
    
    if (bm25Scores.length > 0) {
        for (float score : bm25Scores) {
            bm25Min = Math.min(bm25Min, score);
            bm25Max = Math.max(bm25Max, score);
        }
    } else {
        bm25Min = 0;
        bm25Max = 1;
    }
    
    for (int i = 0; i < chunks.size(); i++) {
        float bm25Raw = bm25Scores[i];
        float bm25Norm;
        
        if (bm25Max != bm25Min) {
            bm25Norm = (bm25Raw - bm25Min) / (bm25Max - bm25Min);
        } else {
            bm25Norm = 0.0f;
        }
        
        float embedSim = embedScoresDict.getOrDefault(String.valueOf(i), 0.0f);
        float finalScore = alpha * bm25Norm + (1 - alpha) * embedSim;
        
        merged.add(Map.of("id", i, "chunk", chunks.get(i), "score", finalScore));
    }
    
    merged.sort(Comparator.comparingDouble(
        (Map<String, Object> m) -> ((Number) m.get("score")).doubleValue()
    ).reversed());
    
    List<Map<String, Object>> topResults = merged.subList(0, Math.min(topK, merged.size()));
    
    return topResults;
}

Here's what's happening in this snippet:

First, we find the minimum and maximum BM25 scores so we can normalize values into the [0, 1] range.
For each chunk, we retrieve its BM25 score and its corresponding embedding similarity.
We compute a final score as the weighted sum of the normalized BM25 score and the embedding similarity, controlled by alpha.
Finally, we sort all chunks by the final combined score (descending) and return the top results.

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 due to repetition. 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.

Putting It All Together

Below is a brief illustration of how you could integrate the above functions into your pipeline:

Java

import java.io.IOException;
import java.nio.file.Paths;
import java.util.List;
import java.util.Map;

public static void main(String[] args) throws IOException, ApiException {
    String currentDir = Paths.get("").toAbsolutePath().toString();
    String datasetFile = Paths.get(currentDir, "src", "main", "resources", "corpus-mini.json").toString();

    List<Map<String, Object>> chunkedDocs = Utils.loadAndChunkDataset(datasetFile, 40);
    BM25Index bm25Index = buildBM25Index(chunkedDocs);
    Collection collection = VectorDB.buildChromaCollection(chunkedDocs, "hybrid_collection");

    String query = "What are the latest AI uses?";
    List<Map<String, Object>> results = hybridRetrieval(query, chunkedDocs, bm25Index, collection, 3, 0.6f);
    
    if (results.isEmpty()) {
        System.out.println("No chunks found. Fallback to a naive or apology answer.");
    } else {
        for (Map<String, Object> r : results) {
            System.out.printf("Chunk ID: %d, Score: %.4f, Text: %s%n", 
                    r.get("id"), r.get("score"), r.get("chunk").toString());
        }
    }
}

import java.io.IOException;
import java.nio.file.Paths;
import java.util.List;
import java.util.Map;

public static void main(String[] args) throws IOException, ApiException {
    String currentDir = Paths.get("").toAbsolutePath().toString();
    String datasetFile = Paths.get(currentDir, "src", "main", "resources", "corpus-mini.json").toString();

    List<Map<String, Object>> chunkedDocs = Utils.loadAndChunkDataset(datasetFile, 40);
    BM25Index bm25Index = buildBM25Index(chunkedDocs);
    Collection collection = VectorDB.buildChromaCollection(chunkedDocs, "hybrid_collection");

    String query = "What are the latest AI uses?";
    List<Map<String, Object>> results = hybridRetrieval(query, chunkedDocs, bm25Index, collection, 3, 0.6f);
    
    if (results.isEmpty()) {
        System.out.println("No chunks found. Fallback to a naive or apology answer.");
    } else {
        for (Map<String, Object> r : results) {
            System.out.printf("Chunk ID: %d, Score: %.4f, Text: %s%n", 
                    r.get("id"), r.get("score"), r.get("chunk").toString());
        }
    }
}

Here, we:

Load or chunk our corpus into manageable pieces.
Build a BM25 index for lexical retrieval.
Create an embedding-based collection for semantic retrieval.
Blend both retrieval methods in one pipeline.
Inspect the top results by their combined scores.

Introduction

Understanding the Okapi BM25 Algorithm

Building a BM25 Index

Join the 1M+ learners on CodeSignal

Merging BM25 and Embedding-Based Retrieval: BM25 Scoring & Similarity Calculation

Merging BM25 and Embedding-Based Retrieval: Normalizing Scores & Final Ranking

Putting It All Together

Choosing the Alpha Parameter

Conclusion and Next Steps