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`