Ensemble Retriever

Author: 3dkids
Peer Review: r14minji, jeongkpa
Proofread : jishin86
This is a part of LangChain Open Tutorial

Overview

This notebook explores the creation and use of an EnsembleRetriever in LangChain to improve information retrieval by combining multiple retrieval methods. The EnsembleRetriever integrates the strengths of sparse and dense retrieval algorithms, using weights and runtime configurations for tailored performance.

Key Features

integrate multiple searchers: take different types of searchers as input and combine results.
result re-ranking: uses the Reciprocal Rank Fusion algorithm to re-rank results.
hybrid search: mainly uses a combination of sparse retriever (e.g. BM25) and dense retriever (e.g. embedding similarity).

Advantages

Sparse retriever: effective for keyword-based searches
Dense retriever: effective for semantic similarity-based searches

Due to these complementary characteristics, EnsembleRetriever can provide improved performance in a variety of search scenarios.

For more information, please refer to the LangChain official documentation

References

Environment Setup

Set up the environment. You may refer to Environment Setup for more details.

[Note]

langchain-opentutorial is a package that provides a set of easy-to-use environment setup, useful functions and utilities for tutorials.
You can checkout the langchain-opentutorial for more details.

%%capture --no-stderr
!pip install langchain-opentutorial

# Install required packages
from langchain_opentutorial import package

package.install(
    [
        "langchain_core",  # Core functionality of LangChain
        "langchain_community",  # Community-supported integrations
        "langchain_openai",  # OpenAI integration for embeddings and models
        "rank_bm25",  # BM25 ranking algorithm for information retrieval
    ],
    verbose=False,  # Suppress detailed installation logs
    upgrade=False,  # Do not upgrade packages if already installed
)

# Set environment variables
from langchain_opentutorial import set_env

set_env(
    {
        "OPENAI_API_KEY": "",
        "LANGCHAIN_API_KEY": "",
        "LANGCHAIN_TRACING_V2": "true",
        "LANGCHAIN_ENDPOINT": "https://api.smith.langchain.com",
        "LANGCHAIN_PROJECT": "Conversation-With-History",
    }
)

Environment variables have been set successfully.

# Configuration file to manage API keys as environment variables
from dotenv import load_dotenv

# Load API key information
load_dotenv(override=True)

False

Creating and Configuring Ensemble Retrievers

Initializing an ensemble retriever Ensemble retrievers combine two discovery mechanisms

Sparse search: Uses BM25Retriever for keyword-based matching.
Dense search: Uses FAISS with OpenAI embedding for semantic similarity.
Initialize EnsembleRetriever to combine the BM25Retriever and FAISS searchers. Set the weights for each searcher.

from langchain.retrievers import BM25Retriever, EnsembleRetriever
from langchain.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings

# list sample documents
doc_list = [
    "I like apples",
    "I like apple company",
    "I like apple's iphone",
    "Apple is my favorite company",
    "I like apple's ipad",
    "I like apple's macbook",
]

# Initialize the bm25 retriever and faiss retriever.
bm25_retriever = BM25Retriever.from_texts(
    doc_list,
)
bm25_retriever.k = 1  # Set the number of search results for BM25Retriever to 1.

embedding = OpenAIEmbeddings()  # Enable OpenAI embedding.

faiss_vectorstore = FAISS.from_texts(
    doc_list,
    embedding,
)
faiss_retriever = faiss_vectorstore.as_retriever(search_kwargs={"k": 1})

# Initialize the ensemble retriever.
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, faiss_retriever],
    weights=[0.7, 0.3],
)

Query Execution

Perform retrieval for a given query using ensemble_retriever and compare results across retrievers.

Call the get_relevant_documents() method of the ensemble_retriever object to retrieve relevant documents.

# Get the search results document.
query = "my favorite fruit is apple"
ensemble_result = ensemble_retriever.invoke(query)
bm25_result = bm25_retriever.invoke(query)
faiss_result = faiss_retriever.invoke(query)

# Output the fetched documents.
print("[Ensemble Retriever]")
for doc in ensemble_result:
    print(f"Content: {doc.page_content}")
    print()

print("[BM25 Retriever]")
for doc in bm25_result:
    print(f"Content: {doc.page_content}")
    print()

print("[FAISS Retriever]")
for doc in faiss_result:
    print(f"Content: {doc.page_content}")
    print()

[Ensemble Retriever]
    Content: Apple is my favorite company
    
    Content: I like apples
    
    [BM25 Retriever]
    Content: Apple is my favorite company
    
    [FAISS Retriever]
    Content: I like apples

# Get the search results document.
query = "Apple company makes my favorite iphone"
ensemble_result = ensemble_retriever.invoke(query)
bm25_result = bm25_retriever.invoke(query)
faiss_result = faiss_retriever.invoke(query)

# Output the fetched documents.
print("[Ensemble Retriever]")
for doc in ensemble_result:
    print(f"Content: {doc.page_content}")
    print()

print("[BM25 Retriever]")
for doc in bm25_result:
    print(f"Content: {doc.page_content}")
    print()

print("[FAISS Retriever]")
for doc in faiss_result:
    print(f"Content: {doc.page_content}")
    print()

[Ensemble Retriever]
    Content: Apple is my favorite company
    
    Content: I like apple's iphone
    
    [BM25 Retriever]
    Content: Apple is my favorite company
    
    [FAISS Retriever]
    Content: I like apple's iphone

Change runtime config

You can also change the properties of a retriever at runtime. This is possible using the ConfigurableField class.

Define the weights parameter as a ConfigurableField object.
- Set the field's ID to “ensemble_weights”.

from langchain_core.runnables import ConfigurableField

ensemble_retriever = EnsembleRetriever(
    # Set the list of retrievers. Here we use bm25_retriever and faiss_retriever.
    retrievers=[bm25_retriever, faiss_retriever],
).configurable_fields(
    weights=ConfigurableField(
        # Set a unique identifier for the search parameter.
        id="ensemble_weights",
        # Set a name for the search parameter.
        name="Ensemble Weights",
        # Write a description of the search parameters.
        description="Ensemble Weights",
    )
)

Specify the search settings via the config parameter when searching.
- Set the weight of the ensemble_weights option to [1, 0] so that all search results are weighted more heavily toward BM25 retriever.

config = {"configurable": {"ensemble_weights": [1, 0]}}

# Use the config parameter to specify search settings.
docs = ensemble_retriever.invoke("my favorite fruit is apple", config=config)
docs  # Print the search result, docs.

[Document(metadata={}, page_content='Apple is my favorite company'),
     Document(id='6280c2a3-b58f-474e-aeb6-d480bb44d49e', metadata={}, page_content='I like apples')]

This time, we want all search results to be weighted more heavily in favor of the FAISS retriever.

config = {"configurable": {"ensemble_weights": [0, 1]}}

# Use the config parameter to specify search settings.
docs = ensemble_retriever.invoke("my favorite fruit is apple", config=config)
docs  # Print the search result, docs.

[Document(id='6280c2a3-b58f-474e-aeb6-d480bb44d49e', metadata={}, page_content='I like apples'),
     Document(metadata={}, page_content='Apple is my favorite company')]

PreviousContextual Compression Retriever NextLong Context Reorder

Last updated 1 month ago