autorag.nodes.passagereranker package

Subpackages

Submodules

autorag.nodes.passagereranker.base module

autorag.nodes.passagereranker.base.passage_reranker_node(func)[source]

autorag.nodes.passagereranker.cohere module

async autorag.nodes.passagereranker.cohere.cohere_rerank_pure(cohere_client: AsyncClient, model: str, query: str, documents: List[str], ids: List[str], top_k: int) → Tuple[List[str], List[str], List[float]][source]

Rerank a list of contents with Cohere rerank models.

Parameters:

  • cohere_client – The Cohere AsyncClient to use for reranking

  • model – The model name for Cohere rerank

  • query – The query to use for reranking

  • documents – The list of contents to rerank

  • ids – The list of ids corresponding to the documents

  • top_k – The number of passages to be retrieved

Returns:

Tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.cohere.cohere_reranker(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, batch: int = 64, model: str = 'rerank-multilingual-v2.0', api_key: str | None = None) → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents with Cohere rerank models. You can get the API key from https://cohere.com/rerank and set it in the environment variable COHERE_API_KEY.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • batch – The number of queries to be processed in a batch

  • model – The model name for Cohere rerank. You can choose between “rerank-multilingual-v2.0” and “rerank-english-v2.0”. Default is “rerank-multilingual-v2.0”.

  • api_key – The API key for Cohere rerank. You can set it in the environment variable COHERE_API_KEY. Or, you can directly set it on the config YAML file using this parameter. Default is env variable “COHERE_API_KEY”.

Returns:

Tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.colbert module

autorag.nodes.passagereranker.colbert.colbert_reranker(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, batch: int = 64, model_name: str = 'colbert-ir/colbertv2.0') → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents with Colbert rerank models. You can get more information about a Colbert model at https://huggingface.co/colbert-ir/colbertv2.0. It uses BERT-based model, so recommend using CUDA gpu for faster reranking.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • batch – The number of queries to be processed in a batch Default is 64.

  • model_name – The model name for Colbert rerank. You can choose colbert model for reranking. Default is “colbert-ir/colbertv2.0”.

Returns:

Tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.colbert.get_colbert_embedding_batch(input_strings: List[str], model, tokenizer, batch_size: int) → List[array][source]
autorag.nodes.passagereranker.colbert.get_colbert_score(query_embedding: array, content_embedding: array) → float[source]
autorag.nodes.passagereranker.colbert.slice_tensor(input_tensor, batch_size)[source]
autorag.nodes.passagereranker.colbert.slice_tokenizer_result(tokenizer_output, batch_size)[source]

autorag.nodes.passagereranker.flag_embedding module

autorag.nodes.passagereranker.flag_embedding.flag_embedding_reranker(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, batch: int = 64, use_fp16: bool = False, model_name: str = 'BAAI/bge-reranker-large') → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents based on their relevance to a query using BAAI normal-Reranker model.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • batch – The number of queries to be processed in a batch Default is 64.

  • use_fp16 – Whether to use fp16 for inference

  • model_name – The name of the BAAI Reranker normal-model name. Default is “BAAI/bge-reranker-large”

Returns:

tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.flag_embedding.flag_embedding_run_model(input_texts, model, batch_size: int)[source]

autorag.nodes.passagereranker.flag_embedding_llm module

autorag.nodes.passagereranker.flag_embedding_llm.flag_embedding_llm_reranker(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, batch: int = 64, use_fp16: bool = False, model_name: str = 'BAAI/bge-reranker-v2-gemma') → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents based on their relevance to a query using BAAI LLM-based-Reranker model.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • batch – The number of queries to be processed in a batch Default is 64.

  • use_fp16 – Whether to use fp16 for inference

  • model_name – The name of the BAAI Reranker LLM-based-model name. Default is “BAAI/bge-reranker-v2-gemma”

Returns:

tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.jina module

autorag.nodes.passagereranker.jina.jina_reranker(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, api_key: str | None = None, model: str = 'jina-reranker-v1-base-en', batch: int = 8) → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents with Jina rerank models. You can get the API key from https://jina.ai/reranker and set it in the environment variable JINAAI_API_KEY.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • api_key – The API key for Jina rerank. You can set it in the environment variable JINAAI_API_KEY. Or, you can directly set it on the config YAML file using this parameter. Default is env variable “JINAAI_API_KEY”.

  • model – The model name for Cohere rerank. You can choose between “jina-reranker-v1-base-en” and “jina-colbert-v1-en”. Default is “jina-reranker-v1-base-en”.

  • batch – The number of queries to be processed in a batch

Returns:

Tuple of lists containing the reranked contents, ids, and scores

async autorag.nodes.passagereranker.jina.jina_reranker_pure(query: str, contents: List[str], scores: List[float], ids: List[str], top_k: int, api_key: str, model: str = 'jina-reranker-v1-base-en') → Tuple[List[str], List[str], List[float]][source]

autorag.nodes.passagereranker.koreranker module

autorag.nodes.passagereranker.koreranker.exp_normalize(x)[source]
autorag.nodes.passagereranker.koreranker.koreranker(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, batch: int = 64) → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents based on their relevance to a query using ko-reranker. ko-reranker is a reranker based on korean (https://huggingface.co/Dongjin-kr/ko-reranker).

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • batch – The number of queries to be processed in a batch Default is 64.

Returns:

tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.koreranker.koreranker_run_model(input_texts, model, tokenizer, device, batch_size: int)[source]

autorag.nodes.passagereranker.monot5 module

autorag.nodes.passagereranker.monot5.monot5(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, model_name: str = 'castorini/monot5-3b-msmarco-10k', batch: int = 64) → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents based on their relevance to a query using MonoT5.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • model_name –

    The name of the MonoT5 model to use for reranking Note: default model name is ‘castorini/monot5-3b-msmarco-10k’

    If there is a ‘/’ in the model name parameter, when we create the file to store the results, the path will be twisted because of the ‘/’. Therefore, it will be received as ‘_’ instead of ‘/’.

  • batch – The number of queries to be processed in a batch

Returns:

tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.monot5.monot5_run_model(input_texts, model, batch_size: int, tokenizer, device, token_false_id, token_true_id)[source]

autorag.nodes.passagereranker.pass_reranker module

autorag.nodes.passagereranker.pass_reranker.pass_reranker(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int)[source]

Do not perform reranking. Return the given top-k passages as is.

autorag.nodes.passagereranker.rankgpt module

class autorag.nodes.passagereranker.rankgpt.AsyncRankGPTRerank(top_n: int = 5, llm: LLM | None = None, verbose: bool = False, rankgpt_rerank_prompt: BasePromptTemplate | None = None)[source]

Bases: RankGPTRerank

async async_postprocess_nodes(nodes: List[NodeWithScore], query_bundle: QueryBundle, ids: List[str] | None = None) → Tuple[List[NodeWithScore], List[str]][source]
async async_run_llm(messages: Sequence[ChatMessage]) → ChatResponse[source]
llm: LLM
model_computed_fields: ClassVar[Dict[str, ComputedFieldInfo]] = {}

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_fields: ClassVar[Dict[str, FieldInfo]] = {'callback_manager': FieldInfo(annotation=CallbackManager, required=False, default_factory=CallbackManager, exclude=True), 'llm': FieldInfo(annotation=LLM, required=False, default_factory=get_default_llm, description='LLM to use for rankGPT'), 'rankgpt_rerank_prompt': FieldInfo(annotation=BasePromptTemplate, required=True, description='rankGPT rerank prompt.', metadata=[SerializeAsAny()]), 'top_n': FieldInfo(annotation=int, required=False, default=5, description='Top N nodes to return from reranking.'), 'verbose': FieldInfo(annotation=bool, required=False, default=False, description='Whether to print intermediate steps.')}

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo] objects.

This replaces Model.__fields__ from Pydantic V1.

rankgpt_rerank_prompt: Annotated[BasePromptTemplate, SerializeAsAny()]
top_n: int
verbose: bool
autorag.nodes.passagereranker.rankgpt.rankgpt(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, llm: LLM | None = None, verbose: bool = False, rankgpt_rerank_prompt: str | None = None, batch: int = 16, **kwargs) → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank given context paragraphs using RankGPT. Return pseudo scores, since the actual scores are not available on RankGPT.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • llm – The LLM model to use for RankGPT rerank. It is LlamaIndex model. Default is OpenAI model with gpt-3.5-turbo-16k.

  • verbose – Whether to print intermediate steps.

  • rankgpt_rerank_prompt – The prompt template for RankGPT rerank. Default is RankGPT’s default prompt.

  • batch – The number of queries to be processed in a batch.

Returns:

Tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.run module

autorag.nodes.passagereranker.run.run_passage_reranker_node(modules: List[Callable], module_params: List[Dict], previous_result: DataFrame, node_line_dir: str, strategies: Dict) → DataFrame[source]

Run evaluation and select the best module among passage reranker node results.

Parameters:

  • modules – Passage reranker modules to run.

  • module_params – Passage reranker module parameters.

  • previous_result – Previous result dataframe. Could be retrieval, reranker modules result. It means it must contain ‘query’, ‘retrieved_contents’, ‘retrieved_ids’, ‘retrieve_scores’ columns.

  • node_line_dir – This node line’s directory.

  • strategies – Strategies for passage reranker node. In this node, we use ‘retrieval_f1’, ‘retrieval_recall’ and ‘retrieval_precision’. You can skip evaluation when you use only one module and a module parameter.

Returns:

The best result dataframe with previous result columns.

autorag.nodes.passagereranker.sentence_transformer module

autorag.nodes.passagereranker.sentence_transformer.sentence_transformer_reranker(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, batch: int = 64, sentence_transformer_max_length: int = 512, model_name: str = 'cross-encoder/ms-marco-MiniLM-L-2-v2') → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents based on their relevance to a query using Sentence Transformer model.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • batch – The number of queries to be processed in a batch

  • sentence_transformer_max_length – The maximum length of the input text for the Sentence Transformer model

  • model_name – The name of the Sentence Transformer model to use for reranking Default is “cross-encoder/ms-marco-MiniLM-L-2-v2”

Returns:

tuple of lists containing the reranked contents, ids, and scores

autorag.nodes.passagereranker.sentence_transformer.sentence_transformer_run_model(input_texts, model, batch_size: int)[source]

autorag.nodes.passagereranker.time_reranker module

autorag.nodes.passagereranker.time_reranker.time_reranker(contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, time_list: List[List[datetime]]) → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank the passages based on merely the datetime of the passage. It uses ‘last_modified_datetime’ key in the corpus metadata, so the metadata should be in the format of {‘last_modified_datetime’: datetime.datetime} at the corpus data file.

Parameters:

  • contents_list – The list of lists of contents

  • scores_list – The list of lists of scores from the initial ranking

  • ids_list – The list of lists of ids

  • top_k – The number of passages to be retrieved after reranking

  • time_list – The metadata list of lists of datetime.datetime It automatically extracts the ‘last_modified_datetime’ key from the metadata in the corpus data.

Returns:

The reranked contents, ids, and scores

autorag.nodes.passagereranker.upr module

class autorag.nodes.passagereranker.upr.UPRScorer(suffix_prompt: str, prefix_prompt: str, use_bf16: bool = False)[source]

Bases: object

compute(query: str, contents: List[str]) → List[float][source]
autorag.nodes.passagereranker.upr.upr(queries: List[str], contents_list: List[List[str]], scores_list: List[List[float]], ids_list: List[List[str]], top_k: int, use_bf16: bool = False, prefix_prompt: str = 'Passage: ', suffix_prompt: str = 'Please write a question based on this passage.') → Tuple[List[List[str]], List[List[str]], List[List[float]]][source]

Rerank a list of contents based on their relevance to a query using UPR. UPR is a reranker based on UPR (https://github.com/DevSinghSachan/unsupervised-passage-reranking). The language model will make a question based on the passage and rerank the passages by the likelihood of the question. The default model is t5-large.

Parameters:

  • queries – The list of queries to use for reranking

  • contents_list – The list of lists of contents to rerank

  • scores_list – The list of lists of scores retrieved from the initial ranking

  • ids_list – The list of lists of ids retrieved from the initial ranking

  • top_k – The number of passages to be retrieved

  • use_bf16 – Whether to use bfloat16 for the model. Default is False.

  • prefix_prompt – The prefix prompt for the language model that generates question for reranking. Default is “Passage: “. The prefix prompt serves as the initial context or instruction for the language model. It sets the stage for what is expected in the output

  • suffix_prompt –

    The suffix prompt for the language model that generates question for reranking. Default is “Please write a question based on this passage.”. The suffix prompt provides a cue or a closing instruction to the language model,

    signaling how to conclude the generated text or what format to follow at the end.

Returns:

tuple of lists containing the reranked contents, ids, and scores

Module contents