chore: import upstream snapshot with attribution

FlagEmbedding/inference/__init__.py

@@ -0,0 +1,29 @@
+from .auto_embedder import FlagAutoModel
+from .auto_reranker import FlagAutoReranker
+from .embedder import (
+    FlagModel, BGEM3FlagModel,
+    FlagICLModel, FlagLLMModel, FlagPseudoMoEModel,
+    EmbedderModelClass
+)
+from .reranker import (
+    FlagReranker,
+    FlagLLMReranker, LayerWiseFlagLLMReranker, LightWeightFlagLLMReranker,
+    RerankerModelClass
+)
+
+
+__all__ = [
+    "FlagAutoModel",
+    "FlagAutoReranker",
+    "EmbedderModelClass",
+    "RerankerModelClass",
+    "FlagModel",
+    "BGEM3FlagModel",
+    "FlagICLModel",
+    "FlagLLMModel",
+    "FlagPseudoMoEModel",
+    "FlagReranker",
+    "FlagLLMReranker",
+    "LayerWiseFlagLLMReranker",
+    "LightWeightFlagLLMReranker",
+]

FlagEmbedding/inference/auto_embedder.py

@@ -0,0 +1,115 @@
+import os
+import logging
+from typing import List, Union, Optional
+
+from FlagEmbedding.inference.embedder.model_mapping import (
+    EmbedderModelClass,
+    AUTO_EMBEDDER_MAPPING, EMBEDDER_CLASS_MAPPING
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FlagAutoModel:
+    """
+    Automatically choose the appropriate class to load the embedding model.
+    """
+    def __init__(self):
+        raise EnvironmentError(
+            "FlagAutoModel is designed to be instantiated using the `FlagAutoModel.from_finetuned(model_name_or_path)` method."
+        )
+
+    @classmethod
+    def from_finetuned(
+        cls,
+        model_name_or_path: str,
+        model_class: Optional[Union[str, EmbedderModelClass]] = None,
+        normalize_embeddings: bool = True,
+        use_fp16: bool = True,
+        use_bf16: bool = False,
+        query_instruction_for_retrieval: Optional[str] = None,
+        devices: Optional[Union[str, List[str]]] = None,
+        pooling_method: Optional[str] = None,
+        trust_remote_code: Optional[bool] = None,
+        query_instruction_format: Optional[str] = None,
+        truncate_dim: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Load a finetuned model according to the provided vars.
+
+        Args:
+            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.
+            model_class (Optional[Union[str, EmbedderModelClass]], optional): The embedder class to use. Defaults to :data:`None`.
+            normalize_embeddings (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor. 
+                Defaults to :data:`True`.
+            use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+                degradation. Defaults to :data:`True`.
+            query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
+                :attr:`query_instruction_format`. Defaults to :data:`None`.
+            devices (Optional[Union[str, List[str]]], optional): Devices to use for model inference. Defaults to :data:`None`.
+            pooling_method (Optional[str], optional): Pooling method to get embedding vector from the last hidden state. Defaults to :data:`None`.
+            trust_remote_code (Optional[bool], optional): trust_remote_code for HF datasets or models. Defaults to :data:`None`.
+            query_instruction_format (Optional[str], optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`None`.
+
+        Raises:
+            ValueError
+
+        Returns:
+            AbsEmbedder: The model class to load model, which is child class of :class:`AbsEmbedder`.
+        """
+        model_name = os.path.basename(model_name_or_path)
+        if model_name.startswith("checkpoint-"):
+            model_name = os.path.basename(os.path.dirname(model_name_or_path))
+
+        if model_class is not None:
+            _model_class = EMBEDDER_CLASS_MAPPING[EmbedderModelClass(model_class)]
+            if pooling_method is None:
+                pooling_method = _model_class.DEFAULT_POOLING_METHOD
+                logger.warning(
+                    f"`pooling_method` is not specified, use default pooling method '{pooling_method}'."
+                )
+            if trust_remote_code is None:
+                trust_remote_code = False
+                logger.warning(
+                    f"`trust_remote_code` is not specified, set to default value '{trust_remote_code}'."
+                )
+            if query_instruction_format is None:
+                query_instruction_format = "{}{}"
+                logger.warning(
+                    f"`query_instruction_format` is not specified, set to default value '{query_instruction_format}'."
+                )
+        else:
+            if model_name not in AUTO_EMBEDDER_MAPPING:
+                raise ValueError(
+                    f"Model name '{model_name}' not found in the model mapping. You can pull request to add the model to "
+                    "`https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/embedder/model_mapping.py`. " 
+                    "If need, you can create a new `<model>.py` file in `https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/embedder/encoder_only` "
+                    "or `https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/embedder/decoder_only`. "
+                    "Welcome to contribute! You can also directly specify the corresponding `model_class` to instantiate the model."
+                )
+
+            model_config = AUTO_EMBEDDER_MAPPING[model_name]
+
+            _model_class = model_config.model_class
+            if pooling_method is None:
+                pooling_method = model_config.pooling_method.value
+            if trust_remote_code is None:
+                trust_remote_code = model_config.trust_remote_code
+            if query_instruction_format is None:
+                query_instruction_format = model_config.query_instruction_format
+
+        return _model_class(
+            model_name_or_path,
+            normalize_embeddings=normalize_embeddings,
+            use_fp16=use_fp16,
+            use_bf16=use_bf16,
+            query_instruction_for_retrieval=query_instruction_for_retrieval,
+            query_instruction_format=query_instruction_format,
+            devices=devices,
+            pooling_method=pooling_method,
+            trust_remote_code=trust_remote_code,
+            truncate_dim=truncate_dim,
+            **kwargs,
+        )

FlagEmbedding/inference/auto_reranker.py

@@ -0,0 +1,81 @@
+import os
+import logging
+from typing import Union, Optional
+
+from FlagEmbedding.inference.reranker.model_mapping import (
+    RerankerModelClass,
+    RERANKER_CLASS_MAPPING,
+    AUTO_RERANKER_MAPPING
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FlagAutoReranker:
+    """
+    Automatically choose the appropriate class to load the reranker model.
+    """
+    def __init__(self):
+        raise EnvironmentError(
+            "FlagAutoReranker is designed to be instantiated using the `FlagAutoReranker.from_finetuned(model_name_or_path)` method."
+        )
+
+    @classmethod
+    def from_finetuned(
+        cls,
+        model_name_or_path: str,
+        model_class: Optional[Union[str, RerankerModelClass]] = None,
+        use_fp16: bool = False,
+        trust_remote_code: Optional[bool] = None,
+        **kwargs,
+    ):
+        """
+        Load a finetuned model according to the provided vars.
+
+        Args:
+            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.
+            model_class (Optional[Union[str, RerankerModelClass]], optional): The reranker class to use.. Defaults to :data:`None`.
+            use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+                degradation. Defaults to :data:`False`.
+            trust_remote_code (Optional[bool], optional): trust_remote_code for HF datasets or models. Defaults to :data:`None`.
+
+        Raises:
+            ValueError
+
+        Returns:
+            AbsReranker: The reranker class to load model, which is child class of :class:`AbsReranker`.
+        """
+        model_name = os.path.basename(model_name_or_path)
+        if model_name.startswith("checkpoint-"):
+            model_name = os.path.basename(os.path.dirname(model_name_or_path))
+
+        if model_class is not None:
+            _model_class = RERANKER_CLASS_MAPPING[RerankerModelClass(model_class)]
+            if trust_remote_code is None:
+                trust_remote_code = False
+                logging.warning(
+                    f"`trust_remote_code` is not specified, set to default value '{trust_remote_code}'."
+                )
+        else:
+            if model_name not in AUTO_RERANKER_MAPPING:
+                raise ValueError(
+                    f"Model name '{model_name}' not found in the model mapping. You can pull request to add the model to "
+                    "`https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/reranker/model_mapping.py`. " 
+                    "If need, you can create a new `<model>.py` file in `https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/reranker/encoder_only` "
+                    "or `https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/reranker/decoder_only`. "
+                    "Welcome to contribute! You can also directly specify the corresponding `model_class` to instantiate the model."
+                )
+
+            model_config = AUTO_RERANKER_MAPPING[model_name]
+
+            _model_class = model_config.model_class
+            if trust_remote_code is None:
+                trust_remote_code = model_config.trust_remote_code
+
+        return _model_class(
+            model_name_or_path,
+            use_fp16=use_fp16,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )

FlagEmbedding/inference/embedder/__init__.py

@@ -0,0 +1,12 @@
+from .encoder_only import FlagModel, BGEM3FlagModel
+from .decoder_only import FlagICLModel, FlagLLMModel, FlagPseudoMoEModel
+from .model_mapping import EmbedderModelClass
+
+__all__ = [
+    "FlagModel",
+    "BGEM3FlagModel",
+    "FlagICLModel",
+    "FlagLLMModel",
+    "FlagPseudoMoEModel",
+    "EmbedderModelClass",
+]

FlagEmbedding/inference/embedder/decoder_only/__init__.py

@@ -0,0 +1,9 @@
+from .base import BaseLLMEmbedder as FlagLLMModel
+from .icl import ICLLLMEmbedder as FlagICLModel
+from .pseudo_moe import PseudoMoELLMEmbedder as FlagPseudoMoEModel
+
+__all__ = [
+    "FlagLLMModel",
+    "FlagICLModel",
+    "FlagPseudoMoEModel",
+]

FlagEmbedding/inference/embedder/decoder_only/base.py

@@ -0,0 +1,301 @@
+from tqdm import tqdm, trange
+from typing import cast, Any, List, Union, Optional
+
+import torch
+import numpy as np
+from transformers import AutoModel, AutoTokenizer
+
+from FlagEmbedding.abc.inference import AbsEmbedder
+
+
+# Pooling function for LLM-based embedding models
+def last_token_pool(last_hidden_states: torch.Tensor,
+                    attention_mask: torch.Tensor) -> torch.Tensor:
+    """Last token pooling method.
+
+    Args:
+        last_hidden_state (torch.Tensor): The last hidden state of the model.
+        attention_mask (torch.Tensor): Attention mask. Defaults to :data:`None`.
+
+    Returns:
+        torch.Tensor: The embedding vectors after pooling.
+    """
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+
+class BaseLLMEmbedder(AbsEmbedder):
+    """Base embedder class for LLM like decoder only models.
+
+    Args:
+        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.
+        normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
+        use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+            degradation. Defaults to :data:`True`.
+        query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
+            with :attr:`query_instruction_format`. Defaults to :data:`None`.
+        query_instruction_format (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"Instruct: {}\nQuery: {}"`.
+        devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
+        trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
+        query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
+        passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
+        convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor. 
+            Defaults to :data:`True`.
+    
+    Attributes:
+        DEFAULT_POOLING_METHOD: The default pooling method when running the model.
+    """
+    DEFAULT_POOLING_METHOD = "last_token"
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        normalize_embeddings: bool = True,
+        use_fp16: bool = True,
+        use_bf16: bool = False,
+        query_instruction_for_retrieval: Optional[str] = None,
+        query_instruction_format: str = "Instruct: {}\nQuery: {}", # specify the format of query_instruction_for_retrieval
+        devices: Optional[Union[str, List[str]]] = None, # specify devices, such as "cuda:0" or ["cuda:0", "cuda:1"]
+        # Additional parameters for BaseLLMEmbedder
+        trust_remote_code: bool = False,
+        cache_dir: Optional[str] = None,
+        # inference
+        batch_size: int = 256,
+        query_max_length: int = 512,
+        passage_max_length: int = 512,
+        convert_to_numpy: bool = True,
+        truncate_dim: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_name_or_path,
+            normalize_embeddings=normalize_embeddings,
+            use_fp16=use_fp16,
+            use_bf16=use_bf16,
+            query_instruction_for_retrieval=query_instruction_for_retrieval,
+            query_instruction_format=query_instruction_format,
+            devices=devices,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            passage_max_length=passage_max_length,
+            convert_to_numpy=convert_to_numpy,
+            truncate_dim=truncate_dim,
+            **kwargs
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            cache_dir=cache_dir
+        )
+        self.model = AutoModel.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            cache_dir=cache_dir,
+            dtype=self.get_model_torch_dtype(),
+        )
+
+        if self.kwargs.get("pooling_method", "last_token") != "last_token":
+            raise ValueError("Pooling method must be 'last_token' for LLM-based models.")
+
+    def encode_queries(
+        self,
+        queries: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the queries.
+
+        Args:
+            queries (Union[List[str], str]): Input queries to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
+        """
+        return super().encode_queries(
+            queries,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+
+    def encode_corpus(
+        self,
+        corpus: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the corpus.
+
+        Args:
+            corpus (Union[List[str], str]): Input corpus to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
+        """
+        return super().encode_corpus(
+            corpus,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+
+    def encode(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the input sentences with the embedding model.
+
+        Args:
+            sentences (Union[List[str], str]): Input sentences to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
+        """
+        return super().encode(
+            sentences,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+
+    @torch.no_grad()
+    def encode_single_device(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: int = 256,
+        max_length: int = 512,
+        convert_to_numpy: bool = True,
+        device: Optional[str] = None,
+        **kwargs: Any   # add `pad_to_multiple_of=8` for bge-multilingual-gemmma2
+    ):
+        """Encode input sentences by a single device.
+
+        Args:
+            sentences (Union[List[str], str]): Input sentences to encode.
+            batch_size (int, optional): Number of sentences for each iter. Defaults to :data:`256`.
+            max_length (int, optional): Maximum length of tokens. Defaults to :data:`512`.
+            convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`True`.
+            device (Optional[str], optional): Device to use for encoding. Defaults to None.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
+        """
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu":
+            self.model.float()
+
+        self.model.to(device)
+        self.model.eval()
+
+        input_was_string = False
+        if isinstance(sentences, str):
+            sentences = [sentences]
+            input_was_string = True
+
+        # tokenize without padding to get the correct length
+        all_inputs = []
+        for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
+                                  disable=len(sentences) < batch_size):
+            sentences_batch = sentences[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer(
+                sentences_batch,
+                truncation=True,
+                max_length=max_length,
+                **kwargs
+            )
+            inputs_batch = [{
+                k: inputs_batch[k][i] for k in inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            all_inputs.extend(inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
+        all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                inputs_batch = self.tokenizer.pad(
+                    all_inputs_sorted[: batch_size],
+                    padding=True,
+                    return_tensors='pt',
+                    **kwargs
+                ).to(device)
+                last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+                embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        # encode
+        all_embeddings = []
+        for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
+                                disable=len(sentences) < batch_size):
+            inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer.pad(
+                inputs_batch,
+                padding=True,
+                return_tensors='pt',
+                **kwargs
+            ).to(device)
+            last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+            embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
+            embeddings = self._truncate_embeddings(embeddings)
+            if self.normalize_embeddings:
+                embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
+            embeddings = cast(torch.Tensor, embeddings)
+
+            if convert_to_numpy:
+                embeddings = self._convert_to_numpy(embeddings, device=device)
+            all_embeddings.append(embeddings)
+
+        if convert_to_numpy:
+            all_embeddings = np.concatenate(all_embeddings, axis=0)
+        else:
+            all_embeddings = torch.cat(all_embeddings, dim=0)
+
+        # adjust the order of embeddings
+        all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
+
+        # return the embeddings
+        if input_was_string:
+            return all_embeddings[0]
+        return all_embeddings

FlagEmbedding/inference/embedder/decoder_only/icl.py

@@ -0,0 +1,567 @@
+from tqdm import tqdm, trange
+from typing import cast, Any, List, Union, Optional
+
+import queue
+from multiprocessing import Queue
+
+import gc
+import torch
+import numpy as np
+from transformers import AutoModel, AutoTokenizer
+
+from FlagEmbedding.abc.inference import AbsEmbedder
+
+
+# Pooling function for LLM-based embedding models
+def last_token_pool(last_hidden_states: torch.Tensor,
+                    attention_mask: torch.Tensor) -> torch.Tensor:
+    """Last token pooling method.
+
+    Args:
+        last_hidden_state (torch.Tensor): The last hidden state of the model.
+        attention_mask (torch.Tensor): Attention mask. Defaults to :data:`None`.
+
+    Returns:
+        torch.Tensor: The embedding vectors after pooling.
+    """
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+
+class ICLLLMEmbedder(AbsEmbedder):
+    """
+    Embedder class for BGE-EN-icl.
+    
+    Args:
+        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.
+        normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
+        use_fp16 (bool, optional) If true, use half-precision floating-point to speed up computation with a slight performance 
+            degradation. Defaults to :data:`True`.
+        query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
+            with :attr:`query_instruction_format`. Defaults to :data:`None`.
+        query_instruction_format (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"{}{}"`.
+        devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
+        examples_for_task (Optional[List[dict]], optional): Few-shot examples for the model to enhance model's ability. 
+            Defaults to :data:`None`.
+        examples_instruction_format (str, optional): Example format when using :attr:`examples_for_task`.
+        trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
+        query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
+        passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
+        convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor. 
+            Defaults to :data:`True`.
+    
+    Attributes:
+        DEFAULT_POOLING_METHOD: The default pooling method when running the model.
+    """
+    DEFAULT_POOLING_METHOD = "last_token"
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        normalize_embeddings: bool = True,
+        use_fp16: bool = True,
+        use_bf16: bool = False,
+        query_instruction_for_retrieval: Optional[str] = None,
+        query_instruction_format: str = "<instruct>{}\n<query>{}", # specify the format of query_instruction_for_retrieval
+        suffix: str = '\n<response>',
+        devices: Optional[Union[str, List[str]]] = None, # specify devices, such as "cuda:0" or ["cuda:0", "cuda:1"]
+        # Additional parameters for ICLLLMEmbedder
+        examples_for_task: Optional[List[dict]] = None,
+        examples_instruction_format: str = "<instruct>{}\n<query>{}\n<response>{}", # specify the format of examples_for_task
+        trust_remote_code: bool = False,
+        cache_dir: Optional[str] = None,
+        # inference
+        batch_size: int = 256,
+        query_max_length: int = 512,
+        passage_max_length: int = 512,
+        convert_to_numpy: bool = True,
+        truncate_dim: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        query_instruction_format = query_instruction_format.replace('\\n', '\n')
+        examples_instruction_format = examples_instruction_format.replace('\\n', '\n')
+        super().__init__(
+            model_name_or_path,
+            normalize_embeddings=normalize_embeddings,
+            use_fp16=use_fp16,
+            use_bf16=use_bf16,
+            query_instruction_for_retrieval=query_instruction_for_retrieval,
+            query_instruction_format=query_instruction_format,
+            devices=devices,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            passage_max_length=passage_max_length,
+            convert_to_numpy=convert_to_numpy,
+            truncate_dim=truncate_dim,
+            **kwargs
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            cache_dir=cache_dir
+        )
+        self.model = AutoModel.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            cache_dir=cache_dir,
+            torch_dtype=self.get_model_torch_dtype(),
+        )
+        self.examples_for_task = examples_for_task
+        self.examples_instruction_format = examples_instruction_format
+
+        if self.kwargs.get("pooling_method", "last_token") != "last_token":
+            raise ValueError("Pooling method must be 'last_token' for LLM-based models.")
+
+        self.set_examples()
+        self.suffix = suffix
+
+        self.query_pool = None
+    
+    def __del__(self):
+        self.stop_self_pool()
+        self.stop_self_query_pool()
+
+    def set_examples(self, examples_for_task: Optional[List[dict]] = None):
+        """Set the prefix to the provided examples.
+
+        Args:
+            examples_for_task (Optional[List[dict]], optional): Few-shot examples for the model to enhance model's ability. 
+                Defaults to :data:`None`.
+        """
+        if examples_for_task is None and self.examples_for_task is None:
+            self.prefix = ''
+        elif examples_for_task is not None:
+            eg_paris = []
+            for i in range(len(examples_for_task)):
+                eg_paris.append(
+                    self.get_detailed_example(
+                        self.examples_instruction_format,
+                        examples_for_task[i].get('instruct', self.query_instruction_for_retrieval),
+                        examples_for_task[i].get('query', ''),
+                        examples_for_task[i].get('response', '')
+                    )
+                )
+            self.prefix = '\n\n'.join(eg_paris) + '\n\n'
+        else:
+            eg_paris = []
+            for i in range(len(self.examples_for_task)):
+                eg_paris.append(
+                    self.get_detailed_example(
+                        self.examples_instruction_format,
+                        self.examples_for_task[i].get('instruct', self.query_instruction_for_retrieval),
+                        self.examples_for_task[i].get('query', ''),
+                        self.examples_for_task[i].get('response', '')
+                    )
+                )
+            self.prefix = '\n\n'.join(eg_paris) + '\n\n'
+
+    @staticmethod
+    def get_detailed_example(instruction_format: str, instruction: str, query: str, response: str):
+        """Combine the instruction and sentence along with the instruction format.
+
+        Args:
+            instruction_format (str): Format for instruction.
+            instruction (str): The text of instruction.
+            query (str): The text of example query.
+            response (str): The text of example response.
+
+        Returns:
+            str: The complete example following the given format.
+        """
+        if "\\n" in instruction_format:
+            instruction_format = instruction_format.replace("\\n", "\n")
+        return instruction_format.format(instruction, query, response)
+
+    def stop_self_query_pool(self):
+        if self.query_pool is not None:
+            self.stop_multi_process_pool(self.query_pool)
+            self.query_pool = None
+        try:
+            self.model.to('cpu')
+            torch.cuda.empty_cache()
+        except:
+            pass
+        gc.collect()
+
+    def encode_queries(
+        self,
+        queries: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the queries.
+
+        Args:
+            queries (Union[List[str], str]): Input queries to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
+        """
+        if batch_size is None: batch_size = self.batch_size
+        if max_length is None: max_length = self.query_max_length
+        if convert_to_numpy is None: convert_to_numpy = self.convert_to_numpy
+        
+        if isinstance(queries, str) or len(self.target_devices) == 1:
+            return self.encode_queries_single_device(
+                queries,
+                batch_size=batch_size,
+                max_length=max_length,
+                convert_to_numpy=convert_to_numpy,
+                device=self.target_devices[0],
+                **kwargs
+            )
+
+        self.stop_self_pool()
+        if self.query_pool is None:
+            self.query_pool = self.start_multi_process_pool(ICLLLMEmbedder._encode_queries_multi_process_worker)
+        embeddings = self.encode_multi_process(
+            queries,
+            self.query_pool,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+        return embeddings
+
+    def encode_corpus(
+        self,
+        corpus: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the corpus.
+
+        Args:
+            corpus (Union[List[str], str]): Input corpus to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
+        """
+        self.stop_self_query_pool()
+        return super().encode_corpus(
+            corpus,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+
+    def encode(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the input sentences with the embedding model.
+
+        Args:
+            sentences (Union[List[str], str]): Input sentences to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
+        """
+        return super().encode(
+            sentences,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+
+    # adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L976
+    @staticmethod
+    def _encode_queries_multi_process_worker(
+        target_device: str, model: 'ICLLLMEmbedder', input_queue: Queue, results_queue: Queue
+    ) -> None:
+        """
+        Internal working process to encode sentences in multi-process setup
+        """
+        while True:
+            try:
+                chunk_id, sentences, kwargs = (
+                    input_queue.get()
+                )
+                embeddings = model.encode_queries_single_device(
+                    sentences,
+                    device=target_device,
+                    **kwargs
+                )
+
+                results_queue.put([chunk_id, embeddings])
+            except queue.Empty:
+                break
+
+    @torch.no_grad()
+    def encode_queries_single_device(
+        self,
+        queries: Union[List[str], str],
+        batch_size: int = 256,
+        max_length: int = 512,
+        convert_to_numpy: bool = True,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ):
+        """Encode queries by a single device.
+
+        Args:
+            queries (Union[List[str], str]): Input queries to encode.
+            batch_size (int, optional): Number of queries for each iter. Defaults to :data:`256`.
+            max_length (int, optional): Maximum length of tokens. Defaults to :data:`512`.
+            convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`True`.
+            device (Optional[str], optional): Device to use for encoding. Defaults to None.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
+        """
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu":
+            self.model.float()
+
+        self.model.to(device)
+        self.model.eval()
+
+        input_was_string = False
+        if isinstance(queries, str):
+            queries = [queries]
+            input_was_string = True
+
+        if self.query_instruction_for_retrieval is not None:
+            if isinstance(queries, str):
+                input_texts = self.get_detailed_instruct(self.query_instruction_format, self.query_instruction_for_retrieval, queries)
+            else:
+                input_texts = [self.get_detailed_instruct(self.query_instruction_format, self.query_instruction_for_retrieval, query) for query in queries]
+        else:
+            input_texts = queries
+
+        prefix_ids = self.tokenizer(self.prefix, add_special_tokens=False)['input_ids']
+        suffix_ids = self.tokenizer(self.suffix, add_special_tokens=False)['input_ids']
+
+        _len_1 = len(self.tokenizer('<s>', add_special_tokens=False)['input_ids'])
+        _len_2 = len(self.tokenizer(f'{self.suffix}</s>', add_special_tokens=False)['input_ids'])
+        new_max_length = (len(prefix_ids) + len(suffix_ids) + max_length + 8) // 8 * 8 + 8
+
+        # tokenize without padding to get the correct length
+        all_inputs = []
+        for start_index in trange(0, len(input_texts), batch_size, desc='pre tokenize',
+                                  disable=len(input_texts) < batch_size):
+            sentences_batch = input_texts[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer(
+                sentences_batch,
+                truncation=True,
+                max_length=max_length - _len_1 - _len_2,
+                add_special_tokens=False,
+                **kwargs
+            )
+            sentences_batch = self.tokenizer.batch_decode(inputs_batch['input_ids'])
+            for i in range(len(sentences_batch)):
+                sentences_batch[i] = self.prefix + sentences_batch[i] + self.suffix
+            inputs_batch = self.tokenizer(
+                sentences_batch,
+                truncation=True,
+                max_length=new_max_length,
+                **kwargs
+            )
+            inputs_batch = [{
+                k: inputs_batch[k][i] for k in inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            all_inputs.extend(inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
+        all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
+        sentences_sorted = [input_texts[i] for i in length_sorted_idx]
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                inputs_batch = self.tokenizer.pad(
+                    all_inputs_sorted[: batch_size],
+                    padding=True,
+                    return_tensors='pt',
+                    **kwargs
+                ).to(device)
+                last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+                embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        # encode
+        all_embeddings = []
+        for start_index in tqdm(range(0, len(sentences_sorted), batch_size), desc="Inference Embeddings",
+                                disable=len(sentences_sorted) < batch_size):
+            inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer.pad(
+                inputs_batch,
+                padding=True,
+                return_tensors='pt',
+                **kwargs
+            ).to(device)
+
+            last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+            embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
+            embeddings = self._truncate_embeddings(embeddings)
+            if self.normalize_embeddings:
+                embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
+            embeddings = cast(torch.Tensor, embeddings)
+
+            if convert_to_numpy:
+                embeddings = self._convert_to_numpy(embeddings, device=device)
+            all_embeddings.append(embeddings)
+
+        if convert_to_numpy:
+            all_embeddings = np.concatenate(all_embeddings, axis=0)
+        else:
+            all_embeddings = torch.cat(all_embeddings, dim=0)
+
+        # adjust the order of embeddings
+        all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
+
+        # return the embeddings
+        if input_was_string:
+            return all_embeddings[0]
+        return all_embeddings
+
+    @torch.no_grad()
+    def encode_single_device(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: int = 256,
+        max_length: int = 512,
+        convert_to_numpy: bool = True,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ):
+        """Encode input sentences by a single device.
+
+        Args:
+            sentences (Union[List[str], str]): Input sentences to encode.
+            batch_size (int, optional): Number of sentences for each iter. Defaults to :data:`256`.
+            max_length (int, optional): Maximum length of tokens. Defaults to :data:`512`.
+            convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`True`.
+            device (Optional[str], optional): Device to use for encoding. Defaults to None.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
+        """
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu":
+            self.model.float()
+
+        self.model.to(device)
+        self.model.eval()
+
+        input_was_string = False
+        if isinstance(sentences, str):
+            sentences = [sentences]
+            input_was_string = True
+
+        # tokenize without padding to get the correct length
+        all_inputs = []
+        for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
+                                  disable=len(sentences) < batch_size):
+            sentences_batch = sentences[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer(
+                sentences_batch,
+                truncation=True,
+                max_length=max_length,
+                **kwargs
+            )
+            inputs_batch = [{
+                k: inputs_batch[k][i] for k in inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            all_inputs.extend(inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
+        all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                inputs_batch = self.tokenizer.pad(
+                    all_inputs_sorted[: batch_size],
+                    padding=True,
+                    return_tensors='pt',
+                    **kwargs
+                ).to(device)
+                last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+                embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        # encode
+        all_embeddings = []
+        for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
+                                disable=len(sentences) < batch_size):
+            inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer.pad(
+                inputs_batch,
+                padding=True,
+                return_tensors='pt',
+                **kwargs
+            ).to(device)
+            last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+            embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
+            embeddings = self._truncate_embeddings(embeddings)
+            if self.normalize_embeddings:
+                embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
+            embeddings = cast(torch.Tensor, embeddings)
+
+            if convert_to_numpy:
+                embeddings = self._convert_to_numpy(embeddings, device=device)
+            all_embeddings.append(embeddings)
+
+        if convert_to_numpy:
+            all_embeddings = np.concatenate(all_embeddings, axis=0)
+        else:
+            all_embeddings = torch.cat(all_embeddings, dim=0)
+
+        # adjust the order of embeddings
+        all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
+
+        # return the embeddings
+        if input_was_string:
+            return all_embeddings[0]
+        return all_embeddings

FlagEmbedding/inference/embedder/decoder_only/pseudo_moe.py

@@ -0,0 +1,194 @@
+from typing import cast, Any, List, Union, Optional
+
+import torch
+import numpy as np
+
+from .base import BaseLLMEmbedder, last_token_pool
+
+
+class PseudoMoELLMEmbedder(BaseLLMEmbedder):
+    """Decoder-only embedder for pseudo MoE checkpoints.
+
+    This class follows the same behavior as :class:`BaseLLMEmbedder`, but supports
+    selecting an active domain (e.g. ``general``, ``coding``, ``reasoning``) during
+    inference when the underlying model implements domain routing.
+
+    Args:
+        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.
+        normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
+        use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+            degradation. Defaults to :data:`True`.
+        query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
+            with :attr:`query_instruction_format`. Defaults to :data:`None`.
+        query_instruction_format (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"Instruct: {}\nQuery: {}"`.
+        devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
+        trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
+        query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
+        passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
+        convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor. 
+            Defaults to :data:`True`.
+        domain_for_pseudo_moe (str, optional): Specifies the active domain for the decoder-only pseudo-MoE model (e.g., "general", "coding", or "reasoning").
+            Defaults to "general".
+    
+    Attributes:
+        DEFAULT_POOLING_METHOD: The default pooling method when running the model.
+    """
+    DEFAULT_POOLING_METHOD = "last_token"
+    
+    def __init__(
+        self,
+        model_name_or_path: str,
+        normalize_embeddings: bool = True,
+        use_fp16: bool = False,
+        use_bf16: bool = True,
+        query_instruction_for_retrieval: Optional[str] = None,
+        query_instruction_format: str = "Instruct: {}\nQuery: {}",
+        devices: Optional[Union[str, List[str]]] = None,
+        trust_remote_code: bool = True,
+        cache_dir: Optional[str] = None,
+        batch_size: int = 256,
+        query_max_length: int = 512,
+        passage_max_length: int = 512,
+        convert_to_numpy: bool = True,
+        truncate_dim: Optional[int] = None,
+        domain_for_pseudo_moe: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        self.domain_for_pseudo_moe = domain_for_pseudo_moe
+        super().__init__(
+            model_name_or_path=model_name_or_path,
+            normalize_embeddings=normalize_embeddings,
+            use_fp16=use_fp16,
+            use_bf16=use_bf16,
+            query_instruction_for_retrieval=query_instruction_for_retrieval,
+            query_instruction_format=query_instruction_format,
+            devices=devices,
+            trust_remote_code=trust_remote_code,
+            cache_dir=cache_dir,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            passage_max_length=passage_max_length,
+            convert_to_numpy=convert_to_numpy,
+            truncate_dim=truncate_dim,
+            **kwargs,
+        )
+
+    def _resolve_domain(self, kwargs: Any) -> Optional[str]:
+        domain = kwargs.pop("domain_for_pseudo_moe", None)
+        if domain is None:
+            domain = kwargs.pop("domain", None)
+        if domain is None:
+            domain = self.domain_for_pseudo_moe
+        return domain
+
+    @torch.no_grad()
+    def encode_single_device(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: int = 256,
+        max_length: int = 512,
+        convert_to_numpy: bool = True,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ):
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu":
+            self.model.float()
+
+        self.model.to(device)
+        self.model.eval()
+
+        input_was_string = False
+        if isinstance(sentences, str):
+            sentences = [sentences]
+            input_was_string = True
+
+        domain = self._resolve_domain(kwargs)
+        if domain is not None and hasattr(self.model, "set_domain"):
+            self.model.set_domain(domain)
+
+        model_forward_kwargs = {"return_dict": True}
+        if domain is not None:
+            model_forward_kwargs["domain"] = domain
+
+        # tokenize without padding to get the correct length
+        all_inputs = []
+        for start_index in range(0, len(sentences), batch_size):
+            sentences_batch = sentences[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer(
+                sentences_batch,
+                truncation=True,
+                max_length=max_length,
+                **kwargs
+            )
+            inputs_batch = [{
+                k: inputs_batch[k][i] for k in inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            all_inputs.extend(inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
+        all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                inputs_batch = self.tokenizer.pad(
+                    all_inputs_sorted[: batch_size],
+                    padding=True,
+                    return_tensors='pt',
+                    **kwargs
+                ).to(device)
+                try:
+                    last_hidden_state = self.model(**inputs_batch, **model_forward_kwargs).last_hidden_state
+                except TypeError:
+                    last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+                _ = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
+                flag = True
+            except RuntimeError:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError:
+                batch_size = batch_size * 3 // 4
+
+        # encode
+        all_embeddings = []
+        for start_index in range(0, len(sentences), batch_size):
+            inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer.pad(
+                inputs_batch,
+                padding=True,
+                return_tensors='pt',
+                **kwargs
+            ).to(device)
+            try:
+                last_hidden_state = self.model(**inputs_batch, **model_forward_kwargs).last_hidden_state
+            except TypeError:
+                last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+            embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
+            embeddings = self._truncate_embeddings(embeddings)
+            embeddings = torch.nan_to_num(embeddings, nan=0.0, posinf=1e4, neginf=-1e4)
+            if self.normalize_embeddings:
+                embeddings = torch.nn.functional.normalize(embeddings.float(), dim=-1)
+            embeddings = cast(torch.Tensor, embeddings)
+
+            if convert_to_numpy:
+                embeddings = self._convert_to_numpy(embeddings, device=device)
+            all_embeddings.append(embeddings)
+
+        if convert_to_numpy:
+            all_embeddings = np.concatenate(all_embeddings, axis=0)
+        else:
+            all_embeddings = torch.cat(all_embeddings, dim=0)
+
+        # adjust the order of embeddings
+        all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
+
+        if input_was_string:
+            return all_embeddings[0]
+        return all_embeddings

FlagEmbedding/inference/embedder/encoder_only/__init__.py

@@ -0,0 +1,7 @@
+from .base import BaseEmbedder as FlagModel
+from .m3 import M3Embedder as BGEM3FlagModel
+
+__all__ = [
+    "FlagModel",
+    "BGEM3FlagModel",
+]

FlagEmbedding/inference/embedder/encoder_only/base.py

@@ -0,0 +1,308 @@
+from tqdm import tqdm, trange
+from typing import cast, Any, List, Union, Optional
+
+import torch
+import numpy as np
+from transformers import AutoModel, AutoTokenizer
+
+from FlagEmbedding.abc.inference import AbsEmbedder
+
+
+class BaseEmbedder(AbsEmbedder):
+    """
+    Base embedder for encoder only models.
+
+    Args:
+        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.
+        normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
+        use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+            degradation. Defaults to :data:`True`.
+        query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
+            with :attr:`query_instruction_format`. Defaults to :data:`None`.
+        query_instruction_format (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"{}{}"`.
+        devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
+        pooling_method (str, optional): Pooling method to get embedding vector from the last hidden state. Defaults to :data:`"cls"`.
+        trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
+        query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
+        passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
+        convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor. 
+            Defaults to :data:`True`.
+    
+    Attributes:
+        DEFAULT_POOLING_METHOD: The default pooling method when running the model.
+    """
+    
+    DEFAULT_POOLING_METHOD = "cls"
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        normalize_embeddings: bool = True,
+        use_fp16: bool = True,
+        use_bf16: bool = False,
+        query_instruction_for_retrieval: Optional[str] = None,
+        query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_retrieval
+        devices: Optional[Union[str, List[str]]] = None, # specify devices, such as "cuda:0" or ["cuda:0", "cuda:1"]
+        # Additional parameters for BaseEmbedder
+        pooling_method: str = "cls",
+        trust_remote_code: bool = False,
+        cache_dir: Optional[str] = None,
+        # inference
+        batch_size: int = 256,
+        query_max_length: int = 512,
+        passage_max_length: int = 512,
+        convert_to_numpy: bool = True,
+        truncate_dim: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_name_or_path,
+            normalize_embeddings=normalize_embeddings,
+            use_fp16=use_fp16,
+            use_bf16=use_bf16,
+            query_instruction_for_retrieval=query_instruction_for_retrieval,
+            query_instruction_format=query_instruction_format,
+            devices=devices,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            passage_max_length=passage_max_length,
+            convert_to_numpy=convert_to_numpy,
+            truncate_dim=truncate_dim,
+            **kwargs
+        )
+        self.pooling_method = pooling_method
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            cache_dir=cache_dir
+        )
+        self.model = AutoModel.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            cache_dir=cache_dir,
+            dtype=self.get_model_torch_dtype(),
+        )
+
+    def encode_queries(
+        self,
+        queries: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the queries.
+
+        Args:
+            queries (Union[List[str], str]): Input queries to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
+        """
+        return super().encode_queries(
+            queries,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+
+    def encode_corpus(
+        self,
+        corpus: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the corpus using the instruction if provided.
+
+        Args:
+            corpus (Union[List[str], str]): Input corpus to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
+        """
+        return super().encode_corpus(
+            corpus,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+
+    def encode(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        convert_to_numpy: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Encode the input sentences with the embedding model.
+
+        Args:
+            sentences (Union[List[str], str]): Input sentences to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`None`.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
+        """
+        return super().encode(
+            sentences,
+            batch_size=batch_size,
+            max_length=max_length,
+            convert_to_numpy=convert_to_numpy,
+            **kwargs
+        )
+
+    @torch.no_grad()
+    def encode_single_device(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: int = 256,
+        max_length: int = 512,
+        convert_to_numpy: bool = True,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ):
+        """Encode input sentences by a single device.
+
+        Args:
+            sentences (Union[List[str], str]): Input sentences to encode.
+            batch_size (int, optional): Number of sentences for each iter. Defaults to :data:`256`.
+            max_length (int, optional): Maximum length of tokens. Defaults to :data:`512`.
+            convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will 
+                be a Torch Tensor. Defaults to :data:`True`.
+            device (Optional[str], optional): Device to use for encoding. Defaults to None.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
+        """
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu":
+            self.model.float()
+
+        self.model.to(device)
+        self.model.eval()
+
+        input_was_string = False
+        if isinstance(sentences, str):
+            sentences = [sentences]
+            input_was_string = True
+
+        # tokenize without padding to get the correct length
+        all_inputs = []
+        for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
+                                  disable=len(sentences) < batch_size):
+            sentences_batch = sentences[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer(
+                sentences_batch,
+                truncation=True,
+                max_length=max_length,
+                **kwargs
+            )
+            inputs_batch = [{
+                k: inputs_batch[k][i] for k in inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            all_inputs.extend(inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
+        all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                inputs_batch = self.tokenizer.pad(
+                    all_inputs_sorted[: batch_size],
+                    padding=True,
+                    return_tensors='pt',
+                    **kwargs
+                ).to(device)
+                last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+                embeddings = self.pooling(last_hidden_state, inputs_batch['attention_mask'])
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        # encode
+        all_embeddings = []
+        for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
+                                disable=len(sentences) < batch_size):
+            inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer.pad(
+                inputs_batch,
+                padding=True,
+                return_tensors='pt',
+                **kwargs
+            ).to(device)
+            last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
+            embeddings = self.pooling(last_hidden_state, inputs_batch['attention_mask'])
+            embeddings = self._truncate_embeddings(embeddings)
+            if self.normalize_embeddings:
+                embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
+            embeddings = cast(torch.Tensor, embeddings)
+
+            if convert_to_numpy:
+                embeddings = self._convert_to_numpy(embeddings, device=device)
+            all_embeddings.append(embeddings)
+
+        if convert_to_numpy:
+            all_embeddings = np.concatenate(all_embeddings, axis=0)
+        else:
+            all_embeddings = torch.cat(all_embeddings, dim=0)
+
+        # adjust the order of embeddings
+        all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
+
+        # return the embeddings
+        if input_was_string:
+            return all_embeddings[0]
+        return all_embeddings
+
+    def pooling(
+        self,
+        last_hidden_state: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None
+    ):
+        """The pooling function.
+
+        Args:
+            last_hidden_state (torch.Tensor): The last hidden state of the model.
+            attention_mask (Optional[torch.Tensor], optional): Attention mask. Defaults to :data:`None`.
+
+        Raises:
+            NotImplementedError: pooling method not implemented.
+
+        Returns:
+            torch.Tensor: The embedding vectors after pooling.
+        """
+        if self.pooling_method == 'cls':
+            return last_hidden_state[:, 0]
+        elif self.pooling_method == 'mean':
+            s = torch.sum(last_hidden_state * attention_mask.unsqueeze(-1).float(), dim=1)
+            d = attention_mask.sum(dim=1, keepdim=True).float()
+            return s / d
+        else:
+            raise NotImplementedError(f"pooling method {self.pooling_method} not implemented")

FlagEmbedding/inference/embedder/encoder_only/m3.py

@@ -0,0 +1,792 @@
+import math
+import torch
+import queue
+import logging
+import numpy as np
+from tqdm import tqdm, trange
+from multiprocessing import Queue
+from collections import defaultdict
+from transformers import AutoTokenizer
+from typing import Any, List, Union, Dict, Literal, Tuple, Optional
+
+from FlagEmbedding.abc.inference import AbsEmbedder
+from FlagEmbedding.finetune.embedder.encoder_only.m3 import (
+    EncoderOnlyEmbedderM3ModelForInference, EncoderOnlyEmbedderM3Runner
+)
+
+logger = logging.getLogger(__name__)
+
+
+class M3Embedder(AbsEmbedder):
+    """ 
+    Embedder class for BGE-M3.
+
+    Args:
+        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.
+        normalize_embeddings (bool, optional): If True, normalize the dense embedding vector. Defaults to :data:`True`.
+        use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+            degradation. Defaults to :data:`True`.
+        query_instruction_for_retrieval: (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
+            with :attr:`query_instruction_format`. Defaults to :data:`None`.
+        query_instruction_format: (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"{}{}"`.
+        devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
+        pooling_method (str, optional): Pooling method to get embedding vector from the last hidden state. Defaults to :data:`"cls"`.
+        trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        cobert_dim (int, optional): Dimension of colbert linear. Return the hidden_size if -1. Defaults to :data:`-1`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
+        query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
+        passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
+        return_dense (bool, optional): If true, will return the dense embedding. Defaults to :data:`True`.
+        return_sparse (bool, optional): If true, will return the sparce embedding. Defaults to :data:`False`.
+        return_colbert_vecs (bool, optional): If true, will return the colbert vectors. Defaults to :data:`False`.
+        
+    Attributes:
+        DEFAULT_POOLING_METHOD: The default pooling method when running the model.
+    """
+    DEFAULT_POOLING_METHOD = "cls"
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        normalize_embeddings: bool = True,
+        use_fp16: bool = True,
+        use_bf16: bool = False,
+        query_instruction_for_retrieval: Optional[str] = None,
+        query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_retrieval
+        devices: Optional[Union[str, List[str]]] = None, # specify devices, such as "cuda:0" or ["cuda:0", "cuda:1"]
+        # Additional parameters for M3Embedder
+        pooling_method: str = "cls",
+        trust_remote_code: bool = False,
+        cache_dir: Optional[str] = None,
+        colbert_dim: int = -1,
+        # inference
+        batch_size: int = 256,
+        query_max_length: int = 512,
+        passage_max_length: int = 512,
+        return_dense: bool = True,
+        return_sparse: bool = False,
+        return_colbert_vecs: bool = False,
+        truncate_dim: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_name_or_path,
+            normalize_embeddings=normalize_embeddings,
+            use_fp16=use_fp16,
+            use_bf16=use_bf16,
+            query_instruction_for_retrieval=query_instruction_for_retrieval,
+            query_instruction_format=query_instruction_format,
+            devices=devices,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            passage_max_length=passage_max_length,
+            return_dense=return_dense,
+            return_sparse=return_sparse,
+            return_colbert_vecs=return_colbert_vecs,
+            truncate_dim=truncate_dim,
+            **kwargs
+        )
+        self.pooling_method = pooling_method
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            cache_dir=cache_dir
+        )
+        self.model = EncoderOnlyEmbedderM3ModelForInference(
+            EncoderOnlyEmbedderM3Runner.get_model(
+                model_name_or_path,
+                trust_remote_code=trust_remote_code,
+                colbert_dim=colbert_dim,
+                cache_dir=cache_dir,
+                torch_dtype=self.get_model_torch_dtype(),
+            ),
+            tokenizer=self.tokenizer,
+            sentence_pooling_method=pooling_method,
+            normalize_embeddings=normalize_embeddings
+        )
+
+    def convert_id_to_token(self, lexical_weights: List[Dict]):
+        """Convert the ids back to tokens.
+
+        Args:
+            lexical_weights (List[Dict]): A list of dictionaries of id & weights.
+
+        Returns:
+            List[Dict]: A list of dictionaries of tokens & weights.
+        """
+        if isinstance(lexical_weights, dict):
+            lexical_weights = [lexical_weights]
+        new_lexical_weights = []
+        for item in lexical_weights:
+            new_item = {}
+            for id, weight in item.items():
+                token = self.tokenizer.decode([int(id)])
+                new_item[token] = weight
+            new_lexical_weights.append(new_item)
+
+        if len(new_lexical_weights) == 1:
+            new_lexical_weights = new_lexical_weights[0]
+        return new_lexical_weights
+
+    def compute_lexical_matching_score(
+        self,
+        lexical_weights_1: Union[Dict[str, float], List[Dict[str, float]]],
+        lexical_weights_2: Union[Dict[str, float], List[Dict[str, float]]]
+    ) -> Union[np.ndarray, float]:
+        """Compute the laxical matching score of two given lexical weights.
+
+        Args:
+            lexical_weights_1 (Union[Dict[str, float], List[Dict[str, float]]]): First array of lexical weights.
+            lexical_weights_2 (Union[Dict[str, float], List[Dict[str, float]]]): Second array of lexical weights.
+
+        Returns:
+            Union[np.ndarray, float]: The computed lexical weights across the two arries of lexical weights.
+        """
+        def _compute_single_lexical_matching_score(lw1: Dict[str, float], lw2: Dict[str, float]):
+            scores = 0
+            for token, weight in lw1.items():
+                if token in lw2:
+                    scores += weight * lw2[token]
+            return scores
+
+        if isinstance(lexical_weights_1, dict) and isinstance(lexical_weights_2, dict):
+            return _compute_single_lexical_matching_score(lexical_weights_1, lexical_weights_2)
+        elif isinstance(lexical_weights_1, list) and isinstance(lexical_weights_2, list):
+            scores_array = []
+            for lw1 in lexical_weights_1:
+                scores_array.append([
+                    _compute_single_lexical_matching_score(lw1, lw2)
+                    for lw2 in lexical_weights_2
+                ])
+            return np.array(scores_array)
+        else:
+            raise ValueError("The input format of lexical_weights is not correct.")
+
+    def colbert_score(self, q_reps, p_reps):
+        """Compute colbert scores of input queries and passages.
+
+        Args:
+            q_reps (np.ndarray): Multi-vector embeddings for queries.
+            p_reps (np.ndarray): Multi-vector embeddings for passages/corpus.
+
+        Returns:
+            torch.Tensor: Computed colbert scores.
+        """
+        q_reps, p_reps = torch.from_numpy(q_reps), torch.from_numpy(p_reps)
+        token_scores = torch.einsum('in,jn->ij', q_reps, p_reps)
+        scores, _ = token_scores.max(-1)
+        scores = torch.sum(scores) / q_reps.size(0)
+        return scores
+
+    def encode_queries(
+        self,
+        queries: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        return_dense: Optional[bool] = None,
+        return_sparse: Optional[bool] = None,
+        return_colbert_vecs: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Dict[
+        Literal["dense_vecs", "lexical_weights", "colbert_vecs"],
+        Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
+    ]:
+        """Encode the queries using the specified way.
+
+        Args:
+            queries (Union[List[str], str]): The input queries to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            return_dense (Optional[bool], optional): If True, compute and return dense embedding. Defaults to :data:`None`.
+            return_sparse (Optional[bool], optional): If True, compute and return sparce embedding. Defaults to :data:`None`.
+            return_colbert_vecs (Optional[bool], optional): If True, compute and return cobert vectors. Defaults to :data:`None`.
+
+        Returns:
+            Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
+        """
+        if batch_size is None: batch_size = self.batch_size
+        if max_length is None: max_length = self.query_max_length
+        if return_dense is None: return_dense = self.return_dense
+        if return_sparse is None: return_sparse = self.return_sparse
+        if return_colbert_vecs is None: return_colbert_vecs = self.return_colbert_vecs
+
+        return super().encode_queries(
+            queries,
+            batch_size=batch_size,
+            max_length=max_length,
+            return_dense=return_dense,
+            return_sparse=return_sparse,
+            return_colbert_vecs=return_colbert_vecs,
+            **kwargs
+        )
+
+    def encode_corpus(
+        self,
+        corpus: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        return_dense: Optional[bool] = None,
+        return_sparse: Optional[bool] = None,
+        return_colbert_vecs: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Dict[
+        Literal["dense_vecs", "lexical_weights", "colbert_vecs"],
+        Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
+    ]:
+        """Encode the corpus using the specified way.
+
+        Args:
+            corpus (Union[List[str], str]): The input corpus to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            return_dense (Optional[bool], optional): If True, compute and return dense embedding. Defaults to :data:`None`.
+            return_sparse (Optional[bool], optional): If True, compute and return sparce embedding. Defaults to :data:`None`.
+            return_colbert_vecs (Optional[bool], optional): If True, compute and return cobert vectors. Defaults to :data:`None`.
+
+        Returns:
+            Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
+        """
+        if batch_size is None: batch_size = self.batch_size
+        if max_length is None: max_length = self.passage_max_length
+        if return_dense is None: return_dense = self.return_dense
+        if return_sparse is None: return_sparse = self.return_sparse
+        if return_colbert_vecs is None: return_colbert_vecs = self.return_colbert_vecs
+
+        return super().encode_corpus(
+            corpus,
+            batch_size=batch_size,
+            max_length=max_length,
+            return_dense=return_dense,
+            return_sparse=return_sparse,
+            return_colbert_vecs=return_colbert_vecs,
+            **kwargs
+        )
+
+    def encode(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        return_dense: Optional[bool] = None,
+        return_sparse: Optional[bool] = None,
+        return_colbert_vecs: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Dict[
+        Literal["dense_vecs", "lexical_weights", "colbert_vecs"],
+        Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
+    ]:
+        """Encode the sentences using the specified way.
+
+        Args:
+            sentences (Union[List[str], str]): The input sentences to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            return_dense (Optional[bool], optional): If True, compute and return dense embedding. Defaults to :data:`None`.
+            return_sparse (Optional[bool], optional): If True, compute and return sparce embedding. Defaults to :data:`None`.
+            return_colbert_vecs (Optional[bool], optional): If True, compute and return cobert vectors. Defaults to :data:`None`.
+
+        Returns:
+            Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
+        """
+        if batch_size is None: batch_size = self.batch_size
+        if max_length is None: max_length = self.passage_max_length
+        if return_dense is None: return_dense = self.return_dense
+        if return_sparse is None: return_sparse = self.return_sparse
+        if return_colbert_vecs is None: return_colbert_vecs = self.return_colbert_vecs
+
+        return super().encode(
+            sentences,
+            batch_size=batch_size,
+            max_length=max_length,
+            return_dense=return_dense,
+            return_sparse=return_sparse,
+            return_colbert_vecs=return_colbert_vecs,
+            **kwargs
+        )
+
+    @torch.no_grad()
+    def encode_single_device(
+        self,
+        sentences: Union[List[str], str],
+        batch_size: int = 256,
+        max_length: int = 512,
+        return_dense: bool = True,
+        return_sparse: bool = False,
+        return_colbert_vecs: bool = False,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ):
+        """Using single device to encode the input sentences.
+
+        Args:
+            sentences (Union[List[str], str]): The input sentences to encode.
+            batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`256`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`512`.
+            return_dense (Optional[bool], optional): If True, compute and return dense embedding. Defaults to :data:`True`.
+            return_sparse (Optional[bool], optional): If True, compute and return sparce embedding. Defaults to :data:`False`.
+            return_colbert_vecs (Optional[bool], optional): If True, compute and return cobert vectors. Defaults to :data:`False`.
+            device (Optional[str], optional): _description_. Defaults to :data:`None`.
+
+        Returns:
+            Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
+        """
+        # pop convert_to_numpy from kwargs
+        kwargs.pop("convert_to_numpy", None)
+
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu":
+            self.model.float()
+
+        self.model.to(device)
+        self.model.eval()
+
+        input_was_string = False
+        if isinstance(sentences, str):
+            sentences = [sentences]
+            input_was_string = True
+
+        def _process_token_weights(token_weights: np.ndarray, input_ids: list):
+            # conver to dict
+            result = defaultdict(int)
+            unused_tokens = set()
+            for _token in ['cls_token', 'eos_token', 'pad_token', 'unk_token']:
+                if _token in self.tokenizer.special_tokens_map:
+                    _token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.special_tokens_map[_token])
+                    unused_tokens.add(_token_id)
+            # token_weights = np.ceil(token_weights * 100)
+            for w, idx in zip(token_weights, input_ids):
+                if idx not in unused_tokens and w > 0:
+                    idx = str(idx)
+                    # w = int(w)
+                    if w > result[idx]:
+                        result[idx] = w
+            return result
+
+        def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
+            # delte the vectors of padding tokens
+            tokens_num = np.sum(attention_mask)
+            return colbert_vecs[:tokens_num - 1]  # we don't use the embedding of cls, so select tokens_num-1
+
+        # tokenize without padding to get the correct length
+        all_inputs = []
+        for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
+                                  disable=len(sentences) < batch_size):
+            sentences_batch = sentences[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer(
+                sentences_batch,
+                truncation=True,
+                max_length=max_length,
+                **kwargs
+            )
+            inputs_batch = [{
+                k: inputs_batch[k][i] for k in inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            all_inputs.extend(inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
+        all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                inputs_batch = self.tokenizer.pad(
+                    all_inputs_sorted[: batch_size],
+                    padding=True,
+                    return_tensors='pt',
+                    **kwargs
+                ).to(device)
+                outputs = self.model(
+                    inputs_batch,
+                    return_dense=return_dense,
+                    return_sparse=return_sparse,
+                    return_colbert_vecs=return_colbert_vecs
+                )
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        # encode
+        all_dense_embeddings, all_lexical_weights, all_colbert_vecs = [], [], []
+        for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
+                                disable=len(sentences) < batch_size):
+            inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
+            inputs_batch = self.tokenizer.pad(
+                inputs_batch,
+                padding=True,
+                return_tensors='pt',
+                **kwargs
+            ).to(device)
+            outputs = self.model(
+                inputs_batch,
+                return_dense=return_dense,
+                return_sparse=return_sparse,
+                return_colbert_vecs=return_colbert_vecs,
+                truncate_dim=self.truncate_dim
+            )
+
+            if return_dense:
+                all_dense_embeddings.append(self._convert_to_numpy(outputs['dense_vecs'], device=device))
+
+            if return_sparse:
+                token_weights = outputs['sparse_vecs'].squeeze(-1)
+                all_lexical_weights.extend(
+                    list(map(
+                        _process_token_weights, 
+                        self._convert_to_numpy(token_weights, device=device),
+                        self._convert_to_numpy(inputs_batch['input_ids'], device=device).tolist()
+                )))
+
+            if return_colbert_vecs:
+                all_colbert_vecs.extend(
+                    list(map(
+                        _process_colbert_vecs,
+                        self._convert_to_numpy(outputs['colbert_vecs'], device=device),
+                        self._convert_to_numpy(inputs_batch['attention_mask'], device=device)
+                )))
+
+        if return_dense:
+            all_dense_embeddings = np.concatenate(all_dense_embeddings, axis=0)
+            # adjust the order of embeddings
+            all_dense_embeddings = all_dense_embeddings[np.argsort(length_sorted_idx)]
+            if input_was_string:
+                all_dense_embeddings = all_dense_embeddings[0]
+        else:
+            all_dense_embeddings = None
+
+        if return_sparse:
+            # adjust the order of lexical weights
+            all_lexical_weights = [all_lexical_weights[i] for i in np.argsort(length_sorted_idx)]
+            if input_was_string:
+                all_lexical_weights = all_lexical_weights[0]
+        else:
+            all_lexical_weights = None
+
+        if return_colbert_vecs:
+            # adjust the order of embeddings
+            all_colbert_vecs = [all_colbert_vecs[i] for i in np.argsort(length_sorted_idx)]
+            if input_was_string:
+                all_colbert_vecs = all_colbert_vecs[0]
+        else:
+            all_colbert_vecs = None
+
+        # return the embeddings
+        return {
+            "dense_vecs": all_dense_embeddings,
+            "lexical_weights": all_lexical_weights,
+            "colbert_vecs": all_colbert_vecs
+        }
+
+    def compute_score(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: Optional[int] = None,
+        max_query_length: Optional[int] = None,
+        max_passage_length: Optional[int] = None,
+        weights_for_different_modes: Optional[List[float]] = None,
+        **kwargs: Any
+    ) -> Dict[
+        Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"],
+        List[float]
+    ]:
+        """Compute the relevance score of different attributes.
+
+        Args:
+            sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]): _description_
+            batch_size (Optional[int], optional): _description_. Defaults to None.
+            max_query_length (Optional[int], optional): _description_. Defaults to None.
+            max_passage_length (Optional[int], optional): _description_. Defaults to None.
+            weights_for_different_modes (Optional[List[float]], optional): _description_. Defaults to None.
+
+        Returns:
+            Dict[Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"], List[float]]
+        """
+        if batch_size is None: batch_size = self.batch_size
+        if max_query_length is None: max_query_length = self.query_max_length
+        if max_passage_length is None: max_passage_length = self.passage_max_length
+
+        if len(self.target_devices) == 1:
+            return self.compute_score_single_device(
+                sentence_pairs,
+                batch_size=batch_size,
+                max_query_length=max_query_length,
+                max_passage_length=max_passage_length,
+                weights_for_different_modes=weights_for_different_modes,
+                device=self.target_devices[0],
+                **kwargs
+            )
+
+        pool = self.start_multi_process_pool(M3Embedder._compute_score_multi_process_worker)
+        embeddings = self.compute_score_multi_process(
+            sentence_pairs,
+            pool,
+            batch_size=batch_size,
+            max_query_length=max_query_length,
+            max_passage_length=max_passage_length,
+            weights_for_different_modes=weights_for_different_modes,
+            **kwargs
+        )
+        self.stop_multi_process_pool(pool)
+        return embeddings
+
+    # adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L877
+    def compute_score_multi_process(
+        self,
+        sentence_pairs: List[Tuple[str, str]],
+        pool: Dict[Literal["input", "output", "processes"], Any],
+        **kwargs
+    ):
+        chunk_size = math.ceil(len(sentence_pairs) / len(pool["processes"]))
+
+        input_queue = pool["input"]
+        last_chunk_id = 0
+        chunk = []
+
+        for sentence_pair in sentence_pairs:
+            chunk.append(sentence_pair)
+            if len(chunk) >= chunk_size:
+                input_queue.put(
+                    [last_chunk_id, chunk, kwargs]
+                )
+                last_chunk_id += 1
+                chunk = []
+
+        if len(chunk) > 0:
+            input_queue.put([last_chunk_id, chunk, kwargs])
+            last_chunk_id += 1
+
+        output_queue = pool["output"]
+        results_list = sorted(
+            [output_queue.get() for _ in trange(last_chunk_id, desc="Chunks")],
+            key=lambda x: x[0],
+        )
+
+        scores_dict = self._concatenate_compute_score_results_from_multi_process([result[1] for result in results_list])
+        return scores_dict
+
+    # adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L976
+    @staticmethod
+    def _compute_score_multi_process_worker(
+        target_device: str, model: 'M3Embedder', input_queue: Queue, results_queue: Queue
+    ) -> None:
+        """
+        Internal working process to encode sentences in multi-process setup
+        """
+        while True:
+            try:
+                chunk_id, sentences, kwargs = (
+                    input_queue.get()
+                )
+                embeddings = model.compute_score_single_device(
+                    sentences,
+                    device=target_device,
+                    **kwargs
+                )
+
+                results_queue.put([chunk_id, embeddings])
+            except queue.Empty:
+                break
+
+    @torch.no_grad()
+    def compute_score_single_device(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: int = 256,
+        max_query_length: int = 512,
+        max_passage_length: int = 512,
+        weights_for_different_modes: Optional[List[float]] = None,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ) -> Dict[
+        Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"],
+        List[float]
+    ]:
+        """Compute the relevance score of different attributes.
+
+        Args:
+            sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]): Pairs of sentences to compute the score.
+            batch_size (Optional[int], optional): _description_. Defaults to :data:`None`.
+            max_query_length (Optional[int], optional): _description_. Defaults to :data:`None`.
+            max_passage_length (Optional[int], optional): _description_. Defaults to :data:`None`.
+            weights_for_different_modes (Optional[List[float]], optional): The weights for different methods. Defaults to :data:`None`.
+            device (Optional[str], optional): The device to use. Defaults to :data:`None`.
+
+        Returns:
+            Dict[Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"], List[float]]
+        """
+        def _tokenize(texts: list, max_length: int):
+            return self.tokenizer(
+                texts,
+                max_length=max_length,
+                padding=True,
+                return_token_type_ids=False,
+                truncation=True,
+                return_tensors='pt',
+                **kwargs
+            )
+
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu":
+            self.model.float()
+
+        self.model.to(device)
+        self.model.eval()
+
+        if isinstance(sentence_pairs, list) and len(sentence_pairs) == 0:
+            return []
+        if isinstance(sentence_pairs[0], str):
+            one_input_pair = True
+            sentence_pairs = [sentence_pairs]
+        else:
+            one_input_pair = False
+
+        all_scores = {
+            'colbert': [],
+            'sparse': [],
+            'dense': [],
+            'sparse+dense': [],
+            'colbert+sparse+dense': []
+        }
+        for start_index in tqdm(range(0, len(sentence_pairs), batch_size), desc="Compute Scores",
+                                disable=len(sentence_pairs) < batch_size):
+            sentences_batch = sentence_pairs[start_index:start_index + batch_size]
+
+            queries_batch = [pair[0] for pair in sentences_batch]
+            corpus_batch = [pair[1] for pair in sentences_batch]
+
+            queries_inputs = _tokenize(queries_batch, max_length=max_query_length).to(device)
+            corpus_inputs = _tokenize(corpus_batch, max_length=max_passage_length).to(device)
+
+            queries_output = self.model(
+                queries_inputs,
+                return_dense=True, return_sparse=True, return_colbert_vecs=True,
+                return_sparse_embedding=True
+            )
+            corpus_output = self.model(
+                corpus_inputs,
+                return_dense=True, return_sparse=True, return_colbert_vecs=True,
+                return_sparse_embedding=True
+            )
+
+            q_dense_vecs, q_sparse_vecs, q_colbert_vecs = queries_output['dense_vecs'], queries_output['sparse_vecs'], \
+            queries_output['colbert_vecs']
+            p_dense_vecs, p_sparse_vecs, p_colbert_vecs = corpus_output['dense_vecs'], corpus_output['sparse_vecs'], \
+            corpus_output['colbert_vecs']
+
+            dense_scores = self.model.compute_dense_score(q_dense_vecs, p_dense_vecs)
+            sparse_scores = self.model.compute_sparse_score(q_sparse_vecs, p_sparse_vecs)
+            colbert_scores = self.model.compute_colbert_score(
+                q_colbert_vecs, p_colbert_vecs,
+                q_mask=queries_inputs['attention_mask']
+            )
+
+            if weights_for_different_modes is None:
+                weights_for_different_modes = [1., 1., 1.]
+                weight_sum = 3
+                logger.info("default weights for dense, sparse, colbert are [1.0, 1.0, 1.0] ")
+            else:
+                assert len(weights_for_different_modes) == 3
+                weight_sum = sum(weights_for_different_modes)
+
+            inx = torch.arange(0, len(sentences_batch))
+            dense_scores, sparse_scores, colbert_scores = dense_scores[inx, inx].float(), sparse_scores[
+                inx, inx].float(), colbert_scores[inx, inx].float()
+
+            all_scores['colbert'].extend(
+                self._convert_to_numpy(colbert_scores, device=device).tolist()
+            )
+            all_scores['sparse'].extend(
+                self._convert_to_numpy(sparse_scores, device=device).tolist()
+            )
+            all_scores['dense'].extend(
+                self._convert_to_numpy(dense_scores, device=device).tolist()
+            )
+            all_scores['sparse+dense'].extend(
+                self._convert_to_numpy(
+                    (sparse_scores * weights_for_different_modes[1] + dense_scores * weights_for_different_modes[0])
+                    / (weights_for_different_modes[1] + weights_for_different_modes[0]),
+                    device=device,
+                ).tolist()
+            )
+            all_scores['colbert+sparse+dense'].extend(
+                self._convert_to_numpy(
+                    (colbert_scores * weights_for_different_modes[2]
+                     + sparse_scores * weights_for_different_modes[1]
+                     + dense_scores * weights_for_different_modes[0]) / weight_sum,
+                    device=device,
+                ).tolist()
+            )
+
+        if one_input_pair:
+            return {k: v[0] for k, v in all_scores.items()}
+        return all_scores
+
+    def _concatenate_results_from_multi_process(
+        self,
+        results_list: List[Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Any]]
+    ):
+        """Concatenate and return the results from all the processes.
+
+        Args:
+            results_list (List[Dict[Literal[&quot;dense_vecs&quot;, &quot;lexical_weights&quot;, &quot;colbert_vecs&quot;], Any]]): 
+                A list of results from all the processes.
+
+        Returns:
+            Dict: The merged encoding results from the multi processes.
+        """
+        merged_results = {
+            "dense_vecs": [],
+            "lexical_weights": [],
+            "colbert_vecs": []
+        }
+        for key in merged_results.keys():
+            for results in results_list:
+                if results[key] is None:
+                    merged_results[key] = None
+                    break
+                else:
+                    if key == "dense_vecs":
+                        merged_results[key].append(results[key])
+                    else:
+                        merged_results[key].extend(results[key])
+
+        if merged_results["dense_vecs"] is not None:
+            merged_results["dense_vecs"] = np.concatenate(merged_results["dense_vecs"], axis=0)
+
+        return merged_results
+
+    def _concatenate_compute_score_results_from_multi_process(
+        self,
+        results_list: List[Dict[Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"], List[float]]]
+    ):
+        """Concatenate and return the results from all the processes.
+
+        Args:
+            results_list (List[Dict[Literal[&quot;colbert&quot;, &quot;sparse&quot;, &quot;dense&quot;, &quot;sparse): 
+                A list of computed scores.
+
+        Returns:
+            Dict: The merged computed scores from the multi processes.
+        """
+        merged_results = {
+            "colbert": [],
+            "sparse": [],
+            "dense": [],
+            "sparse+dense": [],
+            "colbert+sparse+dense": []
+        }
+        for key in merged_results.keys():
+            for results in results_list:
+                merged_results[key].extend(results[key])
+
+        return merged_results

FlagEmbedding/inference/embedder/model_mapping.py

@@ -0,0 +1,274 @@
+from enum import Enum
+from typing import Type, List
+from dataclasses import dataclass
+from collections import OrderedDict
+
+from FlagEmbedding.abc.inference import AbsEmbedder
+from FlagEmbedding.inference.embedder import FlagModel, BGEM3FlagModel, FlagLLMModel, FlagICLModel, FlagPseudoMoEModel
+
+
+class EmbedderModelClass(Enum):
+    ENCODER_ONLY_BASE = "encoder-only-base"
+    ENCODER_ONLY_M3 = "encoder-only-m3"
+    DECODER_ONLY_BASE = "decoder-only-base"
+    DECODER_ONLY_ICL = "decoder-only-icl"
+    DECODER_ONLY_PSEUDO_MOE = "decoder-only-pseudo_moe"
+
+
+EMBEDDER_CLASS_MAPPING = OrderedDict([
+    (EmbedderModelClass.ENCODER_ONLY_BASE, FlagModel),
+    (EmbedderModelClass.ENCODER_ONLY_M3, BGEM3FlagModel),
+    (EmbedderModelClass.DECODER_ONLY_BASE, FlagLLMModel),
+    (EmbedderModelClass.DECODER_ONLY_ICL, FlagICLModel),
+    (EmbedderModelClass.DECODER_ONLY_PSEUDO_MOE, FlagPseudoMoEModel)
+])
+
+
+class PoolingMethod(Enum):
+    LAST_TOKEN = "last_token"
+    CLS = "cls"
+    MEAN = "mean"
+
+
+@dataclass
+class EmbedderConfig:
+    model_class: Type[AbsEmbedder]
+    pooling_method: PoolingMethod
+    trust_remote_code: bool = False
+    query_instruction_format: str = "{}{}"
+
+
+# BGE models mapping
+BGE_MAPPING = OrderedDict([
+    (
+        "bge-reasoner-embed-qwen3-8b-0923",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+    (
+        "bge-code-v1",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="<instruct>{}\n<query>{}")
+    ),
+    (
+        "bge-en-icl", 
+        EmbedderConfig(FlagICLModel, PoolingMethod.LAST_TOKEN, query_instruction_format="<instruct>{}\n<query>{}")
+    ),
+    (
+        "bge-multilingual-gemma2",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="<instruct>{}\n<query>{}")
+    ),
+    (
+        "bge-m3",
+        EmbedderConfig(BGEM3FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-large-en-v1.5",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-base-en-v1.5",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-small-en-v1.5",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-large-zh-v1.5",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-base-zh-v1.5",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-small-zh-v1.5",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-large-en",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-base-en",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-small-en",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-large-zh",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-base-zh",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        "bge-small-zh",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+])
+
+# Qwen3-Embedding models mapping
+QWEN3_EMBEDDING_MAPPING = OrderedDict([
+    (
+        "Qwen3-Embedding-0.6B",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
+    ),
+    (
+        "Qwen3-Embedding-4B",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
+    ),
+    (
+        "Qwen3-Embedding-8B",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
+    ),
+])
+
+
+# E5 models mapping
+E5_MAPPING = OrderedDict([
+    (
+        "e5-mistral-7b-instruct",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+    (
+        "e5-large-v2",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        "e5-base-v2",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        "e5-small-v2",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        "multilingual-e5-large-instruct",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+    (
+        "multilingual-e5-large",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        "multilingual-e5-base",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        "multilingual-e5-small",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        "e5-large",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        "e5-base",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        "e5-small",
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+])
+
+# GTE models mapping
+GTE_MAPPING = OrderedDict([
+    (
+        "gte-Qwen2-7B-instruct",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+    (
+        "gte-Qwen2-1.5B-instruct",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+    (
+        "gte-Qwen1.5-7B-instruct",
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+    (
+        "gte-multilingual-base",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS, trust_remote_code=True)
+    ),
+    (
+        "gte-large-en-v1.5",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS, trust_remote_code=True)
+    ),
+    (
+        "gte-base-en-v1.5",
+        EmbedderConfig(FlagModel, PoolingMethod.CLS, True)
+    ),
+    (
+        'gte-large',
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        'gte-base',
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        'gte-small',
+        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
+    ),
+    (
+        'gte-large-zh',
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        'gte-base-zh',
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+    (
+        'gte-small-zh',
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+])
+
+# SFR models mapping
+SFR_MAPPING = OrderedDict([
+    (
+        'SFR-Embedding-2_R',
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+    (
+        'SFR-Embedding-Mistral',
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+])
+
+# Linq models mapping
+LINQ_MAPPING = OrderedDict([
+    (
+        'Linq-Embed-Mistral',
+        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
+    ),
+])
+
+# BCE models mapping
+BCE_MAPPING = OrderedDict([
+    (
+        'bce-embedding-base_v1',
+        EmbedderConfig(FlagModel, PoolingMethod.CLS)
+    ),
+])
+
+# Combine all mappings
+AUTO_EMBEDDER_MAPPING = OrderedDict()
+AUTO_EMBEDDER_MAPPING.update(BGE_MAPPING)
+AUTO_EMBEDDER_MAPPING.update(QWEN3_EMBEDDING_MAPPING)
+AUTO_EMBEDDER_MAPPING.update(E5_MAPPING)
+AUTO_EMBEDDER_MAPPING.update(GTE_MAPPING)
+AUTO_EMBEDDER_MAPPING.update(SFR_MAPPING)
+AUTO_EMBEDDER_MAPPING.update(LINQ_MAPPING)
+AUTO_EMBEDDER_MAPPING.update(BCE_MAPPING)
+
+# TODO: Add more models, such as Jina, Stella_v5, NV-Embed, etc.
+
+def support_native_bge_model_list()->List[str]:
+    return list(BGE_MAPPING.keys())
+
+def support_model_list()->List[str]:
+    return list(AUTO_EMBEDDER_MAPPING.keys())

FlagEmbedding/inference/reranker/__init__.py

@@ -0,0 +1,11 @@
+from .decoder_only import FlagLLMReranker, LayerWiseFlagLLMReranker, LightWeightFlagLLMReranker
+from .encoder_only import FlagReranker
+from .model_mapping import RerankerModelClass
+
+__all__ = [
+    "FlagReranker",
+    "FlagLLMReranker",
+    "LayerWiseFlagLLMReranker",
+    "LightWeightFlagLLMReranker",
+    "RerankerModelClass",
+]

FlagEmbedding/inference/reranker/decoder_only/__init__.py

@@ -0,0 +1,9 @@
+from .base import BaseLLMReranker as FlagLLMReranker
+from .layerwise import LayerWiseLLMReranker as LayerWiseFlagLLMReranker
+from .lightweight import LightweightLLMReranker as LightWeightFlagLLMReranker
+
+__all__ = [
+    "FlagLLMReranker",
+    "LayerWiseFlagLLMReranker",
+    "LightWeightFlagLLMReranker"
+]

FlagEmbedding/inference/reranker/decoder_only/base.py

@@ -0,0 +1,506 @@
+import torch
+import warnings
+import numpy as np
+from tqdm import tqdm, trange
+from typing import Any, List, Union, Tuple, Optional
+from peft import PeftModel
+from torch import Tensor
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from torch.utils.data import Dataset, DataLoader
+
+from FlagEmbedding.abc.inference import AbsReranker
+from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
+
+
+def last_logit_pool(logits: Tensor,
+                    attention_mask: Tensor) -> Tensor:
+    """Pool the last logit.
+
+    Args:
+        logits (torch.Tensor): The output logits of the model.
+        attention_mask (torch.Tensor): Attention mask.
+
+    Returns:
+        torch.Tensor: The tensor after pooling.
+    """
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return logits[:, -1, :]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = logits.shape[0]
+        return torch.stack([logits[i, sequence_lengths[i], :] for i in range(batch_size)], dim=0)
+
+
+class DatasetForReranker(Dataset):
+    """Prepare the dataset for dataloader.
+
+    Args:
+        all_queries_inputs (_type_): All the input queries.
+        all_passages_inputs (_type_): All the input passages.
+        tokenizer_path (str): Path to the tokenizer to use.
+        max_len (int, optional): Maximum length of tokens. Defaults to :data:`512`.
+        cache_dir (Optional[str], optional): Cache directory for the tokenzier. Defaults to :data:`None`.
+        prompt (Optional[str], optional): Prompt for the specific task, will use the default if not provided.
+            Defaults to `None`.
+    """
+    def __init__(
+        self,
+        all_queries_inputs,
+        all_passages_inputs,
+        tokenizer_path: str,
+        max_len: int = 512,
+        cache_dir: Optional[str] = None,
+        prompt: Optional[str] = None,
+        **kwargs: Any, 
+    ):
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_path,
+            trust_remote_code=True,
+            cache_dir=cache_dir
+        )
+
+        self.all_queries_inputs = all_queries_inputs
+        self.all_passages_inputs = all_passages_inputs
+        self.max_len = max_len
+        self.total_len = len(self.all_queries_inputs)
+        self.kwargs = kwargs
+
+        if prompt is None:
+            prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
+        self.prompt_inputs = self.tokenizer(
+            prompt,
+            return_tensors=None,
+            add_special_tokens=False
+        )['input_ids']
+        sep = "\n"
+        self.sep_inputs = self.tokenizer(
+            sep,
+            return_tensors=None,
+            add_special_tokens=False
+        )['input_ids']
+
+        self.encode_max_length = self.max_len + len(self.sep_inputs) + len(self.prompt_inputs)
+
+    def __len__(self):
+        return self.total_len
+
+    def __getitem__(self, item):
+        query_inputs = self.all_queries_inputs[item]
+        passage_inputs = self.all_passages_inputs[item]
+        if self.tokenizer.bos_token_id is not None and self.tokenizer.bos_token_id != self.tokenizer.pad_token_id:
+            item = self.tokenizer.prepare_for_model(
+                [self.tokenizer.bos_token_id] + query_inputs['input_ids'],
+                self.sep_inputs + passage_inputs['input_ids'],
+                truncation='only_second',
+                max_length=self.encode_max_length,
+                padding=False,
+                return_attention_mask=False,
+                return_token_type_ids=False,
+                add_special_tokens=False
+            )
+        else:
+            item = self.tokenizer.prepare_for_model(
+                query_inputs['input_ids'],
+                self.sep_inputs + passage_inputs['input_ids'],
+                truncation='only_second',
+                max_length=self.encode_max_length,
+                padding=False,
+                return_attention_mask=False,
+                return_token_type_ids=False,
+                add_special_tokens=False
+            )
+        item['input_ids'] = item['input_ids'] + self.sep_inputs + self.prompt_inputs
+        item['attention_mask'] = [1] * len(item['input_ids'])
+        item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+        if 'position_ids' in item.keys():
+            item['position_ids'] = list(range(len(item['input_ids'])))
+
+        return item
+
+
+class Collater:
+    """
+    Collator of the reranker.
+    
+    Args:
+        tokenizer (transformers.AutoTokenizer): The tokenizer for reranker.
+        max_len (int): Maximum length of tokens.
+    """
+    def __init__(self, tokenizer, max_len):
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+        self.pad_to_multiple_of = 8
+        self.label_pad_token_id = -100
+        warnings.filterwarnings("ignore",
+                                message="`max_length` is ignored when `padding`=`True` and there is no truncation strategy.")
+
+    def __call__(self, data):
+        labels = [feature["labels"] for feature in data] if "labels" in data[0].keys() else None
+        # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
+        # same length to return tensors.
+        if labels is not None:
+            max_label_length = max(len(l) for l in labels)
+            if self.pad_to_multiple_of is not None:
+                max_label_length = (
+                        (max_label_length + self.pad_to_multiple_of - 1)
+                        // self.pad_to_multiple_of
+                        * self.pad_to_multiple_of
+                )
+
+            padding_side = self.tokenizer.padding_side
+            for feature in data:
+                remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
+                if isinstance(feature["labels"], list):
+                    feature["labels"] = (
+                        feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"]
+                    )
+                elif padding_side == "right":
+                    feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
+                else:
+                    feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
+
+        return self.tokenizer.pad(
+            data,
+            padding=True,
+            pad_to_multiple_of=8,
+            return_tensors='pt',
+        )
+
+
+class BaseLLMReranker(AbsReranker):
+    """Base reranker class for LLM like decoder only models.
+
+    Args:
+        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 :data:`None`.
+        use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+            degradation. Defaults to :data:`False`. Defaults to :data:`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 :attr:`query_instruction_format`. Defaults to :data:`"A: "`.
+        query_instruction_format (str, optional): The template for :attr:`query_instruction_for_rerank`. Defaults to :data:`"{}{}"`.
+        passage_instruction_for_rerank (str, optional): Passage instruction for retrieval tasks, which will be used with
+            with :attr:`passage_instruction_format`. Defaults to :data:`"B: "`.
+        passage_instruction_format (str, optional): The template for passage. Defaults to "{}{}".
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        trust_remote_code (bool, optional): trust_remote_code. Defaults to :data:`False`.
+        devices (Union[str, List[str], List[int]], optional): Devices to use for model inference, such as ["cuda:0"] or ["0"].
+            Defaults to :data:`None`.
+        prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`128`.
+        query_max_length (int, optional): Maximum length for queries. If not specified, will be 3/4 of :attr:`max_length`.
+            Defaults to :data:`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 :data:`False`.
+    """
+    def __init__(
+        self,
+        model_name_or_path: str,
+        peft_path: Optional[str] = None,
+        use_fp16: bool = False,
+        use_bf16: bool = False,
+        query_instruction_for_rerank: str = "A: ",
+        query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_rerank
+        passage_instruction_for_rerank: str = "B: ",
+        passage_instruction_format: str = "{}{}", # specify the format of passage_instruction_for_rerank
+        cache_dir: Optional[str] = None,
+        trust_remote_code: bool = False,
+        devices: Union[str, List[str], List[int]] = None, # specify devices, such as ["cuda:0"] or ["0"]
+        # inference
+        prompt: Optional[str] = None,
+        batch_size: int = 128,
+        query_max_length: int = None,
+        max_length: int = 512,
+        normalize: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            model_name_or_path=model_name_or_path,
+            use_fp16=use_fp16,
+            query_instruction_for_rerank=query_instruction_for_rerank,
+            query_instruction_format=query_instruction_format,
+            passage_instruction_for_rerank=passage_instruction_for_rerank,
+            passage_instruction_format=passage_instruction_format,
+            devices=devices,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            max_length=max_length,
+            normalize=normalize,
+            prompt=prompt,
+            **kwargs
+        )
+
+        self.prompt = prompt
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            cache_dir=cache_dir,
+            trust_remote_code=trust_remote_code
+        )
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name_or_path,
+            cache_dir=cache_dir,
+            trust_remote_code=trust_remote_code,
+            torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
+        )
+        if peft_path:
+            self.model = PeftModel.from_pretrained(self.model, peft_path)
+            self.model = self.model.merge_and_unload()
+
+        self.yes_loc = self.tokenizer('Yes', add_special_tokens=False)['input_ids'][0]
+
+    @torch.no_grad()
+    def compute_score_single_gpu(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: Optional[int] = None,
+        query_max_length: Optional[int] = None,
+        max_length: Optional[int] = None,
+        prompt: Optional[str] = None,
+        normalize: Optional[bool] = None,
+        use_dataloader: bool = False,
+        num_workers: int = None,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ) -> List[float]:
+        """Compute the relevance scores using a single GPU.
+
+        Args:
+            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 :data:`None`.
+            query_max_length (Optional[int], optional): Maximum length of tokens of queries. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
+            normalize (Optional[bool], optional): If True, use Sigmoid to normalize the results. Defaults to :data:`None`.
+            use_dataloader (bool, optional): If True, will use the dataloader to load the datasets. Defaults to :data:`False`.
+            num_workers (int, optional): Number of workers for dataloader. Defaults to :data:`None`.
+            device (Optional[str], optional): Device to use for computation. Defaults to :data:`None`.
+
+        Returns:
+            List[float]: The computed scores.
+        """
+        if prompt is None: prompt = self.prompt
+        if batch_size is None: batch_size = self.batch_size
+        if max_length is None: max_length = self.max_length
+        if query_max_length is None:
+            if self.query_max_length is not None:
+                query_max_length = self.query_max_length
+            else:
+                query_max_length = max_length * 3 // 4
+        if normalize is None: normalize = self.normalize
+
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu": self.use_fp16 = False
+        if self.use_fp16: self.model.half()
+
+        self.model.to(device)
+        self.model.eval()
+
+        assert isinstance(sentence_pairs, list)
+        if isinstance(sentence_pairs[0], str):
+            sentence_pairs = [sentence_pairs]
+        
+        # tokenize without padding to get the correct length
+        all_queries_inputs = []
+        all_passages_inputs = []
+        for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
+                                  disable=len(sentence_pairs) < batch_size):
+            sentences_batch = sentence_pairs[start_index:start_index + batch_size]
+            queries = [s[0] for s in sentences_batch]
+            passages = [s[1] for s in sentences_batch]
+            queries_inputs_batch = self.tokenizer(
+                queries,
+                return_tensors=None,
+                add_special_tokens=False,
+                max_length=query_max_length,
+                truncation=True,
+                **kwargs
+            )
+            passages_inputs_batch = self.tokenizer(
+                passages,
+                return_tensors=None,
+                add_special_tokens=False,
+                max_length=max_length,
+                truncation=True,
+                **kwargs
+            )
+            queries_inputs_batch = [{
+                k: queries_inputs_batch[k][i] for k in queries_inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            passages_inputs_batch = [{
+                k: passages_inputs_batch[k][i] for k in passages_inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+
+            all_queries_inputs.extend(queries_inputs_batch)
+            all_passages_inputs.extend(passages_inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) - len(y['input_ids']) for (x, y) in zip(all_queries_inputs, all_passages_inputs)])
+        all_queries_inputs_sorted = [all_queries_inputs[i] for i in length_sorted_idx]
+        all_passages_inputs_sorted = [all_passages_inputs[i] for i in length_sorted_idx]
+
+        # other inputs
+        if prompt is None:
+            prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
+        prompt_inputs = self.tokenizer(
+            prompt,
+            return_tensors=None,
+            add_special_tokens=False
+        )['input_ids']
+        sep = "\n"
+        sep_inputs = self.tokenizer(
+            sep,
+            return_tensors=None,
+            add_special_tokens=False
+        )['input_ids']
+        encode_max_length = max_length + len(sep_inputs) + len(prompt_inputs)
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                batch_inputs = []
+                for query_inputs, passage_inputs in zip(
+                    all_queries_inputs_sorted[:min(len(all_queries_inputs_sorted), batch_size)], 
+                    all_passages_inputs_sorted[:min(len(all_passages_inputs_sorted), batch_size)]
+                ):
+                    if self.tokenizer.bos_token_id is not None and self.tokenizer.bos_token_id != self.tokenizer.pad_token_id:
+                        item = self.tokenizer.prepare_for_model(
+                            [self.tokenizer.bos_token_id] + query_inputs['input_ids'],
+                            sep_inputs + passage_inputs['input_ids'],
+                            truncation='only_second',
+                            max_length=encode_max_length,
+                            padding=False,
+                            return_attention_mask=False,
+                            return_token_type_ids=False,
+                            add_special_tokens=False
+                        )
+                    else:
+                        item = self.tokenizer.prepare_for_model(
+                            query_inputs['input_ids'],
+                            sep_inputs + passage_inputs['input_ids'],
+                            truncation='only_second',
+                            max_length=encode_max_length,
+                            padding=False,
+                            return_attention_mask=False,
+                            return_token_type_ids=False,
+                            add_special_tokens=False
+                        )
+                    item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
+                    item['attention_mask'] = [1] * len(item['input_ids'])
+                    item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+                    if 'position_ids' in item.keys():
+                        item['position_ids'] = list(range(len(item['input_ids'])))
+                    batch_inputs.append(item)
+
+                collater_instance = Collater(self.tokenizer, encode_max_length)
+                batch_inputs = collater_instance([{
+                        'input_ids': item['input_ids'],
+                        'attention_mask': item['attention_mask']
+                    } for item in batch_inputs]
+                )
+
+                batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
+
+                self.model(**batch_inputs, output_hidden_states=True)
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        dataset, dataloader = None, None
+        if use_dataloader:
+            if num_workers is None:
+                num_workers = min(batch_size, 16)
+            dataset = DatasetForReranker(
+                all_queries_inputs_sorted,
+                all_passages_inputs_sorted,
+                self.model_name_or_path,
+                max_length,
+                cache_dir=self.cache_dir,
+                prompt=prompt,
+                **kwargs
+            )
+            dataloader = DataLoader(
+                dataset, shuffle=False, batch_size=batch_size, drop_last=False,
+                num_workers=num_workers,
+                collate_fn=Collater(self.tokenizer, encode_max_length)
+            )
+
+        all_scores = []
+        if dataloader is not None:
+            for inputs in tqdm(dataloader):
+                inputs = inputs.to(device)
+
+                outputs = self.model(**inputs, output_hidden_states=True)
+                logits = outputs.logits
+                scores = last_logit_pool(logits, inputs['attention_mask'])
+                scores = scores[:, self.yes_loc]
+                all_scores.extend(scores.cpu().float().tolist())
+        else:
+            for batch_start in trange(0, len(all_queries_inputs_sorted), batch_size):
+                queries_inputs = all_queries_inputs_sorted[batch_start:batch_start+batch_size]
+                passages_inputs = all_passages_inputs_sorted[batch_start:batch_start+batch_size]
+
+                batch_inputs = []
+                for query_inputs, passage_inputs in zip(queries_inputs, passages_inputs):
+                    if self.tokenizer.bos_token_id is not None and self.tokenizer.bos_token_id != self.tokenizer.pad_token_id:
+                        item = self.tokenizer.prepare_for_model(
+                            [self.tokenizer.bos_token_id] + query_inputs['input_ids'],
+                            sep_inputs + passage_inputs['input_ids'],
+                            truncation='only_second',
+                            max_length=encode_max_length,
+                            padding=False,
+                            return_attention_mask=False,
+                            return_token_type_ids=False,
+                            add_special_tokens=False
+                        )
+                    else:
+                        item = self.tokenizer.prepare_for_model(
+                            query_inputs['input_ids'],
+                            sep_inputs + passage_inputs['input_ids'],
+                            truncation='only_second',
+                            max_length=encode_max_length,
+                            padding=False,
+                            return_attention_mask=False,
+                            return_token_type_ids=False,
+                            add_special_tokens=False
+                        )
+                    item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
+                    item['attention_mask'] = [1] * len(item['input_ids'])
+                    item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+                    if 'position_ids' in item.keys():
+                        item['position_ids'] = list(range(len(item['input_ids'])))
+                    batch_inputs.append(item)
+
+                collater_instance = Collater(self.tokenizer, encode_max_length)
+                batch_inputs = collater_instance([{
+                        'input_ids': item['input_ids'],
+                        'attention_mask': item['attention_mask']
+                    } for item in batch_inputs]
+                )
+
+                batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
+
+                outputs = self.model(**batch_inputs, output_hidden_states=True)
+                logits = outputs.logits
+                scores = last_logit_pool(logits, batch_inputs['attention_mask'])
+                scores = scores[:, self.yes_loc]
+                all_scores.extend(scores.cpu().float().tolist())
+
+        all_scores = [all_scores[idx] for idx in np.argsort(length_sorted_idx)]
+
+        if normalize:
+            all_scores = [sigmoid(score) for score in all_scores]
+
+        # if len(all_scores) == 1:
+        #     return all_scores[0]
+
+        return all_scores

FlagEmbedding/inference/reranker/decoder_only/layerwise.py

@@ -0,0 +1,380 @@
+import torch
+import warnings
+import numpy as np
+from tqdm import tqdm, trange
+from typing import Any, List, Union, Tuple, Optional
+from peft import PeftModel
+from torch import Tensor
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from torch.utils.data import DataLoader
+
+from FlagEmbedding.abc.inference import AbsReranker
+from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
+from FlagEmbedding.inference.reranker.decoder_only.base import DatasetForReranker, Collater
+
+from .models.modeling_minicpm_reranker import LayerWiseMiniCPMForCausalLM
+
+
+def last_logit_pool_layerwise(logits: Tensor,
+                              attention_mask: Tensor) -> Tensor:
+    """Pool the last logit.
+
+    Args:
+        logits (torch.Tensor): The output logits of the model.
+        attention_mask (torch.Tensor): Attention mask.
+
+    Returns:
+        torch.Tensor: The tensor after pooling.
+    """
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return logits[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = logits.shape[0]
+        return logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+
+class LayerWiseLLMReranker(AbsReranker):
+    """Base reranker class for layerwise LLM like decoder only models.
+
+    Args:
+        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 :data:`None`.
+        use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+            degradation. Defaults to :data:`False`. Defaults to :data:`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 :attr:`query_instruction_format`. Defaults to :data:`"A: "`.
+        query_instruction_format (str, optional): The template for :attr:`query_instruction_for_rerank`. Defaults to :data:`"{}{}"`.
+        passage_instruction_for_rerank (str, optional): Passage instruction for retrieval tasks, which will be used with
+            with :attr:`passage_instruction_format`. Defaults to :data:`"B: "`.
+        passage_instruction_format (str, optional): The template for passage. Defaults to "{}{}".
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        trust_remote_code (bool, optional): trust_remote_code. Defaults to :data:`False`.
+        devices (Union[str, List[str], List[int]], optional): Devices to use for model inference, such as ["cuda:0"] or ["0"].
+            Defaults to :data:`None`.
+        cutoff_layers (Optional[List[int]]): Pick which layers are used for computing the score. Defaults to :data:`None`.
+        prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`128`.
+        query_max_length (int, optional): Maximum length for queries. If not specified, will be 3/4 of :attr:`max_length`.
+            Defaults to :data:`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 :data:`False`.
+    """
+    def __init__(
+        self,
+        model_name_or_path: str,
+        peft_path: Optional[str] = None,
+        use_fp16: bool = False,
+        use_bf16: bool = False,
+        query_instruction_for_rerank: str = "A: ",
+        query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_rerank
+        passage_instruction_for_rerank: str = "B: ",
+        passage_instruction_format: str = "{}{}", # specify the format of passage_instruction_for_rerank
+        cache_dir: Optional[str] = None,
+        trust_remote_code: bool = False,
+        devices: Optional[Union[str, List[str], List[int]]] = None, # specify devices, such as ["cuda:0"] or ["0"]
+        # inference
+        cutoff_layers: Optional[List[int]] = None, 
+        prompt: Optional[str] = None,
+        batch_size: int = 128,
+        query_max_length: Optional[int] = None,
+        max_length: int = 512,
+        normalize: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            model_name_or_path=model_name_or_path,
+            use_fp16=use_fp16,
+            query_instruction_for_rerank=query_instruction_for_rerank,
+            query_instruction_format=query_instruction_format,
+            passage_instruction_for_rerank=passage_instruction_for_rerank,
+            passage_instruction_format=passage_instruction_format,
+            devices=devices,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            max_length=max_length,
+            normalize=normalize,
+            **kwargs
+        )
+
+        self.cutoff_layers = cutoff_layers
+        self.prompt = prompt
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            cache_dir=cache_dir,
+            trust_remote_code=trust_remote_code
+        )
+
+        if use_bf16 is False and use_fp16 is False:
+            warnings.warn("Due to model constraints, `use_bf16` and `use_fp16` cannot both be `False`. Here, `use_fp16` is set to `True` by default.", UserWarning)
+            self.use_fp16 = True
+        
+        try:
+            self.model = LayerWiseMiniCPMForCausalLM.from_pretrained(
+                model_name_or_path,
+                cache_dir=cache_dir,
+                trust_remote_code=trust_remote_code,
+                torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
+            )
+        except:
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_name_or_path,
+                cache_dir=cache_dir,
+                trust_remote_code=trust_remote_code,
+                torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
+            )
+        if peft_path:
+            self.model = PeftModel.from_pretrained(self.model,peft_path)
+            self.model = self.model.merge_and_unload()
+
+    @torch.no_grad()
+    def compute_score_single_gpu(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: Optional[int] = None,
+        query_max_length: Optional[int] = None,
+        max_length: Optional[int] = None,
+        cutoff_layers: Optional[List[int]] = None, 
+        prompt: Optional[str] = None,
+        normalize: Optional[bool] = None,
+        use_dataloader: bool = False,
+        num_workers: Optional[int] = None,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ) -> List[float]:
+        """Compute the relevance scores using a single GPU.
+
+        Args:
+            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 :data:`None`.
+            query_max_length (Optional[int], optional): Maximum length of tokens of queries. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            cutoff_layers (Optional[List[int]], optional): Pick which layers are used for computing the score. Defaults to :data:`None`.
+            prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
+            normalize (Optional[bool], optional): If True, use Sigmoid to normalize the results. Defaults to :data:`None`.
+            use_dataloader (bool, optional): If True, will use the dataloader to load the datasets. Defaults to :data:`False`.
+            num_workers (int, optional): Number of workers for dataloader. Defaults to :data:`None`.
+            device (Optional[str], optional): Device to use for computation. Defaults to :data:`None`.
+
+        Returns:
+            List[float]: The computed scores.
+        """
+        if cutoff_layers is None: cutoff_layers = self.cutoff_layers
+        if prompt is None: prompt = self.prompt
+        if batch_size is None: batch_size = self.batch_size
+        if max_length is None: max_length = self.max_length
+        if query_max_length is None:
+            if self.query_max_length is not None:
+                query_max_length = self.query_max_length
+            else:
+                query_max_length = max_length * 3 // 4
+        if normalize is None: normalize = self.normalize
+
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu": self.use_fp16 = False
+        if self.use_fp16: self.model.half()
+
+        self.model.to(device)
+        self.model.eval()
+
+        assert isinstance(sentence_pairs, list)
+        if isinstance(sentence_pairs[0], str):
+            sentence_pairs = [sentence_pairs]
+
+        # tokenize without padding to get the correct length
+        all_queries_inputs = []
+        all_passages_inputs = []
+        for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
+                                  disable=len(sentence_pairs) < batch_size):
+            sentences_batch = sentence_pairs[start_index:start_index + batch_size]
+            queries = [s[0] for s in sentences_batch]
+            passages = [s[1] for s in sentences_batch]
+            queries_inputs_batch = self.tokenizer(
+                queries,
+                return_tensors=None,
+                add_special_tokens=False,
+                max_length=query_max_length,
+                truncation=True,
+                **kwargs
+            )
+            passages_inputs_batch = self.tokenizer(
+                passages,
+                return_tensors=None,
+                add_special_tokens=False,
+                max_length=max_length,
+                truncation=True,
+                **kwargs
+            )
+            queries_inputs_batch = [{
+                k: queries_inputs_batch[k][i] for k in queries_inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            passages_inputs_batch = [{
+                k: passages_inputs_batch[k][i] for k in passages_inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+
+            all_queries_inputs.extend(queries_inputs_batch)
+            all_passages_inputs.extend(passages_inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) - len(y['input_ids']) for (x, y) in zip(all_queries_inputs, all_passages_inputs)])
+        all_queries_inputs_sorted = [all_queries_inputs[i] for i in length_sorted_idx]
+        all_passages_inputs_sorted = [all_passages_inputs[i] for i in length_sorted_idx]
+
+        # other inputs
+        if prompt is None:
+            prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
+        prompt_inputs = self.tokenizer(
+            prompt,
+            return_tensors=None,
+            add_special_tokens=False
+        )['input_ids']
+        sep = "\n"
+        sep_inputs = self.tokenizer(
+            sep,
+            return_tensors=None,
+            add_special_tokens=False
+        )['input_ids']
+        encode_max_length = max_length + len(sep_inputs) + len(prompt_inputs)
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                batch_inputs = []
+                for query_inputs, passage_inputs in zip(
+                    all_queries_inputs_sorted[:min(len(all_queries_inputs_sorted), batch_size)], 
+                    all_passages_inputs_sorted[:min(len(all_passages_inputs_sorted), batch_size)]
+                ):
+                    item = self.tokenizer.prepare_for_model(
+                        [self.tokenizer.bos_token_id] + query_inputs['input_ids'],
+                        sep_inputs + passage_inputs['input_ids'],
+                        truncation='only_second',
+                        max_length=encode_max_length,
+                        padding=False,
+                        return_attention_mask=False,
+                        return_token_type_ids=False,
+                        add_special_tokens=False
+                    )
+                    item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
+                    item['attention_mask'] = [1] * len(item['input_ids'])
+                    item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+                    if 'position_ids' in item.keys():
+                        item['position_ids'] = list(range(len(item['input_ids'])))
+                    batch_inputs.append(item)
+
+                collater_instance = Collater(self.tokenizer, encode_max_length)
+                batch_inputs = collater_instance([{
+                        'input_ids': item['input_ids'],
+                        'attention_mask': item['attention_mask']
+                    } for item in batch_inputs]
+                )
+
+                batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
+
+                self.model(**batch_inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        dataset, dataloader = None, None
+        if use_dataloader:
+            if num_workers is None:
+                num_workers = min(batch_size, 16)
+            dataset = DatasetForReranker(
+                all_queries_inputs_sorted,
+                all_passages_inputs_sorted,
+                self.model_name_or_path,
+                max_length,
+                cache_dir=self.cache_dir,
+                prompt=prompt,
+                **kwargs
+            )
+            dataloader = DataLoader(
+                dataset, shuffle=False, batch_size=batch_size, drop_last=False,
+                num_workers=num_workers,
+                collate_fn=Collater(self.tokenizer, encode_max_length)
+            )
+
+        all_scores = []
+        if dataloader is not None:
+            for inputs in tqdm(dataloader):
+                inputs = inputs.to(device)
+
+                outputs = self.model(**inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
+                all_logits = outputs.logits
+                tmp_all_scores = []
+                for logits in all_logits:
+                    scores = last_logit_pool_layerwise(logits, inputs['attention_mask'])
+                    tmp_all_scores.append(scores.contiguous())
+
+                if len(all_scores) == 0:
+                    for _ in range(len(tmp_all_scores)):
+                        all_scores.append([])
+
+                for i in range(len(tmp_all_scores)):
+                    all_scores[i].extend(tmp_all_scores[i].cpu().float().tolist())
+        else:
+            for batch_start in trange(0, len(all_queries_inputs_sorted), batch_size):
+                queries_inputs = all_queries_inputs_sorted[batch_start:batch_start+batch_size]
+                passages_inputs = all_passages_inputs_sorted[batch_start:batch_start+batch_size]
+
+                batch_inputs = []
+                for query_inputs, passage_inputs in zip(queries_inputs, passages_inputs):
+                    item = self.tokenizer.prepare_for_model(
+                        [self.tokenizer.bos_token_id] + query_inputs['input_ids'],
+                        sep_inputs + passage_inputs['input_ids'],
+                        truncation='only_second',
+                        max_length=encode_max_length,
+                        padding=False,
+                        return_attention_mask=False,
+                        return_token_type_ids=False,
+                        add_special_tokens=False
+                    )
+                    item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
+                    item['attention_mask'] = [1] * len(item['input_ids'])
+                    item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+                    if 'position_ids' in item.keys():
+                        item['position_ids'] = list(range(len(item['input_ids'])))
+                    batch_inputs.append(item)
+
+                collater_instance = Collater(self.tokenizer, encode_max_length)
+                batch_inputs = collater_instance([{
+                    'input_ids': item['input_ids'],
+                    'attention_mask': item['attention_mask']
+                    } for item in batch_inputs]
+                )
+
+                batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
+
+                outputs = self.model(**batch_inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
+                all_logits = outputs.logits
+                tmp_all_scores = []
+                for logits in all_logits:
+                    scores = last_logit_pool_layerwise(logits, batch_inputs['attention_mask'])
+                    tmp_all_scores.append(scores.contiguous())
+
+                if len(all_scores) == 0:
+                    for _ in range(len(tmp_all_scores)):
+                        all_scores.append([])
+
+                for i in range(len(tmp_all_scores)):
+                    all_scores[i].extend(tmp_all_scores[i].cpu().float().tolist())
+
+        for i in range(len(all_scores)):
+            all_scores[i] = [all_scores[i][idx] for idx in np.argsort(length_sorted_idx)]
+            if normalize:
+                all_scores[i] = [sigmoid(score) for score in all_scores[i]]
+        
+        if len(all_scores) == 1 and isinstance(all_scores[0], list):
+            all_scores = all_scores[0]
+            
+        return all_scores

FlagEmbedding/inference/reranker/decoder_only/lightweight.py

@@ -0,0 +1,449 @@
+import torch
+import sys
+import warnings
+import numpy as np
+from tqdm import trange
+from typing import Any, List, Union, Tuple, Optional
+from peft import PeftModel
+from torch import Tensor
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from FlagEmbedding.abc.inference import AbsReranker
+from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
+
+
+def last_logit_pool_lightweight(logits: Tensor,
+                    attention_mask: Tensor) -> Tensor:
+    """Pool the last logit.
+
+    Args:
+        logits (torch.Tensor): The output logits of the model.
+        attention_mask (torch.Tensor): Attention mask.
+
+    Returns:
+        torch.Tensor: The tensor after pooling.
+    """
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return logits[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = logits.shape[0]
+        return torch.stack([logits[i, sequence_lengths[i]] for i in range(batch_size)], dim=0)
+
+
+class Collater_for_lightweight:
+    """
+    Collator of the lightweight LLM reranker.
+    
+    Args:
+        tokenizer (transformers.AutoTokenizer): The tokenizer for reranker.
+        max_len (int): Maximum length of tokens.
+    """
+    def __init__(self, tokenizer, max_len):
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+        self.pad_to_multiple_of = 8
+        self.label_pad_token_id = -100
+        warnings.filterwarnings("ignore",
+                                message="`max_length` is ignored when `padding`=`True` and there is no truncation strategy.")
+
+    def __call__(self, data):
+        features = data[0]
+        query_lengths = data[1]
+        prompt_lengths = data[2]
+
+        labels = [feature["labels"] for feature in features] if "labels" in features[0].keys() else None
+        # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
+        # same length to return tensors.
+        if labels is not None:
+            max_label_length = max(len(l) for l in labels)
+            if self.pad_to_multiple_of is not None:
+                max_label_length = (
+                        (max_label_length + self.pad_to_multiple_of - 1)
+                        // self.pad_to_multiple_of
+                        * self.pad_to_multiple_of
+                )
+
+            padding_side = self.tokenizer.padding_side
+            for feature in features:
+                remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
+                if isinstance(feature["labels"], list):
+                    feature["labels"] = (
+                        feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"]
+                    )
+                elif padding_side == "right":
+                    feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
+                else:
+                    feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
+
+        collected = self.tokenizer.pad(
+            features,
+            padding=True,
+            pad_to_multiple_of=8,
+            return_tensors='pt',
+        )
+
+        return collected, query_lengths, prompt_lengths
+
+
+class LightweightLLMReranker(AbsReranker):
+    """Base reranker class for light weight LLM like decoder only models.
+
+    Args:
+        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 :data:`None`.
+        use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance 
+            degradation. Defaults to :data:`False`. Defaults to :data:`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 :attr:`query_instruction_format`. Defaults to :data:`"A: "`.
+        query_instruction_format (str, optional): The template for :attr:`query_instruction_for_rerank`. Defaults to :data:`"{}{}"`.
+        passage_instruction_for_rerank (str, optional): Passage instruction for retrieval tasks, which will be used with
+            with :attr:`passage_instruction_format`. Defaults to :data:`"B: "`.
+        passage_instruction_format (str, optional): The template for passage. Defaults to "{}{}".
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        trust_remote_code (bool, optional): trust_remote_code. Defaults to :data:`False`.
+        devices (Union[str, List[str], List[int]], optional): Devices to use for model inference, such as ["cuda:0"] or ["0"].
+            Defaults to :data:`None`.
+        cutoff_layers (Optional[List[int]]): Pick which layers are used for computing the score. Defaults to :data:`None`.
+        compress_layers (List[int], optional): Choose the layers to compress. Defaults to :data:`[8]`.
+        compress_ratio (int, optional): Ratio to compress the selected layers, supported ratios: :data:`[1, 2, 4, 8]`. 
+            Defaults to :data:`1`.
+        prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`128`.
+        query_max_length (int, optional): Maximum length for queries. If not specified, will be 3/4 of :attr:`max_length`.
+            Defaults to :data:`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 :data:`False`.
+    """
+    def __init__(
+        self,
+        model_name_or_path: str,
+        peft_path: Optional[str] = None,
+        use_fp16: bool = False,
+        use_bf16: bool = False,
+        query_instruction_for_rerank: str = "A: ",
+        query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_rerank
+        passage_instruction_for_rerank: str = "B: ",
+        passage_instruction_format: str = "{}{}", # specify the format of passage_instruction_for_rerank
+        cache_dir: Optional[str] = None,
+        trust_remote_code: bool = False,
+        devices: Union[str, List[str], List[int]] = None, # specify devices, such as ["cuda:0"] or ["0"]
+        # inference
+        cutoff_layers: Optional[List[int]] = None,
+        compress_layers: List[int] = [8],
+        compress_ratio: int = 1,
+        prompt: Optional[str] = None,
+        batch_size: int = 128,
+        query_max_length: Optional[int] = None,
+        max_length: int = 512,
+        normalize: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            from .models.gemma_model import CostWiseGemmaForCausalLM
+        except:
+            print('*') * 20
+            print('*') * 20
+            print('error for load lightweight reranker, please install transformers==4.46.0')
+            print('*') * 20
+            print('*') * 20
+            sys.exit()
+
+        super().__init__(
+            model_name_or_path=model_name_or_path,
+            use_fp16=use_fp16,
+            query_instruction_for_rerank=query_instruction_for_rerank,
+            query_instruction_format=query_instruction_format,
+            passage_instruction_for_rerank=passage_instruction_for_rerank,
+            passage_instruction_format=passage_instruction_format,
+            devices=devices,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            max_length=max_length,
+            normalize=normalize,
+            **kwargs
+        )
+
+        self.cutoff_layers = cutoff_layers
+        self.compress_layers = compress_layers
+        self.compress_ratio = compress_ratio
+        self.prompt = prompt
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            cache_dir=cache_dir,
+            trust_remote_code=trust_remote_code
+        )
+        self.tokenizer.padding_side = 'right'
+
+        if use_bf16 is False and use_fp16 is False:
+            warnings.warn("Due to model constraints, `use_bf16` and `use_fp16` cannot both be `False`. Here, `use_fp16` is set to `True` by default.", UserWarning)
+            use_fp16 = True
+
+        try:
+            self.model = CostWiseGemmaForCausalLM.from_pretrained(
+                model_name_or_path,
+                cache_dir=cache_dir,
+                trust_remote_code=trust_remote_code,
+                torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
+            )
+        except:
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_name_or_path,
+                cache_dir=cache_dir,
+                trust_remote_code=trust_remote_code,
+                torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
+            )
+        if peft_path:
+            self.model = PeftModel.from_pretrained(self.model,peft_path)
+            self.model = self.model.merge_and_unload()
+
+    @torch.no_grad()
+    def compute_score_single_gpu(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: Optional[int] = None,
+        query_max_length: Optional[int] = None,
+        max_length: Optional[int] = None,
+        cutoff_layers: Optional[List[int]] = None,
+        compress_layer: Optional[List[int]] = None,
+        compress_layers: Optional[List[int]] = None,
+        compress_ratio: Optional[int] = None,
+        prompt: Optional[str] = None,
+        normalize: Optional[bool] = None,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ) -> List[float]:
+        """Compute the relevance scores using a single GPU.
+
+        Args:
+            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 :data:`None`.
+            query_max_length (Optional[int], optional): Maximum length of tokens of queries. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            cutoff_layers (Optional[List[int]], optional): Pick which layers are used for computing the score. Defaults to :data:`None`.
+            compress_layer (Optional[List[int]]): Deprecated, use :attr:`compress_layers` instead. Defaults to :data:`None`.
+            compress_layers (Optional[List[int]]): Selected layers to compress. Defaults to :data:`None`.
+            compress_ratio (Optional[int]): Ratio to compress the selected layers, supported ratios: :data:`[1, 2, 4, 8]`. 
+                Defaults to :data:`None`.
+            prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
+            normalize (Optional[bool], optional): If True, use Sigmoid to normalize the results. Defaults to :data:`None`.
+            device (Optional[str], optional): Device to use for computation. Defaults to :data:`None`.
+
+        Returns:
+            List[float]: The computed scores.
+        """
+
+        if cutoff_layers is None: cutoff_layers = self.cutoff_layers
+        if compress_layers is None: compress_layers = self.compress_layers
+        if compress_layer is not None:
+            print('Try not to use the parameter `compress_layer`; use `compress_layers` instead.')
+            compress_layers = compress_layer
+        if compress_ratio is None: compress_ratio = self.compress_ratio
+        if prompt is None: prompt = self.prompt
+        if batch_size is None: batch_size = self.batch_size
+        if max_length is None: max_length = self.max_length
+        if query_max_length is None:
+            if self.query_max_length is not None:
+                query_max_length = self.query_max_length
+            else:
+                query_max_length = max_length * 3 // 4
+        if normalize is None: normalize = self.normalize
+
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu": self.use_fp16 = False
+        if self.use_fp16: self.model.half()
+
+        self.model.to(device)
+        self.model.eval()
+
+        assert isinstance(sentence_pairs, list)
+        if isinstance(sentence_pairs[0], str):
+            sentence_pairs = [sentence_pairs]
+
+        # tokenize without padding to get the correct length
+        all_queries_inputs = []
+        all_passages_inputs = []
+        for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
+                                  disable=len(sentence_pairs) < batch_size):
+            sentences_batch = sentence_pairs[start_index:start_index + batch_size]
+            queries = [s[0] for s in sentences_batch]
+            passages = [s[1] for s in sentences_batch]
+            queries_inputs_batch = self.tokenizer(
+                queries,
+                return_tensors=None,
+                add_special_tokens=False,
+                max_length=query_max_length,
+                truncation=True,
+                **kwargs
+            )
+            passages_inputs_batch = self.tokenizer(
+                passages,
+                return_tensors=None,
+                add_special_tokens=False,
+                max_length=max_length,
+                truncation=True,
+                **kwargs
+            )
+            queries_inputs_batch = [{
+                k: queries_inputs_batch[k][i] for k in queries_inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+            passages_inputs_batch = [{
+                k: passages_inputs_batch[k][i] for k in passages_inputs_batch.keys()
+            } for i in range(len(sentences_batch))]
+
+            all_queries_inputs.extend(queries_inputs_batch)
+            all_passages_inputs.extend(passages_inputs_batch)
+
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) - len(y['input_ids']) for (x, y) in zip(all_queries_inputs, all_passages_inputs)])
+        all_queries_inputs_sorted = [all_queries_inputs[i] for i in length_sorted_idx]
+        all_passages_inputs_sorted = [all_passages_inputs[i] for i in length_sorted_idx]
+
+        # other inputs
+        if prompt is None:
+            prompt = "Predict whether passage B contains an answer to query A."
+        prompt_inputs = self.tokenizer(
+            prompt,
+            return_tensors=None,
+            add_special_tokens=False
+        )['input_ids']
+        sep = "\n"
+        sep_inputs = self.tokenizer(
+            sep,
+            return_tensors=None,
+            add_special_tokens=False
+        )['input_ids']
+        encode_max_length = max_length + len(sep_inputs) + len(prompt_inputs)
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                batch_inputs = []
+                query_lengths = []
+                prompt_lengths = []
+                for query_inputs, passage_inputs in zip(
+                    all_queries_inputs_sorted[:min(len(all_queries_inputs_sorted), batch_size)], 
+                    all_passages_inputs_sorted[:min(len(all_passages_inputs_sorted), batch_size)]
+                ):
+                    item = self.tokenizer.prepare_for_model(
+                        [self.tokenizer.bos_token_id] + query_inputs['input_ids'],
+                        sep_inputs + passage_inputs['input_ids'],
+                        truncation='only_second',
+                        max_length=encode_max_length,
+                        padding=False,
+                        return_attention_mask=False,
+                        return_token_type_ids=False,
+                        add_special_tokens=False
+                    )
+                    item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
+                    item['attention_mask'] = [1] * len(item['input_ids'])
+                    item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+                    if 'position_ids' in item.keys():
+                        item['position_ids'] = list(range(len(item['input_ids'])))
+                    batch_inputs.append(item)
+                    query_lengths.append(len([self.tokenizer.bos_token_id] + query_inputs['input_ids'] + sep_inputs))
+                    prompt_lengths.append(len(sep_inputs + prompt_inputs))
+
+                collater_instance = Collater_for_lightweight(self.tokenizer, max_length)
+                batch_inputs = collater_instance([
+                    [{
+                        'input_ids': item['input_ids'],
+                        'attention_mask': item['attention_mask']
+                    } for item in batch_inputs],
+                    query_lengths,
+                    prompt_lengths
+                ])[0]
+
+                batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
+
+                self.model(
+                    **batch_inputs,
+                    output_hidden_states=True,
+                    compress_layer=compress_layers,
+                    compress_ratio=compress_ratio,
+                    query_lengths=query_lengths,
+                    prompt_lengths=prompt_lengths,
+                    cutoff_layers=cutoff_layers
+                )
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        all_scores = []
+        for batch_start in trange(0, len(all_queries_inputs_sorted), batch_size):
+            queries_inputs = all_queries_inputs_sorted[batch_start:batch_start+batch_size]
+            passages_inputs = all_passages_inputs_sorted[batch_start:batch_start+batch_size]
+
+            batch_inputs = []
+            query_lengths = []
+            prompt_lengths = []
+            for query_inputs, passage_inputs in zip(queries_inputs, passages_inputs):
+                item = self.tokenizer.prepare_for_model(
+                    [self.tokenizer.bos_token_id] + query_inputs['input_ids'],
+                    sep_inputs + passage_inputs['input_ids'],
+                    truncation='only_second',
+                    max_length=encode_max_length,
+                    padding=False,
+                    return_attention_mask=False,
+                    return_token_type_ids=False,
+                    add_special_tokens=False
+                )
+                item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
+                item['attention_mask'] = [1] * len(item['input_ids'])
+                item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
+                if 'position_ids' in item.keys():
+                    item['position_ids'] = list(range(len(item['input_ids'])))
+                batch_inputs.append(item)
+                query_lengths.append(len([self.tokenizer.bos_token_id] + query_inputs['input_ids'] + sep_inputs))
+                prompt_lengths.append(len(sep_inputs + prompt_inputs))
+
+            collater_instance = Collater_for_lightweight(self.tokenizer, max_length)
+            batch_inputs = collater_instance([
+                [{
+                    'input_ids': item['input_ids'],
+                    'attention_mask': item['attention_mask']
+                } for item in batch_inputs],
+                query_lengths,
+                prompt_lengths
+            ])[0]
+
+            batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
+
+            outputs = self.model(
+                **batch_inputs,
+                output_hidden_states=True,
+                compress_layer=compress_layers,
+                compress_ratio=compress_ratio,
+                query_lengths=query_lengths,
+                prompt_lengths=prompt_lengths,
+                cutoff_layers=cutoff_layers
+            )
+            scores = []
+            for i in range(len(outputs.logits)):
+                logits = last_logit_pool_lightweight(outputs.logits[i], outputs.attention_masks[i])
+                scores.append(logits.cpu().float().tolist())
+            if len(all_scores) == 0:
+                for i in range(len(scores)):
+                    all_scores.append([])
+            for i in range(len(scores)):
+                all_scores[i].extend(scores[i])
+
+        for i in range(len(all_scores)):
+            all_scores[i] = [all_scores[i][idx] for idx in np.argsort(length_sorted_idx)]
+            if normalize:
+                all_scores[i] = [sigmoid(score) for score in all_scores[i]]
+    
+        if len(all_scores) == 1 and isinstance(all_scores[0], list):
+            all_scores = all_scores[0]
+
+        return all_scores

FlagEmbedding/inference/reranker/decoder_only/models/configuration_minicpm_reranker.py

@@ -0,0 +1,208 @@
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" MiniCPM model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+MINICPM_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+class LayerWiseMiniCPMConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MiniCPMModel`]. It is used to instantiate an MiniCPM
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the MiniCPM-7B.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32000):
+            Vocabulary size of the MiniCPM model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`MiniCPMModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 11008):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with. MiniCPM 1 supports up to 2048 tokens,
+            MiniCPM 2 up to 4096, CodeMiniCPM up to 16384.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*):
+            Padding token id.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            End of stream token id.
+        pretraining_tp (`int`, *optional*, defaults to 1):
+            Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
+            document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
+            necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
+            issue](https://github.com/pytorch/pytorch/issues/76232).
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+            `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
+            these scaling strategies behave:
+            https://www.reddit.com/r/LocalMiniCPM/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
+            experimental feature, subject to breaking API changes in future versions.
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import MiniCPMModel, MiniCPMConfig
+
+    >>> # Initializing a MiniCPM minicpm-7b style configuration
+    >>> configuration = MiniCPMConfig()
+
+    >>> # Initializing a model from the minicpm-7b style configuration
+    >>> model = MiniCPMModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "minicpm"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        pretraining_tp=1,
+        tie_word_embeddings=True,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        scale_emb=1,
+        dim_model_base=1,
+        scale_depth=1,
+        start_layer=8,
+        head_multi=True,
+        head_type="simple",
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.pretraining_tp = pretraining_tp
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self._rope_scaling_validation()
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.scale_emb = scale_emb
+        self.dim_model_base = dim_model_base
+        self.scale_depth = scale_depth
+
+        self.start_layer = start_layer
+        self.head_multi = head_multi
+        self.head_type = head_type
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        try:
+            import flash_attn
+            self._attn_implementation = "flash_attention_2"
+        except:
+            pass
+
+    def _rope_scaling_validation(self):
+        """
+        Validate the `rope_scaling` configuration.
+        """
+        if self.rope_scaling is None:
+            return
+
+        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
+            raise ValueError(
+                "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
+                f"got {self.rope_scaling}"
+            )
+        rope_scaling_type = self.rope_scaling.get("type", None)
+        rope_scaling_factor = self.rope_scaling.get("factor", None)
+        if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
+            raise ValueError(
+                f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
+            )
+        if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
+            raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")

FlagEmbedding/inference/reranker/decoder_only/models/gemma_config.py

@@ -0,0 +1,67 @@
+#           🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#               This file was automatically generated from <path_to_diff_file.py>.
+#         Do NOT edit this file manually as any edits will be overwritten by the generation of
+#         the file from the diff. If any change should be done, please apply the change to the
+#                                    diff.py file directly.
+#           🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from transformers.models.gemma2.configuration_gemma2 import Gemma2Config
+
+class CostWiseGemmaConfig(Gemma2Config):
+    r"""
+    This is the configuration class to store the configuration of a [`GemmaModel`]. It is used to instantiate an Gemma
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the Gemma-7B.
+    e.g. [google/gemma-7b](https://huggingface.co/google/gemma-7b)
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        start_layer (`int`, *optional*, defaults to 28):
+            The start layer to output score.
+        layer_sep (`int`, *optional*, defaults to 28):
+            The sep layer from the start layer to output score.
+        layer_wise (`bool`, *optional*, defaults to `False`):
+            Whether or not the model should be layerwise.
+    ```python
+    >>> from transformers import Gemma2Model, Gemma2Config
+    >>> # Initializing a Gemma2 gemma2-9b style configuration
+    >>> configuration = Gemma2Config()
+    >>> # Initializing a model from the gemma2-9b style configuration
+    >>> model = Gemma2Model(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "cost_wise_gemma"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+            self,
+            start_layer: int = 28,
+            layer_sep: int = 28,
+            layer_wise: bool = False,
+            **kwargs,
+    ):
+        self.start_layer = start_layer
+        self.layer_sep = layer_sep
+        self.layer_wise = layer_wise
+
+        super().__init__(
+            **kwargs,
+        )

FlagEmbedding/inference/reranker/decoder_only/models/gemma_model.py

@@ -0,0 +1,745 @@
+#           🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#               This file was automatically generated from <path_to_diff_file.py>.
+#         Do NOT edit this file manually as any edits will be overwritten by the generation of
+#         the file from the diff. If any change should be done, please apply the change to the
+#                                    diff.py file directly.
+#           🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from dataclasses import dataclass
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import inspect
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache, StaticCache
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+    ModelOutput,
+)
+from .gemma_config import CostWiseGemmaConfig
+from transformers.models.gemma2.modeling_gemma2 import Gemma2RMSNorm, Gemma2RotaryEmbedding, rotate_half, apply_rotary_pos_emb
+from transformers.models.gemma2.modeling_gemma2 import Gemma2MLP, repeat_kv, Gemma2Attention, Gemma2DecoderLayer, GEMMA2_START_DOCSTRING
+from transformers.models.gemma2.modeling_gemma2 import GEMMA2_INPUTS_DOCSTRING
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+    _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
+
+
+logger = logging.get_logger(__name__)
+
+
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+@add_start_docstrings(
+    "The bare Gemma2 Model outputting raw hidden-states without any specific head on top.",
+    GEMMA2_START_DOCSTRING,
+)
+class CostWiseGemma2PreTrainedModel(PreTrainedModel):
+    config_class = CostWiseGemmaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Gemma2DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = False
+    _supports_quantized_cache = False
+    _supports_static_cache = True
+    _is_stateful = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+_CONFIG_FOR_DOC = "CostWiseGemmaConfig"
+
+@dataclass
+class CostWiseModelOutputWithPast(ModelOutput):
+    last_hidden_state: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    attention_masks: Optional[Tuple[torch.FloatTensor]] = None
+
+@dataclass
+class CostWiseCausalLMOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    attention_masks: Optional[Tuple[torch.FloatTensor]] = None
+
+def token_compress(compress_ratio,
+                   hidden_states,
+                   attention_mask,
+                   query_lengths,
+                   prompt_lengths):
+    """
+        compress_ratio: int
+        hidden_states: (b, s, h)
+        attention_mask: (b, s)
+        query_lengths: (b)
+        prompt_lengths: (b)
+    """
+    # get some specific parameters
+    passage_lengths = torch.sum(attention_mask, dim=1, dtype=torch.int) - query_lengths - prompt_lengths # the raw passage lengths (b)
+    retain_passage_lengths = (passage_lengths + compress_ratio - 1) // compress_ratio # the passage lengths need to be retained (b)
+    final_useful_lengths = query_lengths + prompt_lengths + retain_passage_lengths # the final useful length after compress (b)
+    max_passage_length = torch.max(passage_lengths) # the max passage lengths (1)
+    max_final_lengths = torch.max(final_useful_lengths) # the max useful lengths after compress (1)
+    # make new hidden states and new attention masks
+    new_hidden_states = torch.zeros((hidden_states.shape[0], max_final_lengths,
+                                     hidden_states.shape[-1]), dtype=hidden_states.dtype).to(hidden_states.device) # (b, s', h)
+    new_attention_mask = torch.ones((hidden_states.shape[0], max_final_lengths), dtype=attention_mask.dtype).to(attention_mask.device) # (b, s')
+    # get new attention mask
+    mask_attention_index = torch.arange(max_final_lengths, device=hidden_states.device).unsqueeze(0) >= final_useful_lengths[:, None]
+    new_attention_mask[mask_attention_index] = 0
+    # get new hidden states
+    # add query into new hidden states
+    query_index = torch.arange(max_final_lengths, device=hidden_states.device).unsqueeze(0)
+    mask_query_index = query_index < query_lengths[:, None]
+    new_hidden_states[mask_query_index] = hidden_states[:, : max_final_lengths, :][mask_query_index]
+    # add prompt into new hidden states
+    # get the index of the prompt in new hidden states
+    new_prompt_start_length = query_lengths + retain_passage_lengths
+    new_prompt_end_length = new_prompt_start_length + prompt_lengths
+    new_prompt_index = torch.arange(max_final_lengths, device=hidden_states.device).unsqueeze(0)
+    new_mask_prompt_index_start = new_prompt_index >= new_prompt_start_length[:, None]
+    new_mask_prompt_index_end = new_prompt_index < new_prompt_end_length[:, None]
+    new_mask_prompt_index = new_mask_prompt_index_start & new_mask_prompt_index_end
+    # get the index of the prompt in hidden states
+    raw_prompt_start_length = query_lengths + passage_lengths
+    raw_prompt_end_length = raw_prompt_start_length + prompt_lengths
+    raw_prompt_index = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+    raw_mask_prompt_index_start = raw_prompt_index >= raw_prompt_start_length[:, None]
+    raw_mask_prompt_index_end = raw_prompt_index < raw_prompt_end_length[:, None]
+    raw_mask_prompt_index = raw_mask_prompt_index_start & raw_mask_prompt_index_end
+    # replace the prompt hidden states
+    new_hidden_states[new_mask_prompt_index] = hidden_states[raw_mask_prompt_index]
+    # 以上均没问题
+
+    # print(new_hidden_states.view(len(new_hidden_states), -1))
+    # print(new_attention_mask)
+
+    # get the index of the passage in new hidden states
+    new_passage_start_length = query_lengths
+    new_passage_end_length = new_passage_start_length + retain_passage_lengths
+    new_passage_index = torch.arange(max_final_lengths, device=hidden_states.device).unsqueeze(0)
+    new_mask_passage_index_start = new_passage_index >= new_passage_start_length[:, None]
+    new_mask_passage_index_end = new_passage_index < new_passage_end_length[:, None]
+    new_mask_passage_index = new_mask_passage_index_start & new_mask_passage_index_end
+    # print(query_lengths, prompt_lengths, retain_passage_lengths, final_useful_lengths)
+    # add passage into new hidden states
+    # get mask hidden states
+    psg_start_length = query_lengths
+    psg_end_length = query_lengths + passage_lengths
+    psg_index = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+    mask_psg_index_start = psg_index >= psg_start_length[:, None]
+    mask_psg_index_end = psg_index < psg_end_length[:, None]
+    mask_psg_index = mask_psg_index_start & mask_psg_index_end
+
+    hidden_states = hidden_states * mask_psg_index.unsqueeze(-1)
+    passage_hidden_states = torch.zeros((hidden_states.shape[0],
+                                         (max_passage_length + compress_ratio - 1) // compress_ratio * compress_ratio,
+                                         hidden_states.shape[-1]), dtype=hidden_states.dtype).to(hidden_states.device)
+    passage_end_length = passage_lengths
+    passage_index = torch.arange(passage_hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) # maybe exceed the max passage length
+    mask_passage_index = passage_index < passage_end_length[:, None]
+
+    raw_passage_end_length = query_lengths + passage_lengths
+    raw_passage_start_length = query_lengths
+    raw_passage_index = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+    raw_mask_passage_index_start = raw_passage_index >= raw_passage_start_length[:, None]
+    raw_mask_passage_index_end = raw_passage_index < raw_passage_end_length[:, None]
+    raw_mask_passage_index = raw_mask_passage_index_start & raw_mask_passage_index_end
+    passage_hidden_states[mask_passage_index] = hidden_states[raw_mask_passage_index]
+
+    passage_weights = torch.zeros((hidden_states.shape[0],
+                                   (max_passage_length + compress_ratio - 1) // compress_ratio * compress_ratio)
+                                  , dtype=hidden_states.dtype).to(hidden_states.device)
+    passage_weights[mask_passage_index] = 1
+    passage_weights = passage_weights.view(passage_weights.shape[0], -1, compress_ratio)
+    passage_weights = passage_weights / torch.sum(passage_weights, dim=-1
+                                                  ).view(passage_weights.shape[0], -1, 1)
+    passage_weights = passage_weights.view(passage_weights.shape[0], -1)
+    # passage_weights = torch.where(passage_weights == torch.nan, 0, passage_weights)
+    passage_hidden_states = passage_hidden_states * passage_weights.unsqueeze(-1)
+    passage_hidden_states = passage_hidden_states.view(passage_hidden_states.shape[0], -1, compress_ratio,
+                                                       passage_hidden_states.shape[-1])
+    passage_hidden_states = torch.sum(passage_hidden_states, dim=2)
+    passage_end_length = retain_passage_lengths
+    passage_index = torch.arange(passage_hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+    mask_passage_index = passage_index < passage_end_length[:, None]
+    new_hidden_states[new_mask_passage_index] = passage_hidden_states[mask_passage_index]
+
+    return new_hidden_states, new_attention_mask
+
+@add_start_docstrings(
+    "The bare Gemma2 Model outputting raw hidden-states without any specific head on top.",
+    GEMMA2_START_DOCSTRING,
+)
+class CostWiseGemmaModel(CostWiseGemma2PreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`GemmaDecoderLayer`]
+
+    Args:
+        config: GemmaConfig
+    """
+
+    def __init__(self, config: CostWiseGemmaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Gemma2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Gemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(GEMMA2_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        compress_layer: Optional[int] = None,
+        compress_ratio: Optional[int] = None,
+        cutoff_layers: Optional[List[int]] = None,
+        query_lengths: Optional[int] = None,
+        prompt_lengths: Optional[int] = None,
+    ) -> Union[Tuple, CostWiseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
+        compress_ratio = None if compress_ratio == 1 else compress_ratio
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        if self.config.layer_wise:
+            output_hidden_states = True
+
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if compress_layer is not None and compress_ratio is not None:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with reranker. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if cache_position is None:
+            cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device)
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        )
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        # normalized
+        # Gemma downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5
+        # See https://github.com/huggingface/transformers/pull/29402
+        normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)
+        hidden_states = hidden_states * normalizer
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_attention_masks = ()
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        is_padding_left = (attention_mask[:, -1].sum() == attention_mask.shape[0]) and (
+                torch.sum(attention_mask) != attention_mask.shape[0] * attention_mask.shape[1])
+        query_lengths = [0] * hidden_states.shape[0] if query_lengths is None else query_lengths
+        prompt_lengths = [0] * hidden_states.shape[0] if prompt_lengths is None else prompt_lengths
+        if not isinstance(query_lengths, torch.Tensor):
+            query_lengths = torch.tensor(query_lengths, device=hidden_states.device)
+        if not isinstance(prompt_lengths, torch.Tensor):
+            prompt_lengths = torch.tensor(prompt_lengths, device=hidden_states.device)
+
+        if cutoff_layers is None:
+            max_layer = self.config.num_hidden_layers
+            cutoff_layers = [max_layer]
+        if isinstance(cutoff_layers, int):
+            max_layer = cutoff_layers
+            cutoff_layers = [cutoff_layers]
+        else:
+            max_layer = max(cutoff_layers)
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if self.config.layer_wise:
+                if idx in cutoff_layers and output_hidden_states:
+                    all_hidden_states += (self.norm(hidden_states),)
+                    all_attention_masks += (attention_mask,)
+                if idx == max_layer:
+                    break
+            elif output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if compress_layer is not None and compress_ratio is not None and idx in compress_layer and idx != 0:
+                if is_padding_left:
+                    raise ValueError('You must use right padding...')
+                hidden_states, attention_mask = token_compress(compress_ratio, hidden_states, attention_mask,
+                                                               query_lengths, prompt_lengths)
+                seq_length = hidden_states.shape[1]
+                cache_position = torch.arange(0, seq_length, device=hidden_states.device)
+                position_ids = cache_position.unsqueeze(0)
+                causal_mask = self._update_causal_mask(
+                    attention_mask, hidden_states, cache_position, past_key_values, output_attentions
+                )
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if not self.config.layer_wise:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+                all_attention_masks += (attention_mask,)
+        else:
+            if output_hidden_states and self.config.num_hidden_layers == max_layer:
+                all_hidden_states += (hidden_states,)
+                all_attention_masks += (attention_mask,)
+
+        next_cache = next_decoder_cache if use_cache else None
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return CostWiseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            attention_masks=all_attention_masks
+        )
+
+    def _update_causal_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_tensor: torch.Tensor,
+        cache_position: torch.Tensor,
+        past_key_values: Cache,
+        output_attentions: bool,
+    ):
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+
+        dtype, device = input_tensor.dtype, input_tensor.device
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        if past_key_values is not None:
+            target_length = past_key_values.get_max_length()
+        else:
+            target_length = attention_mask.shape[-1] if attention_mask is not None else input_tensor.shape[1]
+
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing
+            if attention_mask.max() != 0:
+                raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`")
+            causal_mask = attention_mask
+        else:
+            causal_mask = torch.full(
+                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
+            )
+            if sequence_length != 1:
+                causal_mask = torch.triu(causal_mask, diagonal=1)
+            causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+            causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
+            if attention_mask is not None:
+                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
+                padding_mask = padding_mask == 0
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    padding_mask, min_dtype
+                )
+        return causal_mask
+
+
+class CostWiseHead(nn.Module):
+    """Head for sentence-level classification tasks."""
+
+    def __init__(self, input_size, output_size):
+        super().__init__()
+        self.linear_head = nn.Linear(input_size, output_size, bias=False)
+
+    def forward(self, **kwargs):
+        return self.linear_head(**kwargs)
+
+
+class CostWiseGemmaForCausalLM(CostWiseGemma2PreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: CostWiseGemmaConfig):
+        super().__init__(config)
+        self.model = CostWiseGemmaModel(config)
+        self.vocab_size = config.vocab_size
+
+        if not config.layer_wise:
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        else:
+            self.lm_head = nn.ModuleList(
+                [CostWiseHead(config.hidden_size, 1) for _ in range(
+                    config.start_layer, config.num_hidden_layers + 1, config.layer_sep
+                )]
+            )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(GEMMA2_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        compress_layer: Optional[int] = None,
+        compress_ratio: Optional[int] = None,
+        cutoff_layers: Optional[List[int]] = None,
+        query_lengths: Optional[int] = None,
+        prompt_lengths: Optional[int] = None,
+    ) -> Union[Tuple, CostWiseCausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers.,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+         ```python
+        >>> from transformers import AutoTokenizer, GemmaForCausalLM
+
+        >>> model = GemmaForCausalLM.from_pretrained("google/gemma-2-9b")
+        >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
+
+        >>> prompt = "What is your favorite condiment?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "What is your favorite condiment?"
+        ```"""
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if compress_ratio is not None and compress_ratio == 1:
+            compress_ratio = None
+
+        if self.config.layer_wise:
+            if cutoff_layers is None:
+                cutoff_layers = [self.config.num_hidden_layers]
+            elif isinstance(cutoff_layers, int):
+                cutoff_layers = [cutoff_layers]
+            can_use_layers = list(range(self.config.start_layer, self.config.num_hidden_layers + 1, self.config.layer_sep))
+            remove_layers = [i for i in cutoff_layers if i not in can_use_layers]
+            if len(remove_layers) > 0:
+                logger.warning_once(
+                    f"layers {remove_layers} are incompatible with the setting. They will be removed..."
+                )
+            cutoff_layers = [i for i in cutoff_layers if i not in remove_layers]
+            if len(cutoff_layers) == 0:
+                raise ValueError(f"Your cutoff layers must in [{self.config.start_layer}, {self.config.num_hidden_layers}]")
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            compress_layer=compress_layer,
+            compress_ratio=compress_ratio,
+            query_lengths=query_lengths,
+            prompt_lengths=prompt_lengths,
+            cutoff_layers=cutoff_layers,
+        )
+
+        if not self.config.layer_wise:
+            hidden_states = outputs[0]
+            logits = self.lm_head(hidden_states)
+            if self.config.final_logit_softcapping is not None:
+                logits = logits / self.config.final_logit_softcapping
+                logits = torch.tanh(logits)
+                logits = logits * self.config.final_logit_softcapping
+            logits = logits.float()
+            loss = None
+            if labels is not None:
+                # Shift so that tokens < n predict n
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                # Flatten the tokens
+                loss_fct = CrossEntropyLoss()
+                shift_logits = shift_logits.view(-1, self.config.vocab_size)
+                shift_labels = shift_labels.view(-1)
+                # Enable model parallelism
+                shift_labels = shift_labels.to(shift_logits.device)
+                loss = loss_fct(shift_logits, shift_labels)
+        else:
+            hidden_states = outputs.hidden_states
+            logits = ()
+            for i in range(len(hidden_states)):
+                tmp_logits = self.lm_head[i].linear_head(hidden_states[i])
+                if self.config.final_logit_softcapping is not None:
+                    tmp_logits = tmp_logits / self.config.final_logit_softcapping
+                    tmp_logits = torch.tanh(tmp_logits)
+                    tmp_logits = tmp_logits * self.config.final_logit_softcapping
+                tmp_logits = tmp_logits.float()
+                tmp_logits = tmp_logits.reshape(hidden_states[i].shape[0], -1)
+                logits = logits + (tmp_logits,)
+            loss = None
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CostWiseCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            attention_masks=outputs[-1] if self.model.config.layer_wise else outputs[-1][-1]
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        past_length = 0
+        if past_key_values is not None:
+            # Past key values are always initialized with a `Cache` object -> no need for if-else anymore
+            past_length = cache_position[0] if cache_position is not None else torch.tensor(0, device=input_ids.device)
+            max_cache_length = (
+                torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
+                if past_key_values.get_max_length() is not None
+                else None
+            )
+            cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_length == 0:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
+            # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
+            # TODO: use `next_tokens` directly instead.
+            model_inputs = {"input_ids": input_ids.contiguous()}
+
+        input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
+        if cache_position is None:
+            cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
+        elif use_cache:
+            cache_position = cache_position[-input_length:]
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past

FlagEmbedding/inference/reranker/encoder_only/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseReranker as FlagReranker
+
+__all__ = [
+    "FlagReranker",
+]

FlagEmbedding/inference/reranker/encoder_only/base.py

@@ -0,0 +1,195 @@
+import torch
+import numpy as np
+from tqdm import tqdm, trange
+from typing import Any, List, Union, Tuple, Optional
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+from FlagEmbedding.abc.inference import AbsReranker
+
+
+def sigmoid(x):
+    return float(1 / (1 + np.exp(-x)))
+
+
+class BaseReranker(AbsReranker):
+    """Base reranker class for encoder only models.
+
+    Args:
+        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 :data:`False`.
+        query_instruction_for_rerank (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
+            with :attr:`query_instruction_format`. Defaults to :data:`None`.
+        query_instruction_format (str, optional): The template for :attr:`query_instruction_for_rerank`. Defaults to :data:`"{}{}"`.
+        passage_instruction_format (str, optional): The template for passage. Defaults to "{}{}".
+        cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
+        devices (Optional[Union[str, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
+        batch_size (int, optional): Batch size for inference. Defaults to :data:`128`.
+        query_max_length (Optional[int], optional): Maximum length for queries. If not specified, will be 3/4 of :attr:`max_length`.
+            Defaults to :data:`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 :data:`False`.
+    """
+    def __init__(
+        self,
+        model_name_or_path: str,
+        use_fp16: bool = False,
+        query_instruction_for_rerank: Optional[str] = None,
+        query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_rerank
+        passage_instruction_for_rerank: Optional[str] = None,
+        passage_instruction_format: str = "{}{}", # specify the format of passage_instruction_for_rerank
+        trust_remote_code: bool = False,
+        cache_dir: Optional[str] = None,
+        devices: Optional[Union[str, List[str], List[int]]] = None, # specify devices, such as ["cuda:0"] or ["0"]
+        # inference
+        batch_size: int = 128,
+        query_max_length: Optional[int] = None,
+        max_length: int = 512,
+        normalize: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_name_or_path=model_name_or_path,
+            use_fp16=use_fp16,
+            query_instruction_for_rerank=query_instruction_for_rerank,
+            query_instruction_format=query_instruction_format,
+            passage_instruction_for_rerank=passage_instruction_for_rerank,
+            passage_instruction_format=passage_instruction_format,
+            devices=devices,
+            batch_size=batch_size,
+            query_max_length=query_max_length,
+            max_length=max_length,
+            normalize=normalize,
+            **kwargs
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path, 
+            trust_remote_code=trust_remote_code, 
+            cache_dir=cache_dir
+        )
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            model_name_or_path, 
+            trust_remote_code=trust_remote_code, 
+            cache_dir=cache_dir
+        )
+
+    @torch.no_grad()
+    def compute_score_single_gpu(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: Optional[int] = None,
+        query_max_length: Optional[int] = None,
+        max_length: Optional[int] = None,
+        normalize: Optional[bool] = None,
+        device: Optional[str] = None,
+        **kwargs: Any
+    ) -> List[float]:
+        """_summary_
+
+        Args:
+            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 :data:`None`.
+            query_max_length (Optional[int], optional): Maximum length of tokens of queries. Defaults to :data:`None`.
+            max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
+            normalize (Optional[bool], optional): If True, use Sigmoid to normalize the results. Defaults to :data:`None`.
+            device (Optional[str], optional): Device to use for computation. Defaults to :data:`None`.
+
+        Returns:
+            List[float]: Computed scores of queries and passages.
+        """
+        if batch_size is None: batch_size = self.batch_size
+        if max_length is None: max_length = self.max_length
+        if query_max_length is None:
+            if self.query_max_length is not None:
+                query_max_length = self.query_max_length
+            else:
+                query_max_length = max_length * 3 // 4
+        if normalize is None: normalize = self.normalize
+
+        if device is None:
+            device = self.target_devices[0]
+
+        if device == "cpu": self.use_fp16 = False
+        if self.use_fp16: self.model.half()
+
+        self.model.to(device)
+        self.model.eval()
+
+        assert isinstance(sentence_pairs, list)
+        if isinstance(sentence_pairs[0], str):
+            sentence_pairs = [sentence_pairs]
+        
+        # tokenize without padding to get the correct length
+        all_inputs = []
+        for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
+                                  disable=len(sentence_pairs) < batch_size):
+            sentences_batch = sentence_pairs[start_index:start_index + batch_size]
+            queries = [s[0] for s in sentences_batch]
+            passages = [s[1] for s in sentences_batch]
+            queries_inputs_batch = self.tokenizer(
+                queries,
+                return_tensors=None,
+                add_special_tokens=False,
+                max_length=query_max_length,
+                truncation=True,
+                **kwargs
+            )['input_ids']
+            passages_inputs_batch = self.tokenizer(
+                passages,
+                return_tensors=None,
+                add_special_tokens=False,
+                max_length=max_length,
+                truncation=True,
+                **kwargs
+            )['input_ids']
+            for q_inp, d_inp in zip(queries_inputs_batch, passages_inputs_batch):
+                item = self.tokenizer.prepare_for_model(
+                    q_inp,
+                    d_inp,
+                    truncation='only_second',
+                    max_length=max_length,
+                    padding=False,
+                )
+                all_inputs.append(item)
+        # sort by length for less padding
+        length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
+        all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
+
+        # adjust batch size
+        flag = False
+        while flag is False:
+            try:
+                test_inputs_batch = self.tokenizer.pad(
+                    all_inputs_sorted[:min(len(all_inputs_sorted), batch_size)],
+                    padding=True,
+                    return_tensors='pt',
+                    **kwargs
+                ).to(device)
+                scores = self.model(**test_inputs_batch, return_dict=True).logits.view(-1, ).float()
+                flag = True
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.cuda.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
+        all_scores = []
+        for start_index in tqdm(range(0, len(all_inputs_sorted), batch_size), desc="Compute Scores",
+                                disable=len(all_inputs_sorted) < batch_size):
+            sentences_batch = all_inputs_sorted[start_index:start_index + batch_size]
+            inputs = self.tokenizer.pad(
+                sentences_batch,
+                padding=True,
+                return_tensors='pt',
+                **kwargs
+            ).to(device)
+
+            scores = self.model(**inputs, return_dict=True).logits.view(-1, ).float()
+            all_scores.extend(scores.cpu().numpy().tolist())
+
+        all_scores = [all_scores[idx] for idx in np.argsort(length_sorted_idx)]
+
+        if normalize:
+            all_scores = [sigmoid(score) for score in all_scores]
+
+        return all_scores

FlagEmbedding/inference/reranker/model_mapping.py

@@ -0,0 +1,75 @@
+from enum import Enum
+from typing import Type
+from dataclasses import dataclass
+from collections import OrderedDict
+
+from FlagEmbedding.abc.inference import AbsReranker
+from FlagEmbedding.inference.reranker import FlagReranker, FlagLLMReranker, LayerWiseFlagLLMReranker, LightWeightFlagLLMReranker
+
+
+class RerankerModelClass(Enum):
+    ENCODER_ONLY_BASE = "encoder-only-base"
+    DECODER_ONLY_BASE = "decoder-only-base"
+    DECODER_ONLY_LAYERWISE = "decoder-only-layerwise"
+    DECODER_ONLY_LIGHTWEIGHT = "decoder-only-lightweight"
+
+
+RERANKER_CLASS_MAPPING = OrderedDict([
+    (RerankerModelClass.ENCODER_ONLY_BASE, FlagReranker),
+    (RerankerModelClass.DECODER_ONLY_BASE, FlagLLMReranker),
+    (RerankerModelClass.DECODER_ONLY_LAYERWISE, LayerWiseFlagLLMReranker),
+    (RerankerModelClass.DECODER_ONLY_LIGHTWEIGHT, LightWeightFlagLLMReranker)
+])
+
+
+@dataclass
+class RerankerConfig:
+    model_class: Type[AbsReranker]
+    trust_remote_code: bool = False
+
+
+AUTO_RERANKER_MAPPING = OrderedDict([
+    # ============================== BGE ==============================
+    (
+        "bge-reranker-base", 
+        RerankerConfig(FlagReranker)
+    ),
+    (
+        "bge-reranker-large", 
+        RerankerConfig(FlagReranker)
+    ),
+    (
+        "bge-reranker-v2-m3",
+        RerankerConfig(FlagReranker)
+    ),
+    (
+        "bge-reranker-v2-gemma",
+        RerankerConfig(FlagLLMReranker)
+    ),
+    (
+        "bge-reranker-v2-minicpm-layerwise",
+        RerankerConfig(LayerWiseFlagLLMReranker)
+    ),
+    (
+        "bge-reranker-v2.5-gemma2-lightweight",
+        RerankerConfig(LightWeightFlagLLMReranker)
+    ),
+    # others
+    (
+        "jinaai/jina-reranker-v2-base-multilingual",
+        RerankerConfig(FlagReranker)
+    ),
+    (
+        "Alibaba-NLP/gte-multilingual-reranker-base",
+        RerankerConfig(FlagReranker)
+    ),
+    (
+        "maidalun1020/bce-reranker-base_v1",
+        RerankerConfig(FlagReranker)
+    ),
+    (
+        "jinaai/jina-reranker-v1-turbo-en",
+        RerankerConfig(FlagReranker)
+    ),
+    # TODO: Add more models.
+])