LayerWiseLLMReranker#

class FlagEmbedding.inference.reranker.decoder_only.layerwise.LayerWiseLLMReranker(model_name_or_path: str, peft_path: str | None = None, use_fp16: bool = False, use_bf16: bool = False, query_instruction_for_rerank: str = 'A: ', query_instruction_format: str = '{}{}', passage_instruction_for_rerank: str = 'B: ', passage_instruction_format: str = '{}{}', cache_dir: str | None = None, trust_remote_code: bool = False, devices: str | List[str] | List[int] | None = None, cutoff_layers: List[int] | None = None, prompt: str | None = None, batch_size: int = 128, query_max_length: int | None = None, max_length: int = 512, normalize: bool = False, **kwargs: Any)[source]#

Base reranker class for layerwise LLM like decoder only models.

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.

  • peft_path (Optional[str], optional) – Path to the PEFT config. Defaults to None.

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

  • use_bf16 (bool, optional) – Another type of half-precision floating-point, you can use bf16 if the hardware supports. Defaults to :data:False.

  • query_instruction_for_rerank (str, optional) – Query instruction for retrieval tasks, which will be used with with query_instruction_format. Defaults to "A: ".

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

  • passage_instruction_for_rerank (str, optional) – Passage instruction for retrieval tasks, which will be used with with passage_instruction_format. Defaults to "B: ".

  • passage_instruction_format (str, optional) – The template for passage. Defaults to “{}{}”.

  • cache_dir (Optional[str], optional) – Cache directory for the model. Defaults to None.

  • trust_remote_code (bool, optional) – trust_remote_code. Defaults to False.

  • devices (Union[str, List[str], List[int]], optional) – Devices to use for model inference, such as [“cuda:0”] or [“0”]. Defaults to None.

  • cutoff_layers (Optional[List[int]]) – Pick which layers are used for computing the score. Defaults to None.

  • prompt (Optional[str], optional) – Prompt for the specific task. Defaults to None.

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

  • query_max_length (int, optional) – Maximum length for queries. If not specified, will be 3/4 of max_length. Defaults to None.

  • max_length (int, optional) – Maximum length of passages. Defaults to :data`512`.

  • normalize (bool, optional) – If True, use Sigmoid to normalize the results. Defaults to False.

Methods#

class FlagEmbedding.inference.reranker.decoder_only.layerwise.LayerWiseLLMReranker.compute_score_single_gpu(self, sentence_pairs: List[Tuple[str, str]] | Tuple[str, str], batch_size: int | None = None, query_max_length: int | None = None, max_length: int | None = None, cutoff_layers: List[int] | None = None, prompt: str | None = None, normalize: bool | None = None, use_dataloader: bool = False, num_workers: int | None = None, device: str | None = None, **kwargs: Any)#

Compute the relevance scores using a single GPU.

Parameters:

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

  • batch_size (Optional[int], optional) – Number of inputs for each iter. Defaults to None.

  • query_max_length (Optional[int], optional) – Maximum length of tokens of queries. Defaults to None.

  • max_length (Optional[int], optional) – Maximum length of tokens. Defaults to None.

  • cutoff_layers (Optional[List[int]], optional) – Pick which layers are used for computing the score. Defaults to None.

  • prompt (Optional[str], optional) – Prompt for the specific task. Defaults to None.

  • normalize (Optional[bool], optional) – If True, use Sigmoid to normalize the results. Defaults to None.

  • use_dataloader (bool, optional) – If True, will use the dataloader to load the datasets. Defaults to False.

  • num_workers (int, optional) – Number of workers for dataloader. Defaults to None.

  • device (Optional[str], optional) – Device to use for computation. Defaults to None.

Returns:

The computed scores.

Return type:

List[float]