315 lines
11 KiB
Python
315 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm import PoolingParams
|
|
from vllm.config import PoolerConfig
|
|
from vllm.entrypoints.pooling.embed.io_processor import (
|
|
JinaRankingTokenEmbedIOProcessor,
|
|
TokenEmbedIOProcessor,
|
|
)
|
|
from vllm.entrypoints.pooling.late_chunking import build_late_chunking_metadata
|
|
from vllm.entrypoints.pooling.offline import PoolingOfflineMixin
|
|
from vllm.entrypoints.pooling.typing import (
|
|
OfflineEncodeInputsContext,
|
|
OfflineOutputsContext,
|
|
)
|
|
from vllm.exceptions import VLLMValidationError
|
|
from vllm.outputs import (
|
|
LateChunk,
|
|
PoolingOutput,
|
|
PoolingRequestOutput,
|
|
RequestError,
|
|
)
|
|
from vllm.pooling_params import LateChunkingParams
|
|
from vllm.renderers import TokenizeParams
|
|
|
|
|
|
@pytest.fixture
|
|
def late_chunk_processor():
|
|
processor = object.__new__(TokenEmbedIOProcessor)
|
|
processor.model_config = SimpleNamespace(
|
|
is_encoder_decoder=False,
|
|
architecture="NomicBertModel",
|
|
model_impl="auto",
|
|
is_matryoshka=False,
|
|
hf_config=SimpleNamespace(),
|
|
pooler_config=PoolerConfig(seq_pooling_type="MEAN", tok_pooling_type="ALL"),
|
|
)
|
|
processor.vllm_config = SimpleNamespace(
|
|
model_config=processor.model_config,
|
|
cache_config=SimpleNamespace(enable_prefix_caching=False),
|
|
scheduler_config=SimpleNamespace(enable_chunked_prefill=False),
|
|
lora_config=None,
|
|
)
|
|
processor.renderer = Mock()
|
|
processor.renderer.default_cmpl_tok_params = TokenizeParams(max_total_tokens=32)
|
|
processor.renderer.render_cmpl.side_effect = lambda **_: [
|
|
{
|
|
"type": "token",
|
|
"prompt_token_ids": [101, 1, 2, 102],
|
|
"prompt_token_offsets": [(0, 0), (0, 1), (2, 3), (0, 0)],
|
|
}
|
|
]
|
|
return processor
|
|
|
|
|
|
def _late_chunk_context(prompts="a b", **kwargs):
|
|
return OfflineEncodeInputsContext(
|
|
pooling_task="token_embed",
|
|
tokenization_kwargs=kwargs or None,
|
|
lora_request=None,
|
|
priorities=None,
|
|
prompts=prompts,
|
|
pooling_params=PoolingParams(
|
|
late_chunking_params=LateChunkingParams(chunk_size=2)
|
|
),
|
|
)
|
|
|
|
|
|
def test_late_chunk_render_requests_offsets_once(late_chunk_processor):
|
|
processor = late_chunk_processor
|
|
ctx = _late_chunk_context()
|
|
factory, count = processor.get_request_factory_offline(ctx)
|
|
render_params = next(factory())
|
|
result = processor.render(render_params)
|
|
assert count == 1
|
|
processor.renderer.render_cmpl.assert_called_once()
|
|
assert processor.renderer.render_cmpl.call_args.kwargs[
|
|
"tok_params"
|
|
].return_token_offsets
|
|
assert not render_params["tok_params"].return_token_offsets
|
|
assert "prompt_token_offsets" not in result["prompts"]
|
|
assert "late_chunking" not in result
|
|
assert [c.char_range for c in ctx.late_chunking[0].metadata.chunks] == [
|
|
(0, 1),
|
|
(2, 3),
|
|
]
|
|
assert ctx.late_chunking[0].metadata.input_tokens == 4
|
|
|
|
|
|
@pytest.mark.parametrize("as_sequence", [False, True])
|
|
@pytest.mark.parametrize("chunk_size", [None, 2])
|
|
def test_jina_ranking_rejects_chunking_before_tokenization(
|
|
late_chunk_processor, as_sequence, chunk_size
|
|
):
|
|
processor = object.__new__(JinaRankingTokenEmbedIOProcessor)
|
|
processor.__dict__.update(late_chunk_processor.__dict__)
|
|
ctx = _late_chunk_context(prompts=["document", "query"])
|
|
params = PoolingParams(
|
|
late_chunking_params=LateChunkingParams(chunk_size) if chunk_size else None
|
|
)
|
|
ctx.pooling_params = [params] if as_sequence else params
|
|
factory, count = processor.get_request_factory_offline(ctx)
|
|
assert count == 1
|
|
request = next(factory())
|
|
if chunk_size is not None:
|
|
with pytest.raises(VLLMValidationError, match="JinaForRanking.*late chunking"):
|
|
processor.render(request)
|
|
late_chunk_processor.renderer.render_cmpl.assert_not_called()
|
|
else:
|
|
result = processor.render(request)
|
|
assert result["prompts"]["type"] == "token"
|
|
assert result["prompts"]["prompt_token_ids"] == [101, 1, 2, 102]
|
|
late_chunk_processor.renderer.render_cmpl.assert_called_once()
|
|
assert not request["tok_params"].return_token_offsets
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"prompts, kwargs, error",
|
|
[
|
|
pytest.param("", {}, "plain-text", id="empty"),
|
|
pytest.param([1, 2], {}, "plain-text", id="token-ids"),
|
|
pytest.param(
|
|
{"prompt": "abc", "multi_modal_data": {}},
|
|
{},
|
|
"plain-text",
|
|
id="multimodal",
|
|
),
|
|
]
|
|
+ [
|
|
pytest.param("abc", {key: value}, "does not support", id=key)
|
|
for key, value in (
|
|
("truncate_prompt_tokens", 2),
|
|
("truncation", True),
|
|
("pad_prompt_tokens", 8),
|
|
("padding", True),
|
|
("do_lower_case", True),
|
|
)
|
|
],
|
|
)
|
|
def test_late_chunking_rejects_unsupported_input_before_render(
|
|
late_chunk_processor, prompts, kwargs, error
|
|
):
|
|
processor = late_chunk_processor
|
|
with pytest.raises(VLLMValidationError, match=error):
|
|
factory, _ = processor.get_request_factory_offline(
|
|
_late_chunk_context(prompts, **kwargs)
|
|
)
|
|
processor.render(next(factory()))
|
|
processor.renderer.render_cmpl.assert_not_called()
|
|
|
|
|
|
def test_late_chunk_ranges_keep_unicode_overlaps_and_special_only_chunks():
|
|
text = "中 😀 e\u0301"
|
|
offsets = [(0, 0), (0, 1), (2, 3), (2, 3), (4, 6), (0, 0)]
|
|
metadata = build_late_chunking_metadata(text, len(offsets), offsets, 1)
|
|
assert [c.char_range for c in metadata.chunks] == [
|
|
None,
|
|
(0, 1),
|
|
(2, 3),
|
|
(2, 3),
|
|
(4, 6),
|
|
None,
|
|
]
|
|
assert [
|
|
text[slice(*c.char_range)] if c.char_range else None for c in metadata.chunks
|
|
] == [None, "中", "😀", "😀", "e\u0301", None]
|
|
chunks = build_late_chunking_metadata(text, len(offsets), offsets, 4).chunks
|
|
assert [(c.token_range, c.char_range) for c in chunks] == [
|
|
((0, 4), (0, 3)),
|
|
((4, 6), (4, 6)),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"offsets", [None, [(0, 1)], [(-1, 1), (0, 0)], [(0, 4), (0, 0)], [(2, 3), (0, 1)]]
|
|
)
|
|
def test_late_chunk_ranges_reject_missing_or_invalid_offsets(offsets):
|
|
with pytest.raises(VLLMValidationError, match="offsets"):
|
|
build_late_chunking_metadata("abc", 2, offsets, 2)
|
|
|
|
|
|
def _mock_chunk_llm(processor, outputs):
|
|
llm = mock.Mock(spec=PoolingOfflineMixin)
|
|
llm.encode = PoolingOfflineMixin.encode.__get__(llm)
|
|
llm._run_tiling_engine = PoolingOfflineMixin._run_tiling_engine.__get__(llm)
|
|
llm.pooling_io_processors = {"token_embed": processor}
|
|
llm._executor = SimpleNamespace(map=map)
|
|
llm.llm_engine = mock.Mock()
|
|
llm.llm_engine.vllm_config.scheduler_config.max_num_seqs = 2
|
|
llm.llm_engine.has_unfinished_requests.return_value = False
|
|
llm.llm_engine.step.side_effect = outputs
|
|
llm._render_and_add_requests = mock.Mock(
|
|
side_effect=lambda **kw: [
|
|
str(i) + "-internal" for i in range(len(kw["params"]))
|
|
]
|
|
)
|
|
return llm
|
|
|
|
|
|
def _chunk_output(request_id, rows=2, **kwargs):
|
|
return PoolingRequestOutput(
|
|
str(request_id),
|
|
PoolingOutput(torch.ones(rows, 4)),
|
|
[101, 1, 2, 102],
|
|
0,
|
|
True,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
def test_late_chunk_mapping_preserves_order_mixed_requests_and_errors(
|
|
late_chunk_processor,
|
|
):
|
|
error = RequestError("test_error", "original error")
|
|
failed = _chunk_output(3, rows=0, error=error)
|
|
llm = _mock_chunk_llm(late_chunk_processor, RuntimeError("step failed"))
|
|
params = [
|
|
PoolingParams(late_chunking_params=LateChunkingParams(2)),
|
|
PoolingParams(),
|
|
PoolingParams(late_chunking_params=LateChunkingParams(1)),
|
|
PoolingParams(late_chunking_params=LateChunkingParams(2)),
|
|
]
|
|
# A failed encode call must not leave metadata for the following batch.
|
|
with pytest.raises(RuntimeError, match="step failed"):
|
|
llm.encode(
|
|
["a b"] * 2,
|
|
pooling_task="token_embed",
|
|
pooling_params=params[:2],
|
|
use_tqdm=False,
|
|
)
|
|
llm.llm_engine.step.side_effect = [
|
|
[
|
|
_chunk_output(2, rows=4),
|
|
failed,
|
|
_chunk_output(1, rows=4),
|
|
_chunk_output(0),
|
|
]
|
|
]
|
|
outputs = llm.encode(
|
|
["a b"] * 4, pooling_task="token_embed", pooling_params=params, use_tqdm=False
|
|
)
|
|
assert outputs[0].late_chunking.chunks == [
|
|
LateChunk((0, 2), (0, 1)),
|
|
LateChunk((2, 4), (2, 3)),
|
|
]
|
|
assert outputs[1].late_chunking is None
|
|
assert len(outputs[2].late_chunking.chunks) == 4
|
|
assert outputs[3].late_chunking is None
|
|
assert outputs[3].error is error
|
|
|
|
|
|
def test_late_chunk_mapping_rejects_successful_output_with_wrong_row_count(
|
|
late_chunk_processor,
|
|
):
|
|
llm = _mock_chunk_llm(late_chunk_processor, [[_chunk_output(0, rows=0)]])
|
|
with pytest.raises(ValueError, match="does not match"):
|
|
llm.encode(
|
|
"a b",
|
|
pooling_task="token_embed",
|
|
pooling_params=PoolingParams(late_chunking_params=LateChunkingParams(2)),
|
|
use_tqdm=False,
|
|
)
|
|
# All engine requests have finished before post-processing validates the output.
|
|
llm.llm_engine.abort_request.assert_not_called()
|
|
|
|
|
|
def test_late_chunk_contexts_remain_isolated_when_rendering_interleaves(
|
|
late_chunk_processor,
|
|
):
|
|
processor = late_chunk_processor
|
|
first, second = _late_chunk_context("xy z"), _late_chunk_context()
|
|
shared = PoolingParams(late_chunking_params=LateChunkingParams(2))
|
|
first.pooling_params = second.pooling_params = shared
|
|
first_factory, _ = processor.get_request_factory_offline(first)
|
|
second_factory, _ = processor.get_request_factory_offline(second)
|
|
first_request, second_request = next(first_factory()), next(second_factory())
|
|
processor.render(second_request)
|
|
assert first.late_chunking[0].metadata is None
|
|
processor.renderer.render_cmpl.side_effect = lambda **_: [
|
|
{
|
|
"type": "token",
|
|
"prompt_token_ids": [101, 1, 2, 102],
|
|
"prompt_token_offsets": [(0, 0), (0, 2), (3, 4), (0, 0)],
|
|
}
|
|
]
|
|
processor.render(first_request)
|
|
first_output = processor.post_process_offline(
|
|
OfflineOutputsContext([_chunk_output(0)], late_chunking=first.late_chunking)
|
|
)[0]
|
|
second_output = processor.post_process_offline(
|
|
OfflineOutputsContext([_chunk_output(0)], late_chunking=second.late_chunking)
|
|
)[0]
|
|
assert first_output.late_chunking.chunks[0].char_range == (0, 2)
|
|
assert second_output.late_chunking.chunks[0].char_range == (0, 1)
|
|
|
|
|
|
def test_plain_token_embed_skips_chunk_metadata(late_chunk_processor):
|
|
ctx = _late_chunk_context()
|
|
ctx.pooling_params = PoolingParams()
|
|
factory, _ = late_chunk_processor.get_request_factory_offline(ctx)
|
|
request = next(factory())
|
|
result = late_chunk_processor.render(request)
|
|
assert not ctx.late_chunking
|
|
assert result["params"].late_chunking_params is None
|
|
assert not late_chunk_processor.renderer.render_cmpl.call_args.kwargs[
|
|
"tok_params"
|
|
].return_token_offsets
|