diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 148aea73c7..ee1658640b 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -608,6 +608,11 @@ steps: - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py - pip uninstall bge_m3_sparse_plugin -y # END: `bge_m3_sparse io_processor` test + # BEGIN: `colbert_query io_processor` test + - pip install -e ./plugins/colbert_query_plugin + - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py + - pip uninstall colbert_query_plugin -y + # END: `colbert_query io_processor` test # BEGIN: `stat_logger` plugins test - pip install -e ./plugins/vllm_add_dummy_stat_logger - pytest -v -s plugins_tests/test_stats_logger_plugins.py diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 21e3572fc7..310c2a8fd2 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -27,6 +27,10 @@ steps: - pip install -e ./plugins/bge_m3_sparse_plugin - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py - pip uninstall bge_m3_sparse_plugin -y + # test colbert_query io_processor plugin + - pip install -e ./plugins/colbert_query_plugin + - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py + - pip uninstall colbert_query_plugin -y # end io_processor plugins test # begin stat_logger plugins test - pip install -e ./plugins/vllm_add_dummy_stat_logger diff --git a/tests/plugins/colbert_query_plugin/colbert_query_processor/__init__.py b/tests/plugins/colbert_query_plugin/colbert_query_processor/__init__.py new file mode 100644 index 0000000000..021a6764d3 --- /dev/null +++ b/tests/plugins/colbert_query_plugin/colbert_query_processor/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +def register_colbert_query_embedding_processor(): + return "colbert_query_processor.query_embedding_processor.ColBERTQueryEmbeddingProcessor" # noqa: E501 diff --git a/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py b/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py new file mode 100644 index 0000000000..b56807ec15 --- /dev/null +++ b/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterator, Sequence +from typing import cast + +from vllm.config import VllmConfig +from vllm.entrypoints.openai.engine.protocol import UsageInfo +from vllm.inputs import PromptType, TokensPrompt +from vllm.outputs import PoolingRequestOutput +from vllm.plugins.io_processors.interface import IOProcessor +from vllm.pooling_params import PoolingParams +from vllm.renderers import BaseRenderer +from vllm.utils.collection_utils import is_list_of + +from .types import ( + QUERY_MAXLEN, + ColBERTEmbeddingCompletionRequestMixin, + ColBERTEmbeddingResponse, + ColBERTEmbeddingResponseData, +) + +QUERY_MARKER_TOKEN = "[QueryMarker]" +DOCUMENT_MARKER_TOKEN = "[DocumentMarker]" + + +class ColBERTQueryEmbeddingProcessor( + IOProcessor[ColBERTEmbeddingCompletionRequestMixin, ColBERTEmbeddingResponse] +): + """This IO processor only supports the ColBERT-style model jinaai/jina-colbert-v2. + It does not support all ColBERT-style variants (e.g. colbert-ir/colbertv2.0). + """ + + def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer): + super().__init__(vllm_config, renderer) + self.requests_cache: dict[str, ColBERTEmbeddingCompletionRequestMixin] = {} + self.renderer: BaseRenderer = renderer + # Context window (8192 for jinaai/jina-colbert-v2); caps document + # content length minus the 3 special-token slots. + self.max_model_len = vllm_config.model_config.max_model_len + self._query_marker_id: int | None = None + self._document_marker_id: int | None = None + + def __repr__(self) -> str: + return ( + f"ColBERTQueryEmbeddingProcessor(" + f"query_maxlen={QUERY_MAXLEN}, " + f"doc_maxlen={self.max_model_len}, " + f"query_marker_token={QUERY_MARKER_TOKEN!r}, " + f"document_marker_token={DOCUMENT_MARKER_TOKEN!r})" + ) + + def _resolve_marker_ids(self, tokenizer) -> tuple[int, int]: + if self._query_marker_id is not None and self._document_marker_id is not None: + return self._query_marker_id, self._document_marker_id + + unk_id = getattr(tokenizer, "unk_token_id", None) + marker_ids: list[int] = [] + for marker in (QUERY_MARKER_TOKEN, DOCUMENT_MARKER_TOKEN): + marker_id = tokenizer.convert_tokens_to_ids(marker) + if marker_id is None or marker_id == unk_id: + raise ValueError( + f"Marker token {marker!r} not found in the tokenizer " + "vocabulary. This plugin requires a ColBERT model whose " + "tokenizer defines both " + f"{QUERY_MARKER_TOKEN!r} and {DOCUMENT_MARKER_TOKEN!r} " + "(e.g. jinaai/jina-colbert-v2)." + ) + marker_ids.append(marker_id) + + self._query_marker_id, self._document_marker_id = marker_ids + return self._query_marker_id, self._document_marker_id + + def _iter_content_token_ids( + self, + tokenizer, + request_input: list[int] | list[list[int]] | str | list[str], + ) -> Iterator[list[int]]: + if isinstance(request_input, str): + yield tokenizer.encode(request_input, add_special_tokens=False) + return + + if not isinstance(request_input, list) or not request_input: + raise ValueError("input must be a non-empty string or list") + + if is_list_of(request_input, int): + yield list(cast(list[int], request_input)) + return + + for item in request_input: + if isinstance(item, str): + yield tokenizer.encode(item, add_special_tokens=False) + else: + yield list(cast(list[int], item)) + + def _build_query_prompt( + self, + tokenizer, + content_ids: list[int], + ) -> TokensPrompt: + """[CLS] [QueryMarker] [SEP] [MASK]... up to QUERY_MAXLEN.""" + query_marker_id, _ = self._resolve_marker_ids(tokenizer) + mask_token_id = tokenizer.mask_token_id + if mask_token_id is None: + raise ValueError( + "Tokenizer has no mask token; cannot perform query expansion." + ) + + # [CLS], marker and [SEP] take 3 slots. + content_ids = content_ids[: QUERY_MAXLEN - 3] + token_ids = [ + tokenizer.cls_token_id, + query_marker_id, + *content_ids, + tokenizer.sep_token_id, + ] + token_ids += [mask_token_id] * (QUERY_MAXLEN - len(token_ids)) + return TokensPrompt(prompt_token_ids=token_ids) + + def _build_document_prompt( + self, + tokenizer, + content_ids: list[int], + ) -> TokensPrompt: + """[CLS] [DocumentMarker] [SEP]""" + _, document_marker_id = self._resolve_marker_ids(tokenizer) + + content_ids = content_ids[: self.max_model_len - 3] + token_ids = [ + tokenizer.cls_token_id, + document_marker_id, + *content_ids, + tokenizer.sep_token_id, + ] + return TokensPrompt(prompt_token_ids=token_ids) + + def parse_data(self, data: object) -> ColBERTEmbeddingCompletionRequestMixin: + if isinstance(data, dict): + return ColBERTEmbeddingCompletionRequestMixin(**data) + raise TypeError("request data should be a dictionary") + + def pre_process( + self, + prompt: ColBERTEmbeddingCompletionRequestMixin, + request_id: str | None = None, + **kwargs, + ) -> PromptType | Sequence[PromptType]: + cache_key = request_id or "offline" + assert cache_key not in self.requests_cache, "request_id duplicated" + self.requests_cache[cache_key] = prompt + + tokenizer = self.renderer.get_tokenizer() + prompts: list[TokensPrompt] = [] + for content_ids in self._iter_content_token_ids(tokenizer, prompt.input): + if prompt.input_type == "query": + prompts.append(self._build_query_prompt(tokenizer, content_ids)) + else: + prompts.append(self._build_document_prompt(tokenizer, content_ids)) + return prompts + + def merge_pooling_params( + self, + params: PoolingParams | None = None, + ) -> PoolingParams: + if params is None: + params = PoolingParams() + params.task = "token_embed" + params.skip_reading_prefix_cache = True + return params + + def post_process( + self, + model_output: Sequence[PoolingRequestOutput], + request_id: str | None = None, + **kwargs, + ) -> ColBERTEmbeddingResponse: + raw_request = self.requests_cache.pop(request_id or "offline") + + num_prompt_tokens = 0 + response_data: list[ColBERTEmbeddingResponseData] = [] + for idx, output in enumerate(model_output): + num_prompt_tokens += len(output.prompt_token_ids) + response_data.append( + ColBERTEmbeddingResponseData( + index=idx, + input_type=raw_request.input_type, + embedding=output.outputs.data.tolist(), + ) + ) + + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + total_tokens=num_prompt_tokens, + ) + return ColBERTEmbeddingResponse(data=response_data, usage=usage) diff --git a/tests/plugins/colbert_query_plugin/colbert_query_processor/types.py b/tests/plugins/colbert_query_plugin/colbert_query_processor/types.py new file mode 100644 index 0000000000..9cf0700653 --- /dev/null +++ b/tests/plugins/colbert_query_plugin/colbert_query_processor/types.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Literal, get_args + +from pydantic import BaseModel, Field + +from vllm.entrypoints.openai.engine.protocol import UsageInfo +from vllm.entrypoints.pooling.base.protocol import CompletionRequestMixin + +InputType = Literal["query", "document"] +INPUT_TYPES: tuple[InputType, ...] = get_args(InputType) +QUERY_MAXLEN = 32 + + +class ColBERTEmbeddingCompletionRequestMixin(CompletionRequestMixin): + input_type: InputType = Field( + description="Whether to encode the input as a ColBERT 'query' " + f"(query marker + [mask] expansion to {QUERY_MAXLEN} tokens) or as a " + "'document' (document marker only). Required.", + ) + + +class ColBERTEmbeddingResponseData(BaseModel): + index: int + object: str = "embedding" + input_type: InputType + embedding: list[list[float]] + + +class ColBERTEmbeddingResponse(BaseModel): + data: list[ColBERTEmbeddingResponseData] + usage: UsageInfo diff --git a/tests/plugins/colbert_query_plugin/setup.py b/tests/plugins/colbert_query_plugin/setup.py new file mode 100644 index 0000000000..993c32cd02 --- /dev/null +++ b/tests/plugins/colbert_query_plugin/setup.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from setuptools import setup + +setup( + name="colbert-query-plugin", + version="0.1", + packages=["colbert_query_processor"], + entry_points={ + "vllm.io_processor_plugins": [ + "colbert_query_plugin = colbert_query_processor:register_colbert_query_embedding_processor", # noqa: E501 + ] + }, +) diff --git a/tests/plugins_tests/test_colbert_query_io_processor_plugins.py b/tests/plugins_tests/test_colbert_query_io_processor_plugins.py new file mode 100644 index 0000000000..930c493fdd --- /dev/null +++ b/tests/plugins_tests/test_colbert_query_io_processor_plugins.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from typing import TypedDict + +import pytest +import requests + +from tests.utils import RemoteOpenAIServer +from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse + + +# Test configuration for ColBERT query plugin +class ModelConfig(TypedDict): + model_name: str + plugin: str + query_input: str + document_input: str + hf_overrides: str + embedding_dim: int + query_maxlen: int + + +model_config: ModelConfig = { + "model_name": "jinaai/jina-colbert-v2", + "plugin": "colbert_query_plugin", + "query_input": "What is machine learning?", + "document_input": "Machine learning is a subset of artificial intelligence.", + "hf_overrides": json.dumps({"architectures": ["ColBERTJinaRobertaModel"]}), + "embedding_dim": 128, + "query_maxlen": 32, +} + + +def _get_attr_or_val(obj: object | dict, key: str): + if isinstance(obj, dict) and key in obj: + return obj[key] + return getattr(obj, key, None) + + +def _check_token_embeddings(entry, expected_input_type: str): + assert _get_attr_or_val(entry, "object") == "embedding" + assert _get_attr_or_val(entry, "input_type") == expected_input_type + + embedding = _get_attr_or_val(entry, "embedding") + assert isinstance(embedding, list) and len(embedding) > 0 + for token_embedding in embedding: + assert isinstance(token_embedding, list) + assert len(token_embedding) == model_config["embedding_dim"] + return embedding + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--runner", + "pooling", + "--enforce-eager", + "--max-num-seqs", + "32", + "--trust-remote-code", + "--hf_overrides", + model_config["hf_overrides"], + "--io-processor-plugin", + model_config["plugin"], + ] + + with RemoteOpenAIServer(model_config["model_name"], args) as remote_server: + yield remote_server + + +def _post_pooling(server: RemoteOpenAIServer, data: dict): + request_payload = { + "model": model_config["model_name"], + "task": "plugin", + "data": data, + } + ret = requests.post(server.url_for("pooling"), json=request_payload) + ret.raise_for_status() + response = ret.json() + parsed_response = IOProcessorResponse(**response).data + assert parsed_response + return parsed_response + + +def test_colbert_query_plugin_query_online(server: RemoteOpenAIServer): + """Queries are expanded to exactly query_maxlen token vectors.""" + parsed_response = _post_pooling( + server, {"input": model_config["query_input"], "input_type": "query"} + ) + + data = _get_attr_or_val(parsed_response, "data") + assert len(data) == 1 + + embedding = _check_token_embeddings(data[0], "query") + assert len(embedding) == model_config["query_maxlen"] + + usage = _get_attr_or_val(parsed_response, "usage") + assert _get_attr_or_val(usage, "prompt_tokens") == model_config["query_maxlen"] + + +def test_colbert_query_plugin_document_online(server: RemoteOpenAIServer): + """Documents return one vector per token, with no mask expansion.""" + parsed_response = _post_pooling( + server, {"input": model_config["document_input"], "input_type": "document"} + ) + + data = _get_attr_or_val(parsed_response, "data") + assert len(data) == 1 + + embedding = _check_token_embeddings(data[0], "document") + # No query expansion: number of vectors tracks the input length. + assert len(embedding) != model_config["query_maxlen"] + + usage = _get_attr_or_val(parsed_response, "usage") + assert _get_attr_or_val(usage, "prompt_tokens") == len(embedding) + + +def test_colbert_query_plugin_missing_input_type_online(server: RemoteOpenAIServer): + """input_type is required; omitting it is rejected.""" + request_payload = { + "model": model_config["model_name"], + "task": "plugin", + "data": {"input": model_config["document_input"]}, + } + ret = requests.post(server.url_for("pooling"), json=request_payload) + assert ret.status_code == 400 + + +def test_colbert_query_plugin_batch_online(server: RemoteOpenAIServer): + """A list input returns one entry per prompt.""" + queries = ["What is machine learning?", "What is deep learning?"] + parsed_response = _post_pooling(server, {"input": queries, "input_type": "query"}) + + data = _get_attr_or_val(parsed_response, "data") + assert len(data) == len(queries) + for i, entry in enumerate(data): + assert _get_attr_or_val(entry, "index") == i + embedding = _check_token_embeddings(entry, "query") + assert len(embedding) == model_config["query_maxlen"] + + +@pytest.mark.parametrize("input_type", ["query", "document"]) +def test_colbert_query_plugin_offline(vllm_runner, input_type: str): + """Test the ColBERT query plugin in offline mode.""" + input_text = ( + model_config["query_input"] + if input_type == "query" + else model_config["document_input"] + ) + prompt = { + "data": { + "input": input_text, + "input_type": input_type, + } + } + + with vllm_runner( + model_config["model_name"], + runner="pooling", + enforce_eager=True, + max_num_seqs=32, + trust_remote_code=True, + io_processor_plugin=model_config["plugin"], + hf_overrides=json.loads(model_config["hf_overrides"]), + default_torch_num_threads=1, + ) as llm_runner: + llm = llm_runner.get_llm() + pooler_output = llm.encode(prompt, pooling_task="plugin") + + response = pooler_output[0].outputs + assert len(response.data) == 1 + + embedding = _check_token_embeddings(response.data[0], input_type) + if input_type == "query": + assert len(embedding) == model_config["query_maxlen"] + else: + assert len(embedding) != model_config["query_maxlen"] + + assert response.usage.prompt_tokens == len(embedding) + assert response.usage.total_tokens == response.usage.prompt_tokens + + +def test_colbert_query_plugin_offline_multiple_inputs(vllm_runner): + """Test the ColBERT query plugin with multiple inputs in offline mode.""" + queries = [ + "What is machine learning?", + "What is deep learning?", + "Why?", + ] + prompts = { + "data": { + "input": queries, + "input_type": "query", + } + } + + with vllm_runner( + model_config["model_name"], + runner="pooling", + enforce_eager=True, + max_num_seqs=32, + trust_remote_code=True, + io_processor_plugin=model_config["plugin"], + hf_overrides=json.loads(model_config["hf_overrides"]), + default_torch_num_threads=1, + ) as llm_runner: + llm = llm_runner.get_llm() + pooler_output = llm.encode(prompts, pooling_task="plugin") + + response = pooler_output[0].outputs + assert len(response.data) == len(queries) + + for i, entry in enumerate(response.data): + assert entry.index == i + embedding = _check_token_embeddings(entry, "query") + assert len(embedding) == model_config["query_maxlen"] + + expected_tokens = model_config["query_maxlen"] * len(queries) + assert response.usage.prompt_tokens == expected_tokens + assert response.usage.total_tokens == response.usage.prompt_tokens