Reranker

Modules concerning the rerankers used in Gen AI applications.

CohereBedrockReranker(model_name, access_key_id=None, secret_access_key=None, region_name='us-east-1', fallback_to_original=True, **kwargs)

Bases: BaseReranker

A reranker to rerank search results using AWS Bedrock's Cohere service.

Examples:

Basic usage with explicit credentials:

reranker = CohereBedrockReranker(
    model_name="cohere.rerank-v3-5:0",
    access_key_id="your_access_key_id",
    secret_access_key="your_secret_access_key"
)

chunks = [Chunk(content="Python programming"), Chunk(content="ML algorithms")]
reranked = await reranker.rerank(chunks, "machine learning")

With explicit region:

reranker = CohereBedrockReranker(
    model_name="cohere.rerank-v3-5:0",
    access_key_id="your_key",
    secret_access_key="your_secret",
    region_name="us-west-2"
)

Graceful fallback on errors:

reranker = CohereBedrockReranker(
    model_name="cohere.rerank-v3-5:0",
    access_key_id="invalid_key",
    secret_access_key="invalid_secret"
)
result = await reranker.rerank(chunks, "query")

Note

Requires AWS credentials with Bedrock permissions. Currently supported model: 1. cohere.rerank-v3-5:0

See: https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html

Initializes a new instance of the CohereBedrockReranker class.

Parameters:

Name	Type	Description	Default
model_name	str	The model to use for reranking. Currently supported: "cohere.rerank-v3-5:0"	required
access_key_id	str | None	The AWS access key ID. Defaults to None, in which case the AWS_ACCESS_KEY_ID environment variable will be used.	None
secret_access_key	str | None	The AWS secret access key. Defaults to None, in which case the AWS_SECRET_ACCESS_KEY environment variable will be used.	None
region_name	str	The AWS region name. Defaults to "us-east-1".	'us-east-1'
fallback_to_original	bool	Whether to fallback to the original chunks if the reranking fails. Defaults to True.	True
**kwargs	Any	Additional keyword arguments.	{}

FlagEmbeddingReranker(model_path, use_fp16=True)

Bases: BaseReranker

Reranks a list of chunks using a flag embedding.

Requires the FlagEmbedding package to be installed.

Examples:

reranker = FlagEmbeddingReranker(
    model_path="path/to/your/model",
    use_fp16=True
)

chunks = [Chunk(content="Python programming"), Chunk(content="ML algorithms")]
reranked = await reranker.rerank(chunks, "machine learning")

Attributes:

Name	Type	Description
reranker	FlagReranker	The flag embedding reranker model.

Initializes a new instance of the FlagEmbeddingReranker class.

Parameters:

Name	Type	Description	Default
model_path	str	A path containing the reranker model.	required
use_fp16	bool	Whether to reduce the model size to FP16 or not. Defaults to True.	True

SimilarityBasedReranker(embeddings, similarity_func=lambda x, y: 1 - cosine_distance(x, y))

Bases: BaseReranker

A class to rerank chunks based on their similarity to a given query.

Examples:

reranker = SimilarityBasedReranker(
    embeddings=OpenAIEMInvoker(...),
    similarity_func=lambda x, y: 1 - cosine_distance(x, y)
)

chunks = [Chunk(content="Python programming"), Chunk(content="ML algorithms")]
reranked = await reranker.rerank(chunks, "machine learning")

Attributes:

Name	Type	Description
embeddings	BaseEMInvoker	An instance of the BaseEMInvoker class to embed text.
similarity_func	Callable[[list[float], list[float]], float]	A callback function to calculate the similarity.

Initializes the SimilarityBasedReranker class with a similarity function.

This constructor method initializes an instance of the SimilarityBasedReranker class, setting up the similarity function that will be used to rerank chunks based on their similarity to a query.

Parameters:

Name	Type	Description	Default
embeddings	BaseEMInvoker	An instance of the BaseEMInvoker class that will be used to calculate the embeddings of the query and the chunks.	required
similarity_func	Callable[[list[float], list[float]], float]	A callback function that takes two parameters (the query embeddings and chunk embeddings), and returns a similarity score as a float. Defaults to cosine similarity.	lambda x, y: 1 - cosine(x, y)

Returns:

Type	Description
None	None

TEIReranker(url, username=None, password=None, api_key=None, timeout=10, fallback_to_original=False)

Bases: BaseReranker

A class to rerank chunks by utilizing a reranker model hosted on Text Embedding Inference (TEI).

Examples:

reranker = TEIReranker(
    url="https://api-inference.huggingface.co/models/your-model-id"
)

chunks = [Chunk(content="Python programming"), Chunk(content="ML algorithms")]
reranked = await reranker.rerank(chunks, "machine learning")

Attributes:

Name	Type	Description
url	str	The URL of the TEI endpoint that hosts the reranker model.
auth_kwargs	dict[str, Any]	Additional authentication arguments for the TEI reranker model.
timeout	int	The timeout for the TEI reranker model in seconds.
fallback_to_original	bool	Whether to fallback to the original chunks if the reranking fails.

Initializes a new instance of the TEIReranker class.

Parameters:

Name	Type	Description	Default
url	str	The URL of the TEI endpoint that hosts the reranker model.	required
username	str | None	The username for the TEI reranker model. Defaults to None.	None
password	str | None	The password for the TEI reranker model. Defaults to None.	None
api_key	str | None	The API key for the TEI reranker model. Defaults to None.	None
timeout	int	The timeout for the TEI reranker model in seconds. Defaults to 10.	10
fallback_to_original	bool	Whether to fallback to the original chunks if the reranking fails. Defaults to False.	False

close() async

Closes the aiohttp ClientSession if it exists.

This method should be called when the reranker is no longer needed to properly clean up resources.