AbsReranker#

class FlagEmbedding.abc.inference.AbsReranker(model_name_or_path: str, use_fp16: bool = False, query_instruction_for_rerank: str | None = None, query_instruction_format: str = '{}{}', passage_instruction_for_rerank: str | None = None, passage_instruction_format: str = '{}{}', devices: str | int | List[str] | List[int] | None = None, batch_size: int = 128, query_max_length: int | None = None, max_length: int = 512, normalize: bool = False, **kwargs: Any)[source]#

Base class for Reranker. Extend this class and implement compute_score_single_gpu() for custom rerankers.

Parameters:

  • model_name_or_path (str) – If it’s a path to a local model, it loads the model from the path. Otherwise tries to download and load a model from HuggingFace Hub with the name.

  • use_fp16 (bool, optional) – If true, use half-precision floating-point to speed up computation with a slight performance degradation. Defaults to False.

  • query_instruction_for_rerank – (Optional[str], optional): Query instruction for reranking, which will be used with with query_instruction_format. Defaults to None.

  • query_instruction_format – (str, optional): The template for query_instruction_for_rerank. Defaults to "{}{}".

  • passage_instruction_for_rerank (Optional[str], optional) – Passage instruction for reranking. Defaults to None.

  • passage_instruction_format (str, optional) – Passage instruction format when using passage_instruction_for_rerank. Defaults to "{}{}".

  • devices (Optional[Union[str, int, List[str], List[int]]], optional) – Devices to use for model inference. Defaults to None.

  • batch_size (int, optional) – Batch size for inference. Defaults to 128.

  • query_max_length (int, optional) – Maximum length for query. Defaults to None.

  • max_length (int, optional) – Maximum length. Defaults to 512.

  • normalize (bool, optional) – If true, normalize the result. Defaults to False.

  • kwargs (Dict[Any], optional) – Additional parameters for HuggingFace Transformers config or children classes.

Methods#

static AbsReranker.get_target_devices(devices: str | int | List[str] | List[int]) → List[str][source]#
Parameters:

devices (Union[str, int, List[str], List[int]]) – Specified devices, can be str, int, list of str, or list of int.

Raises:

ValueError – Devices should be a string or an integer or a list of strings or a list of integers.

Returns:

A list of target devices in format

Return type:

List[str]

AbsReranker.get_detailed_instruct(instruction_format: str, instruction: str, sentence: str)[source]#

Combine the instruction and sentence along with the instruction format.

Parameters:

  • instruction_format (str) – Format for instruction.

  • instruction (str) – The text of instruction.

  • sentence (str) – The sentence to concatenate with.

Returns:

The complete sentence with instruction

Return type:

str

AbsReranker.get_detailed_inputs(sentence_pairs: str | List[str])[source]#

get detailed instruct for all the inputs

Parameters:

sentence_pairs (Union[str, List[str]]) – Input sentence pairs

Returns:

The complete sentence pairs with instruction

Return type:

list[list[str]]

AbsReranker.compute_score(sentence_pairs: List[Tuple[str, str]] | Tuple[str, str], **kwargs)[source]#

Compute score for each sentence pair

Parameters:

sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]) – Input sentence pairs to compute.

Returns:

scores of all the sentence pairs.

Return type:

numpy.ndarray

abstractmethod AbsReranker.compute_score_single_gpu(sentence_pairs: List[Tuple[str, str]] | Tuple[str, str], batch_size: int = 256, query_max_length: int | None = None, max_length: int = 512, normalize: bool = False, device: str | None = None, **kwargs: Any)[source]#

This method should compute the scores of sentence_pair and return scores.

AbsReranker.start_multi_process_pool() → Dict[Literal['input', 'output', 'processes'], Any][source]#

Starts a multi-process pool to process the encoding with several independent processes via SentenceTransformer.encode_multi_process.

This method is recommended if you want to encode on multiple GPUs or CPUs. It is advised to start only one process per GPU. This method works together with encode_multi_process and stop_multi_process_pool.

Returns:

A dictionary with the target processes, an input queue, and an output queue.

Return type:

Dict[str, Any]

AbsReranker.encode_multi_process(sentence_pairs: List, pool: Dict[Literal['input', 'output', 'processes'], Any], **kwargs) → ndarray[source]#
static AbsReranker._encode_multi_process_worker(target_device: str, model: AbsReranker, input_queue: Queue, results_queue: Queue) → None[source]#

Internal working process to encode sentences in multi-process setup

static AbsReranker.stop_multi_process_pool(pool: Dict[Literal['input', 'output', 'processes'], Any]) → None[source]#

Stops all processes started with start_multi_process_pool.

Parameters:

pool (Dict[str, object]) – A dictionary containing the input queue, output queue, and process list.

Returns:

None