[Misc][Model] add io processor for query/document embeddings from ColBERT (jinaai/jina-colbert-v2) (#45210)
Signed-off-by: thomas <thomas.varghese@columbia.edu>
This commit is contained in:
xx-thomas
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
+194
@@ -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] <tokens> [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] <tokens> [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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
]
|
||||
},
|
||||
)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user