Support bge-m3 sparse embeddings

Now with the pooling task framework

Signed-off-by: Max de Bayser <mbayser@br.ibm.com>

Author: maxdebayser

vllm/entrypoints/openai/protocol.py

@@ -1366,8 +1366,25 @@ def to_pooling_params(self):
 
 EmbeddingRequest = Union[EmbeddingCompletionRequest, EmbeddingChatRequest]
 
-PoolingCompletionRequest = EmbeddingCompletionRequest
-PoolingChatRequest = EmbeddingChatRequest
+
+class PoolingCompletionRequest(EmbeddingCompletionRequest):
+    task: Optional[str] = None
+
+    def to_pooling_params(self):
+        return PoolingParams(dimensions=self.dimensions,
+                             normalize=self.normalize,
+                             task=self.task)
+
+
+class PoolingChatRequest(EmbeddingChatRequest):
+    task: Optional[str] = None
+
+    def to_pooling_params(self):
+        return PoolingParams(dimensions=self.dimensions,
+                             normalize=self.normalize,
+                             task=self.task)
+
+
 PoolingRequest = Union[PoolingCompletionRequest, PoolingChatRequest]
 
 

vllm/entrypoints/openai/serving_pooling.py

@@ -140,7 +140,8 @@ async def create_pooling(
             pooling_params = request.to_pooling_params()
 
             try:
-                pooling_params.verify("encode", self.model_config)
+                task = request.task if request.task is not None else "encode"
+                pooling_params.verify(task, self.model_config)
             except ValueError as e:
                 return self.create_error_response(str(e))
 

vllm/model_executor/models/registry.py

@@ -173,6 +173,7 @@
     "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
     "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
     "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
+    "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
     # [Multimodal]
     "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
     "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),

vllm/model_executor/models/roberta.py

@@ -1,26 +1,33 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import itertools
 from collections.abc import Iterable
 from typing import Optional, Union
 
 import torch
 from torch import nn
 from transformers import RobertaConfig
 
-from vllm.config import VllmConfig
+from vllm.config import PoolerConfig, VllmConfig
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.pooler import (ClassifierPooler, CLSPool,
-                                               DispatchPooler, Pooler)
+                                               DispatchPooler, Pooler,
+                                               PoolerOutput, PoolingMetadata,
+                                               PoolingParamsUpdate,
+                                               PoolingTask, build_output)
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding)
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.bert import (TOKEN_TYPE_SHIFT,
                                              BertEmbeddingModel, BertModel,
                                              _decode_token_type_ids,
                                              _encode_token_type_ids)
 from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
                                               maybe_prefix)
 from vllm.sequence import IntermediateTensors
+from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata
 
 from .bert_with_rope import BertWithRope, JinaRobertaModel
 from .interfaces import SupportsCrossEncoding, default_pooling_type
@@ -150,6 +157,130 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         return loader.load_weights(weights_list, mapper=mapper)
 
 
+class M3SparsePooler(Pooler):
+    """A pooler that implements M3 sparse pooling
+
+    This layer does the following:
+    1. By default returns dense embeddings.
+    2. If the pooling params "additional_data" contain
+       "sparse_embeddings", return sparse embeddings
+
+    Attributes:
+        dense_pooler: The default pooler.
+        sparse_linear: the linear module applied to the
+          logits to obtain the token weights
+        bos_token_id and eos_token_id: The special tokens
+          inserted by the tokenizer. These are removed for
+          sparse embeddings
+    """
+
+    def __init__(self, sparse_linear: nn.Module, bos_token_id: int,
+                 eos_token_id: int) -> None:
+        super().__init__()
+        self.sparse_linear = sparse_linear
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+
+    def get_supported_tasks(self) -> set[PoolingTask]:
+        return {"embed-sparse"}
+
+    def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
+        return PoolingParamsUpdate(requires_token_ids=True)
+
+    def forward(
+        self,
+        hidden_states: Union[torch.Tensor, list[torch.Tensor]],
+        pooling_metadata: PoolingMetadata,
+    ) -> PoolerOutput:
+
+        assert isinstance(pooling_metadata, V1PoolingMetadata), \
+            "BGE-M3 sparse embeddding are only support with V1"
+        assert isinstance(hidden_states, list)
+
+        pooled_outputs = []
+
+        for i, hidden_state in enumerate(hidden_states):
+            pooled_data = torch.squeeze(torch.relu(
+                self.sparse_linear(hidden_state)),
+                                        dim=0)
+            token_ids = pooling_metadata.prompt_token_ids[
+                i, :pooling_metadata.prompt_lens[i]]
+            if token_ids[0] == self.bos_token_id:
+                pooled_data = pooled_data[1:]
+            if token_ids[-1] == self.eos_token_id:
+                pooled_data = pooled_data[:-1]
+            pooled_outputs.append(pooled_data)
+
+        return PoolerOutput(outputs=build_output(pooled_outputs))
+
+
+def filter_secondary_weights(
+    all_weights: Iterable[tuple[str, torch.Tensor]],
+    secondary_weights: list[str],
+) -> tuple[Iterable[tuple[str, torch.Tensor]], Iterable[tuple[str,
+                                                              torch.Tensor]]]:
+    all_weights1, all_weights2 = itertools.tee(all_weights)
+
+    def filtered(n):
+        return any(n.startswith(f) for f in secondary_weights)
+
+    return ((n, w) for n, w in all_weights1 if filtered(n)), \
+           ((n, w) for n, w in all_weights2 if not filtered(n))
+
+
+class BgeM3EmbeddingModel(RobertaEmbeddingModel):
+    """A model that extends RobertaEmbeddingModel with sparse embeddings.
+
+   This class supports loading an additional sparse_linear.pt file
+   to create sparse embeddings as described in https://arxiv.org/abs/2402.03216
+   """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+
+        self.hidden_size = vllm_config.model_config.hf_config.hidden_size
+
+        self.bos_token_id = vllm_config.model_config.hf_config.bos_token_id
+        self.eos_token_id = vllm_config.model_config.hf_config.eos_token_id
+
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        self.secondary_weight_prefix = "sparse_linear."
+
+        self.secondary_weights = [
+            DefaultModelLoader.Source(
+                model_or_path=vllm_config.model_config.model,
+                revision=None,
+                prefix=self.secondary_weight_prefix,
+                allow_patterns_overrides=["sparse_linear.pt"])
+        ]
+
+    def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
+        self.sparse_linear = nn.Linear(self.hidden_size, 1)
+        return DispatchPooler({
+            "encode":
+            Pooler.for_encode(pooler_config),
+            "embed":
+            Pooler.for_embed(pooler_config),
+            "embed-sparse":
+            M3SparsePooler(self.sparse_linear, self.bos_token_id,
+                           self.eos_token_id),
+        })
+
+    def load_weights(self, all_weights: Iterable[tuple[str, torch.Tensor]]):
+        secondary, weights = filter_secondary_weights(
+            all_weights, [self.secondary_weight_prefix])
+
+        super().load_weights(weights)
+
+        params_dict = dict(self.named_parameters())
+
+        for name, loaded_weight in secondary:
+            if name.startswith(self.secondary_weight_prefix):
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+
 @default_pooling_type("CLS")
 class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
     """A model that uses Roberta to provide embedding functionalities.

vllm/pooling_params.py

@@ -62,6 +62,7 @@ def all_parameters(self) -> list[str]:
     def valid_parameters(self):
         return {
             "embed": ["dimensions", "normalize"],
+            "embed-sparse": ["dimensions", "normalize"],
             "classify": ["activation"],
             "score": ["activation"],
             "encode": ["softmax", "step_tag_id", "returned_token_ids"],
@@ -111,7 +112,7 @@ def _merge_default_parameters(self,
                 setattr(self, k, getattr(pooler_config, k))
 
     def _set_default_parameters(self, model_config: Optional["ModelConfig"]):
-        if self.task == "embed":
+        if self.task in ["embed", "embed-sparse"]:
             if self.normalize is None:
                 self.normalize = True
 

vllm/tasks.py

@@ -5,7 +5,7 @@
 GenerationTask = Literal["generate", "transcription"]
 GENERATION_TASKS = get_args(GenerationTask)
 
-PoolingTask = Literal["encode", "embed", "classify", "score"]
+PoolingTask = Literal["encode", "embed", "embed-sparse", "classify", "score"]
 POOLING_TASKS = get_args(PoolingTask)
 
 SupportedTask = Literal[GenerationTask, PoolingTask]