[Core] Set pooling params based on task and model (vllm-project#21128)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Author: DarkLight1337

tests/models/language/pooling/test_gritlm.py

@@ -2,9 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from __future__ import annotations
 
-import importlib.util
-from array import array
-
+import numpy as np
 import openai
 import pytest
 from scipy.spatial.distance import cosine
@@ -14,10 +12,6 @@
 
 from ....utils import RemoteOpenAIServer
 
-# GritLM embedding implementation is only supported by XFormers backend.
-pytestmark = pytest.mark.skipif(not importlib.util.find_spec("xformers"),
-                                reason="GritLM requires XFormers")
-
 MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
 MAX_MODEL_LEN = 4000
 
@@ -26,11 +20,11 @@ def _arr(arr):
     """
     Convert a list of integers to an array of integers.
     """
-    return array("i", arr)
+    return np.array(arr)
 
 
 def test_find_array():
-    from vllm.model_executor.models.gritlm import GritLMPooler
+    from vllm.model_executor.models.gritlm import GritLMMeanPool
 
     model_config = ModelConfig(
         MODEL_NAME,
@@ -41,17 +35,19 @@ def test_find_array():
         dtype="bfloat16",
         seed=0,
     )
-    pooler = GritLMPooler(model_config=model_config)
+    pooling = GritLMMeanPool(model_config=model_config)
 
     arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
 
-    assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3
-    assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3
-    assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1
-    assert pooler._find_array(arr, _arr([3, 5]), start_idx=0) == -1
+    assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3
+    assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3
+    assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1
+    assert pooling._find_array(arr, _arr([3, 4, 5]), end_idx=3) == -1
+    assert pooling._find_array(arr, _arr([3, 4, 5]), end_idx=4) == 3
+    assert pooling._find_array(arr, _arr([3, 5]), start_idx=0) == -1
 
     with pytest.raises(ValueError):
-        pooler._find_array(arr, _arr([3, 4, 5]), start_idx=-1)
+        pooling._find_array(arr, _arr([3, 4, 5]), start_idx=-1)
 
 
 def run_llm_encode(

vllm/entrypoints/llm.py

@@ -44,7 +44,7 @@
 from vllm.outputs import (ClassificationRequestOutput, EmbeddingRequestOutput,
                           PoolingRequestOutput, RequestOutput,
                           ScoringRequestOutput)
-from vllm.pooling_params import PoolingParams
+from vllm.pooling_params import PoolingParams, PoolingTask
 from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import (BeamSearchParams, GuidedDecodingParams,
                                   RequestOutputKind, SamplingParams)
@@ -964,6 +964,7 @@ def encode(
         use_tqdm: Union[bool, Callable[..., tqdm]] = True,
         lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        pooling_task: PoolingTask = "encode",
     ) -> list[PoolingRequestOutput]:
         ...
 
@@ -979,6 +980,7 @@ def encode(
         use_tqdm: Union[bool, Callable[..., tqdm]] = True,
         lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        pooling_task: PoolingTask = "encode",
     ) -> list[PoolingRequestOutput]:
         ...
 
@@ -994,6 +996,7 @@ def encode(
         use_tqdm: Union[bool, Callable[..., tqdm]] = True,
         lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        pooling_task: PoolingTask = "encode",
     ) -> list[PoolingRequestOutput]:
         ...
 
@@ -1010,6 +1013,7 @@ def encode(
         use_tqdm: Union[bool, Callable[..., tqdm]] = True,
         lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        pooling_task: PoolingTask = "encode",
     ) -> list[PoolingRequestOutput]:
         ...
 
@@ -1026,6 +1030,7 @@ def encode(
         use_tqdm: Union[bool, Callable[..., tqdm]] = True,
         lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        pooling_task: PoolingTask = "encode",
     ) -> list[PoolingRequestOutput]:
         ...
 
@@ -1040,6 +1045,7 @@ def encode(
         use_tqdm: Union[bool, Callable[..., tqdm]] = True,
         lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        pooling_task: PoolingTask = "encode",
     ) -> list[PoolingRequestOutput]:
         ...
 
@@ -1059,6 +1065,7 @@ def encode(
         use_tqdm: Union[bool, Callable[..., tqdm]] = True,
         lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        pooling_task: PoolingTask = "encode",
     ) -> list[PoolingRequestOutput]:
         """Apply pooling to the hidden states corresponding to the input
         prompts.
@@ -1080,6 +1087,7 @@ def encode(
             lora_request: LoRA request to use for generation, if any.
             prompt_adapter_request: Prompt Adapter request to use for
                 generation, if any.
+            pooling_task: Override the pooling task to use.
 
         Returns:
             A list of `PoolingRequestOutput` objects containing the
@@ -1116,11 +1124,12 @@ def encode(
         if pooling_params is None:
             # Use default pooling params.
             pooling_params = PoolingParams()
-        elif isinstance(pooling_params, PoolingParams):
-            pooling_params.verify(model_config)
+
+        if isinstance(pooling_params, PoolingParams):
+            pooling_params.verify(pooling_task, model_config)
         else:
             for pooling_param in pooling_params:
-                pooling_param.verify(model_config)
+                pooling_param.verify(pooling_task, model_config)
 
         tokenization_kwargs = dict[str, Any]()
         _validate_truncation_size(model_config.max_model_len,
@@ -1181,12 +1190,15 @@ def embed(
             raise ValueError("Embedding API is not supported by this model. "
                              "Please set `--task embed`.")
 
-        items = self.encode(prompts,
-                            truncate_prompt_tokens=truncate_prompt_tokens,
-                            use_tqdm=use_tqdm,
-                            pooling_params=pooling_params,
-                            lora_request=lora_request,
-                            prompt_adapter_request=prompt_adapter_request)
+        items = self.encode(
+            prompts,
+            truncate_prompt_tokens=truncate_prompt_tokens,
+            use_tqdm=use_tqdm,
+            pooling_params=pooling_params,
+            lora_request=lora_request,
+            prompt_adapter_request=prompt_adapter_request,
+            pooling_task="embed",
+        )
 
         return [EmbeddingRequestOutput.from_base(item) for item in items]
 
@@ -1228,10 +1240,13 @@ def classify(
                 "Classification API is not supported by this model. "
                 "Please set `--task classify`.")
 
-        items = self.encode(prompts,
-                            use_tqdm=use_tqdm,
-                            lora_request=lora_request,
-                            prompt_adapter_request=prompt_adapter_request)
+        items = self.encode(
+            prompts,
+            use_tqdm=use_tqdm,
+            lora_request=lora_request,
+            prompt_adapter_request=prompt_adapter_request,
+            pooling_task="classify",
+        )
 
         return [ClassificationRequestOutput.from_base(item) for item in items]
 
@@ -1251,7 +1266,9 @@ def _embedding_score(
             truncate_prompt_tokens=truncate_prompt_tokens,
             use_tqdm=use_tqdm,
             lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request)
+            prompt_adapter_request=prompt_adapter_request,
+            pooling_task="embed",
+        )
 
         encoded_output_1: list[PoolingRequestOutput] = encoded_output[
             0:len(text_1)]
@@ -1287,7 +1304,7 @@ def _cross_encoding_score(
         if len(data_1) == 1:
             data_1 = data_1 * len(data_2)
 
-        pooling_params = PoolingParams(use_cross_encoder=True)
+        pooling_params = PoolingParams(task="score")
         tokenization_kwargs: dict[str, Any] = {}
         _validate_truncation_size(self.llm_engine.model_config.max_model_len,
                                   truncate_prompt_tokens, tokenization_kwargs)

vllm/entrypoints/openai/protocol.py

@@ -1347,8 +1347,8 @@ class ScoreRequest(OpenAIBaseModel):
 
     # --8<-- [end:score-extra-params]
 
-    def to_pooling_params(self, *, use_cross_encoder: bool = False):
-        return PoolingParams(use_cross_encoder=use_cross_encoder)
+    def to_pooling_params(self):
+        return PoolingParams()
 
 
 class RerankRequest(OpenAIBaseModel):
@@ -1375,8 +1375,8 @@ class RerankRequest(OpenAIBaseModel):
 
     # --8<-- [end:rerank-extra-params]
 
-    def to_pooling_params(self, *, use_cross_encoder: bool = False):
-        return PoolingParams(use_cross_encoder=use_cross_encoder)
+    def to_pooling_params(self):
+        return PoolingParams()
 
 
 class RerankDocument(BaseModel):

vllm/entrypoints/openai/serving_classification.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 from fastapi import Request
+from typing_extensions import override
 
 from vllm.config import ModelConfig
 from vllm.engine.protocol import EngineClient
@@ -21,12 +22,14 @@
 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
 from vllm.logger import init_logger
 from vllm.outputs import ClassificationOutput, PoolingRequestOutput
+from vllm.pooling_params import PoolingParams
 
 logger = init_logger(__name__)
 
 
 class ClassificationMixin(OpenAIServing):
 
+    @override
     async def _preprocess(
         self,
         ctx: ServeContext,
@@ -75,6 +78,7 @@ async def _preprocess(
             logger.exception("Error in preprocessing prompt inputs")
             return self.create_error_response(str(e))
 
+    @override
     def _build_response(
         self,
         ctx: ServeContext,
@@ -158,3 +162,31 @@ async def create_classify(
         )
 
         return await super().handle(ctx)  # type: ignore
+
+    @override
+    def _validate_request(
+        self,
+        ctx: ClassificationServeContext,
+    ) -> Optional[ErrorResponse]:
+        if error := super()._validate_request(ctx):
+            return error
+
+        ctx.truncate_prompt_tokens = ctx.request.truncate_prompt_tokens
+
+        return None
+
+    @override
+    def _create_pooling_params(
+        self,
+        ctx: ClassificationServeContext,
+    ) -> Union[PoolingParams, ErrorResponse]:
+        pooling_params = super()._create_pooling_params(ctx)
+        if isinstance(pooling_params, ErrorResponse):
+            return pooling_params
+
+        try:
+            pooling_params.verify("classify", self.model_config)
+        except ValueError as e:
+            return self.create_error_response(str(e))
+
+        return pooling_params

vllm/entrypoints/openai/serving_embedding.py

@@ -32,7 +32,8 @@
 from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
 from vllm.logger import init_logger
 from vllm.outputs import (EmbeddingOutput, EmbeddingRequestOutput,
-                          PoolingRequestOutput, RequestOutput)
+                          PoolingRequestOutput)
+from vllm.pooling_params import PoolingParams
 
 logger = init_logger(__name__)
 
@@ -54,6 +55,7 @@ def _get_embedding(
 
 class EmbeddingMixin(OpenAIServing):
 
+    @override
     async def _preprocess(
         self,
         ctx: ServeContext,
@@ -106,6 +108,7 @@ async def _preprocess(
             logger.exception("Error in preprocessing prompt inputs")
             return self.create_error_response(str(e))
 
+    @override
     def _build_response(
         self,
         ctx: ServeContext,
@@ -906,11 +909,20 @@ def _validate_request(
 
         ctx.truncate_prompt_tokens = ctx.request.truncate_prompt_tokens
 
-        pooling_params = ctx.request.to_pooling_params()
+        return None
+
+    @override
+    def _create_pooling_params(
+        self,
+        ctx: ServeContext[EmbeddingRequest],
+    ) -> Union[PoolingParams, ErrorResponse]:
+        pooling_params = super()._create_pooling_params(ctx)
+        if isinstance(pooling_params, ErrorResponse):
+            return pooling_params
 
         try:
-            pooling_params.verify(self.model_config)
+            pooling_params.verify("embed", self.model_config)
         except ValueError as e:
             return self.create_error_response(str(e))
 
-        return None
+        return pooling_params

vllm/entrypoints/openai/serving_engine.py

@@ -305,6 +305,16 @@ def _validate_request(self, ctx: ServeContext) -> Optional[ErrorResponse]:
                     " Please, select a smaller truncation size.")
         return None
 
+    def _create_pooling_params(
+        self,
+        ctx: ServeContext,
+    ) -> Union[PoolingParams, ErrorResponse]:
+        if not hasattr(ctx.request, "to_pooling_params"):
+            return self.create_error_response(
+                "Request type does not support pooling parameters")
+
+        return ctx.request.to_pooling_params()
+
     async def _prepare_generators(
         self,
         ctx: ServeContext,
@@ -318,11 +328,9 @@ async def _prepare_generators(
             trace_headers = (None if ctx.raw_request is None else await
                              self._get_trace_headers(ctx.raw_request.headers))
 
-            if not hasattr(ctx.request, "to_pooling_params"):
-                return self.create_error_response(
-                    "Request type does not support pooling parameters")
-
-            pooling_params = ctx.request.to_pooling_params()
+            pooling_params = self._create_pooling_params(ctx)
+            if isinstance(pooling_params, ErrorResponse):
+                return pooling_params
 
             if ctx.engine_prompts is None:
                 return self.create_error_response(

vllm/entrypoints/openai/serving_pooling.py

@@ -142,6 +142,11 @@ async def create_pooling(
         try:
             pooling_params = request.to_pooling_params()
 
+            try:
+                pooling_params.verify("encode", self.model_config)
+            except ValueError as e:
+                return self.create_error_response(str(e))
+
             for i, engine_prompt in enumerate(engine_prompts):
                 request_id_item = f"{request_id}-{i}"