Add vLLM v0.18.1 source tree with KV transfer abort fix
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:
vllm/v1/core/sched/scheduler.py:
Replace fatal assert with graceful skip when KV transfer callback
arrives for an already-aborted request during PD disaggregated serving.
Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
0
third_party/vllm/tests/entrypoints/pooling/__init__.py
vendored
Normal file
0
third_party/vllm/tests/entrypoints/pooling/__init__.py
vendored
Normal file
0
third_party/vllm/tests/entrypoints/pooling/basic/__init__.py
vendored
Normal file
0
third_party/vllm/tests/entrypoints/pooling/basic/__init__.py
vendored
Normal file
92
third_party/vllm/tests/entrypoints/pooling/basic/test_encode.py
vendored
Normal file
92
third_party/vllm/tests/entrypoints/pooling/basic/test_encode.py
vendored
Normal file
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM, PoolingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "intfloat/multilingual-e5-small"
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
TOKEN_IDS = [
|
||||
# Using ID={0, 1, 2, 3} results in NaN values,
|
||||
# so we add this offset of 1000
|
||||
[1000],
|
||||
[1000, 1001],
|
||||
[1000, 1002, 1001],
|
||||
[1000, 1003, 1001, 1002],
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
# ROCm: Use FLEX_ATTENTION backend as it's the only attention backend
|
||||
# that supports encoder-only models on ROCm.
|
||||
attention_config = None
|
||||
if current_platform.is_rocm():
|
||||
attention_config = {"backend": "FLEX_ATTENTION"}
|
||||
|
||||
# pytest caches the fixture so we use weakref.proxy to
|
||||
# enable garbage collection
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
max_num_batched_tokens=32768,
|
||||
tensor_parallel_size=1,
|
||||
gpu_memory_utilization=0.75,
|
||||
enforce_eager=True,
|
||||
seed=0,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_multiple_pooling_params(llm: LLM):
|
||||
pooling_params = [
|
||||
PoolingParams(),
|
||||
PoolingParams(),
|
||||
PoolingParams(),
|
||||
PoolingParams(),
|
||||
]
|
||||
|
||||
# Multiple PoolingParams should be matched with each prompt
|
||||
outputs = llm.encode(PROMPTS, pooling_params=pooling_params, pooling_task="embed")
|
||||
assert len(PROMPTS) == len(outputs)
|
||||
|
||||
# Exception raised, if the size of params does not match the size of prompts
|
||||
with pytest.raises(ValueError):
|
||||
outputs = llm.encode(
|
||||
PROMPTS, pooling_params=pooling_params[:3], pooling_task="embed"
|
||||
)
|
||||
|
||||
# Single PoolingParams should be applied to every prompt
|
||||
single_pooling_params = PoolingParams()
|
||||
outputs = llm.encode(
|
||||
PROMPTS, pooling_params=single_pooling_params, pooling_task="embed"
|
||||
)
|
||||
assert len(PROMPTS) == len(outputs)
|
||||
|
||||
# pooling_params is None, default params should be applied
|
||||
outputs = llm.encode(PROMPTS, pooling_params=None, pooling_task="embed")
|
||||
assert len(PROMPTS) == len(outputs)
|
||||
|
||||
|
||||
def test_right_side_truncation(llm: LLM):
|
||||
# Embeddings models should truncate the end of the prompt
|
||||
tokenizer = llm.get_tokenizer()
|
||||
assert tokenizer.truncation_side == "right"
|
||||
104
third_party/vllm/tests/entrypoints/pooling/basic/test_truncation.py
vendored
Normal file
104
third_party/vllm/tests/entrypoints/pooling/basic/test_truncation.py
vendored
Normal file
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import Any
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
|
||||
max_model_len = 128
|
||||
|
||||
input = """Immerse yourself in the enchanting chronicle of calculus, a
|
||||
mathematical domain that has radically transformed our comprehension of
|
||||
change and motion. Despite its roots in ancient civilizations, the
|
||||
formal birth of calculus predominantly occurred in the 17th century,
|
||||
primarily under the influential guidance of Sir Isaac Newton and Gottfried
|
||||
Wilhelm Leibniz. The earliest traces of calculus concepts are found in
|
||||
ancient Greek mathematics,most notably in the works of Eudoxus and
|
||||
Archimedes, around 300 BCE. They utilized the 'method of exhaustion'—a
|
||||
technique for computing areas and volumes through the use of finite sums.
|
||||
This methodology laid crucial foundational work for integral calculus.
|
||||
In the 17th century, both Newton and Leibniz independently pioneered
|
||||
calculus, each contributing unique perspectives that would shape this new
|
||||
field."""
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
str(max_model_len),
|
||||
]
|
||||
|
||||
# ROCm: Use Flex Attention to support encoder-only self-attention.
|
||||
if current_platform.is_rocm():
|
||||
args.extend(["--attention-backend", "FLEX_ATTENTION"])
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_smaller_truncation_size(client: openai.AsyncOpenAI):
|
||||
truncation_size = 10
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": MODEL_NAME,
|
||||
"input": input,
|
||||
"truncate_prompt_tokens": truncation_size,
|
||||
}
|
||||
|
||||
response = await client.post(path="embeddings", cast_to=object, body={**kwargs})
|
||||
|
||||
assert response["usage"]["prompt_tokens"] == truncation_size
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bigger_truncation_size(client: openai.AsyncOpenAI):
|
||||
truncation_size = max_model_len + 1
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": MODEL_NAME,
|
||||
"input": input,
|
||||
"truncate_prompt_tokens": truncation_size,
|
||||
}
|
||||
|
||||
with pytest.raises(openai.BadRequestError) as err:
|
||||
await client.post(path="embeddings", cast_to=object, body={**kwargs})
|
||||
|
||||
assert err.value.status_code == 400
|
||||
error_details = err.value.response.json()["error"]
|
||||
assert error_details["type"] == "BadRequestError"
|
||||
expected_message = (
|
||||
"truncate_prompt_tokens value is "
|
||||
"greater than max_model_len."
|
||||
" Please, select a smaller truncation size."
|
||||
)
|
||||
assert error_details["message"] == expected_message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_truncation_size(client: openai.AsyncOpenAI):
|
||||
truncation_size = -1
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": MODEL_NAME,
|
||||
"input": input,
|
||||
"truncate_prompt_tokens": truncation_size,
|
||||
}
|
||||
|
||||
response = await client.post(path="embeddings", cast_to=object, body={**kwargs})
|
||||
|
||||
assert response["usage"]["prompt_tokens"] == max_model_len
|
||||
0
third_party/vllm/tests/entrypoints/pooling/classify/__init__.py
vendored
Normal file
0
third_party/vllm/tests/entrypoints/pooling/classify/__init__.py
vendored
Normal file
114
third_party/vllm/tests/entrypoints/pooling/classify/test_offline.py
vendored
Normal file
114
third_party/vllm/tests/entrypoints/pooling/classify/test_offline.py
vendored
Normal file
@@ -0,0 +1,114 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.models.utils import softmax
|
||||
from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.tasks import PoolingTask
|
||||
|
||||
MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
|
||||
|
||||
prompt = "The chef prepared a delicious meal."
|
||||
prompt_token_ids = [785, 29706, 10030, 264, 17923, 15145, 13]
|
||||
num_labels = 2
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
# pytest caches the fixture so we use weakref.proxy to
|
||||
# enable garbage collection
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
max_num_batched_tokens=32768,
|
||||
tensor_parallel_size=1,
|
||||
gpu_memory_utilization=0.75,
|
||||
enforce_eager=True,
|
||||
seed=0,
|
||||
)
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_str_prompts(llm: LLM):
|
||||
outputs = llm.classify(prompt, use_tqdm=False)
|
||||
assert len(outputs) == 1
|
||||
assert isinstance(outputs[0], ClassificationRequestOutput)
|
||||
assert outputs[0].prompt_token_ids == prompt_token_ids
|
||||
assert len(outputs[0].outputs.probs) == num_labels
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_ids_prompts(llm: LLM):
|
||||
outputs = llm.classify([prompt_token_ids], use_tqdm=False)
|
||||
assert len(outputs) == 1
|
||||
assert isinstance(outputs[0], ClassificationRequestOutput)
|
||||
assert outputs[0].prompt_token_ids == prompt_token_ids
|
||||
assert len(outputs[0].outputs.probs) == num_labels
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_list_prompts(llm: LLM):
|
||||
outputs = llm.classify([prompt, prompt_token_ids], use_tqdm=False)
|
||||
assert len(outputs) == 2
|
||||
for i in range(len(outputs)):
|
||||
assert isinstance(outputs[i], ClassificationRequestOutput)
|
||||
assert outputs[i].prompt_token_ids == prompt_token_ids
|
||||
assert len(outputs[i].outputs.probs) == num_labels
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_classify(llm: LLM):
|
||||
outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
|
||||
assert len(outputs) == 1
|
||||
assert isinstance(outputs[0], PoolingRequestOutput)
|
||||
assert outputs[0].prompt_token_ids == prompt_token_ids
|
||||
assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(use_activation):
|
||||
outputs = llm.classify(
|
||||
prompt,
|
||||
pooling_params=PoolingParams(use_activation=use_activation),
|
||||
use_tqdm=False,
|
||||
)
|
||||
return torch.tensor([x.outputs.probs for x in outputs])
|
||||
|
||||
default = get_outputs(use_activation=None)
|
||||
w_activation = get_outputs(use_activation=True)
|
||||
wo_activation = get_outputs(use_activation=False)
|
||||
|
||||
assert torch.allclose(default, w_activation, atol=1e-2), (
|
||||
"Default should use activation."
|
||||
)
|
||||
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
|
||||
"wo_activation should not use activation."
|
||||
)
|
||||
assert torch.allclose(softmax(wo_activation), w_activation, atol=1e-2), (
|
||||
"w_activation should be close to activation(wo_activation)."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_score_api(llm: LLM):
|
||||
err_msg = "Score API is only enabled for num_labels == 1."
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.score("ping", "pong", use_tqdm=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
|
||||
err_msg = f"Unsupported task: '{task}' Supported tasks.+"
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
472
third_party/vllm/tests/entrypoints/pooling/classify/test_online.py
vendored
Normal file
472
third_party/vllm/tests/entrypoints/pooling/classify/test_online.py
vendored
Normal file
@@ -0,0 +1,472 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse
|
||||
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
|
||||
|
||||
MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
|
||||
DTYPE = "float32" # Use float32 to avoid NaN issue
|
||||
input_text = "This product was excellent and exceeded my expectations"
|
||||
input_tokens = [1986, 1985, 572, 9073, 323, 33808, 847, 16665]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_basic(server: RemoteOpenAIServer, model_name: str):
|
||||
# test /v1/models
|
||||
response = requests.get(server.url_for("/v1/models"))
|
||||
served_model = response.json()["data"][0]["id"]
|
||||
assert served_model == MODEL_NAME
|
||||
|
||||
# test /tokenize
|
||||
response = requests.post(
|
||||
server.url_for("/tokenize"),
|
||||
json={"model": model_name, "prompt": input_text},
|
||||
)
|
||||
assert response.json()["tokens"] == input_tokens
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_completion_request(server: RemoteOpenAIServer, model_name: str):
|
||||
# test input: str
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": input_text},
|
||||
)
|
||||
|
||||
classification_response.raise_for_status()
|
||||
output = ClassificationResponse.model_validate(classification_response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data) == 1
|
||||
assert hasattr(output.data[0], "label")
|
||||
assert hasattr(output.data[0], "probs")
|
||||
|
||||
# test input: list[int]
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": input_tokens},
|
||||
)
|
||||
|
||||
classification_response.raise_for_status()
|
||||
output = ClassificationResponse.model_validate(classification_response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data) == 1
|
||||
assert hasattr(output.data[0], "label")
|
||||
assert hasattr(output.data[0], "probs")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str):
|
||||
N = 10
|
||||
|
||||
# test input: list[str]
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": [input_text] * N},
|
||||
)
|
||||
output = ClassificationResponse.model_validate(classification_response.json())
|
||||
|
||||
assert len(output.data) == N
|
||||
for i, item in enumerate(output.data):
|
||||
assert item.index == i
|
||||
assert hasattr(item, "label")
|
||||
assert hasattr(item, "probs")
|
||||
assert len(item.probs) == item.num_classes
|
||||
assert item.label in ["Default", "Spoiled"]
|
||||
|
||||
# test input: list[list[int]]
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": [input_tokens] * N},
|
||||
)
|
||||
output = ClassificationResponse.model_validate(classification_response.json())
|
||||
|
||||
assert len(output.data) == N
|
||||
for i, item in enumerate(output.data):
|
||||
assert item.index == i
|
||||
assert hasattr(item, "label")
|
||||
assert hasattr(item, "probs")
|
||||
assert len(item.probs) == item.num_classes
|
||||
assert item.label in ["Default", "Spoiled"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_empty_input_error(server: RemoteOpenAIServer, model_name: str):
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": ""},
|
||||
)
|
||||
|
||||
error = classification_response.json()
|
||||
assert classification_response.status_code == 400
|
||||
assert "error" in error
|
||||
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": []},
|
||||
)
|
||||
|
||||
error = classification_response.json()
|
||||
assert classification_response.status_code == 400
|
||||
assert "error" in error
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_truncate_prompt_tokens(server: RemoteOpenAIServer, model_name: str):
|
||||
long_text = "hello " * 600
|
||||
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": long_text, "truncate_prompt_tokens": 5},
|
||||
)
|
||||
|
||||
classification_response.raise_for_status()
|
||||
output = ClassificationResponse.model_validate(classification_response.json())
|
||||
|
||||
assert len(output.data) == 1
|
||||
assert output.data[0].index == 0
|
||||
assert hasattr(output.data[0], "probs")
|
||||
assert output.usage.prompt_tokens == 5
|
||||
assert output.usage.total_tokens == 5
|
||||
|
||||
# invalid_truncate_prompt_tokens
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": "test", "truncate_prompt_tokens": 513},
|
||||
)
|
||||
|
||||
error = classification_response.json()
|
||||
assert classification_response.status_code == 400
|
||||
assert "truncate_prompt_tokens" in error["error"]["message"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_add_special_tokens(server: RemoteOpenAIServer, model_name: str):
|
||||
# The add_special_tokens parameter doesn't seem to be working with this model.
|
||||
# working with papluca/xlm-roberta-base-language-detection
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": input_text, "add_special_tokens": False},
|
||||
)
|
||||
response.raise_for_status()
|
||||
ClassificationResponse.model_validate(response.json())
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "input": input_text, "add_special_tokens": True},
|
||||
)
|
||||
response.raise_for_status()
|
||||
ClassificationResponse.model_validate(response.json())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_chat_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "The cat sat on the mat.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "A feline was resting on a rug.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Stars twinkle brightly in the night sky.",
|
||||
},
|
||||
]
|
||||
|
||||
# test chat request basic usage
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = ClassificationResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data) == 1
|
||||
assert hasattr(output.data[0], "label")
|
||||
assert hasattr(output.data[0], "probs")
|
||||
assert output.usage.prompt_tokens == 51
|
||||
|
||||
# test add_generation_prompt
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "messages": messages, "add_generation_prompt": True},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = ClassificationResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data) == 1
|
||||
assert hasattr(output.data[0], "label")
|
||||
assert hasattr(output.data[0], "probs")
|
||||
assert output.usage.prompt_tokens == 54
|
||||
|
||||
# test continue_final_message
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"continue_final_message": True,
|
||||
},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = ClassificationResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data) == 1
|
||||
assert hasattr(output.data[0], "label")
|
||||
assert hasattr(output.data[0], "probs")
|
||||
assert output.usage.prompt_tokens == 49
|
||||
|
||||
# test add_special_tokens
|
||||
# The add_special_tokens parameter doesn't seem to be working with this model.
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "messages": messages, "add_special_tokens": True},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = ClassificationResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data) == 1
|
||||
assert hasattr(output.data[0], "label")
|
||||
assert hasattr(output.data[0], "probs")
|
||||
assert output.usage.prompt_tokens == 51
|
||||
|
||||
# test continue_final_message with add_generation_prompt
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"continue_final_message": True,
|
||||
"add_generation_prompt": True,
|
||||
},
|
||||
)
|
||||
assert (
|
||||
"Cannot set both `continue_final_message` and `add_generation_prompt` to True."
|
||||
in response.json()["error"]["message"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invocations_completion_request(server: RemoteOpenAIServer):
|
||||
request_args = {
|
||||
"model": MODEL_NAME,
|
||||
"input": input_text,
|
||||
}
|
||||
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"), json=request_args
|
||||
)
|
||||
classification_response.raise_for_status()
|
||||
|
||||
invocation_response = requests.post(
|
||||
server.url_for("invocations"), json=request_args
|
||||
)
|
||||
invocation_response.raise_for_status()
|
||||
|
||||
classification_output = classification_response.json()
|
||||
invocation_output = invocation_response.json()
|
||||
|
||||
assert classification_output.keys() == invocation_output.keys()
|
||||
for classification_data, invocation_data in zip(
|
||||
classification_output["data"], invocation_output["data"]
|
||||
):
|
||||
assert classification_data.keys() == invocation_data.keys()
|
||||
assert classification_data["probs"] == pytest.approx(
|
||||
invocation_data["probs"], rel=0.01
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invocations_chat_request(server: RemoteOpenAIServer):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "The cat sat on the mat.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "A feline was resting on a rug.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Stars twinkle brightly in the night sky.",
|
||||
},
|
||||
]
|
||||
|
||||
request_args = {"model": MODEL_NAME, "messages": messages}
|
||||
|
||||
classification_response = requests.post(
|
||||
server.url_for("classify"), json=request_args
|
||||
)
|
||||
classification_response.raise_for_status()
|
||||
|
||||
invocation_response = requests.post(
|
||||
server.url_for("invocations"), json=request_args
|
||||
)
|
||||
invocation_response.raise_for_status()
|
||||
|
||||
classification_output = classification_response.json()
|
||||
invocation_output = invocation_response.json()
|
||||
|
||||
assert classification_output.keys() == invocation_output.keys()
|
||||
for classification_data, invocation_data in zip(
|
||||
classification_output["data"], invocation_output["data"]
|
||||
):
|
||||
assert classification_data.keys() == invocation_data.keys()
|
||||
assert classification_data["probs"] == pytest.approx(
|
||||
invocation_data["probs"], rel=0.01
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
|
||||
async def get_outputs(use_activation):
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"use_activation": use_activation,
|
||||
},
|
||||
)
|
||||
outputs = response.json()
|
||||
return torch.tensor([x["probs"] for x in outputs["data"]])
|
||||
|
||||
default = await get_outputs(use_activation=None)
|
||||
w_activation = await get_outputs(use_activation=True)
|
||||
wo_activation = await get_outputs(use_activation=False)
|
||||
|
||||
assert torch.allclose(default, w_activation, atol=1e-2), (
|
||||
"Default should use activation."
|
||||
)
|
||||
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
|
||||
"wo_activation should not use activation."
|
||||
)
|
||||
assert torch.allclose(F.softmax(wo_activation, dim=-1), w_activation, atol=1e-2), (
|
||||
"w_activation should be close to activation(wo_activation)."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_score(server: RemoteOpenAIServer, model_name: str):
|
||||
# score api is only enabled for num_labels == 1.
|
||||
response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"queries": "ping",
|
||||
"documents": "pong",
|
||||
},
|
||||
)
|
||||
assert response.json()["detail"] == "Not Found"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_rerank(server: RemoteOpenAIServer, model_name: str):
|
||||
# rerank api is only enabled for num_labels == 1.
|
||||
response = requests.post(
|
||||
server.url_for("rerank"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"query": "ping",
|
||||
"documents": ["pong"],
|
||||
},
|
||||
)
|
||||
assert response.json()["detail"] == "Not Found"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": "classify",
|
||||
},
|
||||
)
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "token_classify"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == 8
|
||||
assert len(poolings.data[0].data[0]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")
|
||||
146
third_party/vllm/tests/entrypoints/pooling/classify/test_online_vision.py
vendored
Normal file
146
third_party/vllm/tests/entrypoints/pooling/classify/test_online_vision.py
vendored
Normal file
@@ -0,0 +1,146 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
MODEL_NAME = "muziyongshixin/Qwen2.5-VL-7B-for-VideoCls"
|
||||
MAXIMUM_VIDEOS = 1
|
||||
|
||||
HF_OVERRIDES = {"architectures": ["Qwen2_5_VLForSequenceClassification"]}
|
||||
input_text = "This product was excellent and exceeded my expectations"
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
image_base64 = {"url": encode_image_url(fetch_image(image_url))}
|
||||
video_url = "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--max-model-len",
|
||||
"5000",
|
||||
"--enforce-eager",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"video": MAXIMUM_VIDEOS}),
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, args, override_hf_configs=HF_OVERRIDES
|
||||
) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_chat_text_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please classify this text request.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": input_text,
|
||||
},
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output = ClassificationResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == model_name
|
||||
assert len(output.data) == 1
|
||||
assert len(output.data[0].probs) == 2
|
||||
assert output.usage.prompt_tokens == 35
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_chat_image_url_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Please classify this image."},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output = ClassificationResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == model_name
|
||||
assert len(output.data) == 1
|
||||
assert len(output.data[0].probs) == 2
|
||||
assert output.usage.prompt_tokens == 47
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_chat_image_base64_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Please classify this image."},
|
||||
{"type": "image_url", "image_url": image_base64},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output = ClassificationResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == model_name
|
||||
assert len(output.data) == 1
|
||||
assert len(output.data[0].probs) == 2
|
||||
assert output.usage.prompt_tokens == 47
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Please classify this video."},
|
||||
{"type": "video_url", "video_url": {"url": video_url}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("classify"),
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output = ClassificationResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert output.model == model_name
|
||||
assert len(output.data) == 1
|
||||
assert len(output.data[0].probs) == 2
|
||||
assert output.usage.prompt_tokens == 4807
|
||||
0
third_party/vllm/tests/entrypoints/pooling/embed/__init__.py
vendored
Normal file
0
third_party/vllm/tests/entrypoints/pooling/embed/__init__.py
vendored
Normal file
28
third_party/vllm/tests/entrypoints/pooling/embed/conftest.py
vendored
Normal file
28
third_party/vllm/tests/entrypoints/pooling/embed/conftest.py
vendored
Normal file
@@ -0,0 +1,28 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Pytest configuration for vLLM pooling embed tests."""
|
||||
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(config, items):
|
||||
"""Configure ROCm-specific settings based on collected tests."""
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
# Disable Flash/MemEfficient SDP on ROCm to avoid HF Transformers
|
||||
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
|
||||
# TODO: Remove once ROCm SDP accuracy issues are resolved on HuggingFace
|
||||
torch.backends.cuda.enable_flash_sdp(False)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(False)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
warnings.warn(
|
||||
"ROCm: Disabled flash_sdp and mem_efficient_sdp, enabled math_sdp "
|
||||
"to avoid HuggingFace Transformers accuracy issues",
|
||||
UserWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
310
third_party/vllm/tests/entrypoints/pooling/embed/test_cohere_online.py
vendored
Normal file
310
third_party/vllm/tests/entrypoints/pooling/embed/test_cohere_online.py
vendored
Normal file
@@ -0,0 +1,310 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the Cohere /v2/embed API with generic (non-Cohere) models.
|
||||
|
||||
Validates that the Cohere v2 embed endpoint works correctly with standard
|
||||
embedding models, covering text embedding, embedding type conversions,
|
||||
response structure, batching, normalisation, and semantic similarity.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
MODELS: list[tuple[str, list[str]]] = [
|
||||
("intfloat/multilingual-e5-small", []),
|
||||
(
|
||||
"Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
[
|
||||
"--trust_remote_code",
|
||||
"--hf_overrides",
|
||||
'{"matryoshka_dimensions":[256]}',
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=MODELS, ids=lambda m: m[0])
|
||||
def model_config(request):
|
||||
return request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_name(model_config):
|
||||
return model_config[0]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(model_config):
|
||||
name, extra_args = model_config
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--gpu-memory-utilization",
|
||||
"0.02",
|
||||
] + extra_args
|
||||
with RemoteOpenAIServer(name, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
model_name: str,
|
||||
texts: list[str] | None = None,
|
||||
images: list[str] | None = None,
|
||||
input_type: str | None = None,
|
||||
embedding_types: list[str] | None = None,
|
||||
) -> dict:
|
||||
body: dict = {"model": model_name}
|
||||
if input_type is not None:
|
||||
body["input_type"] = input_type
|
||||
if texts is not None:
|
||||
body["texts"] = texts
|
||||
if images is not None:
|
||||
body["images"] = images
|
||||
if embedding_types is not None:
|
||||
body["embedding_types"] = embedding_types
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _openai_embed(
|
||||
server: RemoteOpenAIServer, model_name: str, texts: list[str]
|
||||
) -> dict:
|
||||
body = {"model": model_name, "input": texts, "encoding_format": "float"}
|
||||
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _cosine_sim(a: list[float], b: list[float]) -> float:
|
||||
va, vb = np.array(a), np.array(b)
|
||||
return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
|
||||
|
||||
|
||||
# -----------------------------------------------------------
|
||||
# Text embedding tests
|
||||
# -----------------------------------------------------------
|
||||
|
||||
|
||||
def test_basic_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server, model_name, texts=["hello world"], embedding_types=["float"]
|
||||
)
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
|
||||
|
||||
def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str):
|
||||
"""An input_type not defined in the model's prompt config should be
|
||||
rejected with a 400 error."""
|
||||
body = {
|
||||
"model": model_name,
|
||||
"input_type": "nonexistent_type",
|
||||
"texts": ["hello world"],
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 400
|
||||
assert "Unsupported input_type" in resp.json()["error"]["message"]
|
||||
|
||||
|
||||
def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str):
|
||||
"""Omitting input_type should always work (no prompt prefix applied)."""
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": ["hello world"],
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str):
|
||||
"""v1 (OpenAI) and v2 (Cohere) endpoints should produce the same
|
||||
float embeddings for a generic model."""
|
||||
texts = ["hello world"]
|
||||
v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
|
||||
v1 = _openai_embed(server, model_name, texts)
|
||||
cos = _cosine_sim(v2["embeddings"]["float"][0], v1["data"][0]["embedding"])
|
||||
assert cos > 0.9999, f"v1/v2 parity failed, cosine={cos}"
|
||||
|
||||
|
||||
def test_embedding_types(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["test"],
|
||||
embedding_types=["float", "binary", "ubinary"],
|
||||
)
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
assert len(r["embeddings"]["binary"][0]) == dim // 8
|
||||
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
|
||||
|
||||
|
||||
def test_response_structure(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(server, model_name, texts=["test"], embedding_types=["float"])
|
||||
assert "id" in r
|
||||
assert "embeddings" in r
|
||||
assert "texts" in r
|
||||
assert r["texts"] == ["test"]
|
||||
assert "meta" in r
|
||||
assert r["meta"]["api_version"]["version"] == "2"
|
||||
assert "billed_units" in r["meta"]
|
||||
assert r["meta"]["billed_units"]["input_tokens"] > 0
|
||||
assert r["meta"]["billed_units"]["image_tokens"] == 0
|
||||
|
||||
|
||||
def test_batch(server: RemoteOpenAIServer, model_name: str):
|
||||
texts = ["apple", "banana", "cherry"]
|
||||
r = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
|
||||
assert len(r["embeddings"]["float"]) == 3
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
for emb in r["embeddings"]["float"]:
|
||||
assert len(emb) == dim
|
||||
|
||||
|
||||
def test_l2_normalized(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server, model_name, texts=["hello world"], embedding_types=["float"]
|
||||
)
|
||||
emb = np.array(r["embeddings"]["float"][0])
|
||||
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
|
||||
|
||||
|
||||
def test_semantic_similarity(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["machine learning", "deep learning", "chocolate cake recipe"],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
embs = r["embeddings"]["float"]
|
||||
cos_related = _cosine_sim(embs[0], embs[1])
|
||||
cos_unrelated = _cosine_sim(embs[0], embs[2])
|
||||
assert cos_related > cos_unrelated
|
||||
|
||||
|
||||
def test_missing_input_returns_error(server: RemoteOpenAIServer, model_name: str):
|
||||
body = {"model": model_name}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_base64_embedding_type(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["test encoding"],
|
||||
embedding_types=["float", "base64"],
|
||||
)
|
||||
float_emb = r["embeddings"]["float"][0]
|
||||
b64_str = r["embeddings"]["base64"][0]
|
||||
decoded = struct.unpack(f"<{len(float_emb)}f", base64.b64decode(b64_str))
|
||||
np.testing.assert_allclose(float_emb, decoded, rtol=1e-5)
|
||||
|
||||
|
||||
# -----------------------------------------------------------
|
||||
# Truncation tests
|
||||
# -----------------------------------------------------------
|
||||
|
||||
|
||||
def _cohere_embed_raw(
|
||||
server: RemoteOpenAIServer,
|
||||
body: dict,
|
||||
) -> requests.Response:
|
||||
return requests.post(server.url_for("/v2/embed"), json=body)
|
||||
|
||||
|
||||
def test_truncate_end_succeeds(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=END should silently truncate long input."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "END",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_truncate_start_succeeds(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=START should silently truncate long input from the start."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "START",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_truncate_none_rejects_long_input(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=NONE should error when input exceeds model context."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "NONE",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_truncate_start_vs_end_differ(server: RemoteOpenAIServer, model_name: str):
|
||||
"""START and END truncation should produce different embeddings
|
||||
when the input is long enough to actually be truncated.
|
||||
|
||||
We construct input with distinct tokens at the start vs end
|
||||
so that keeping different halves produces different embeddings.
|
||||
"""
|
||||
start_words = " ".join([f"alpha{i}" for i in range(300)])
|
||||
end_words = " ".join([f"omega{i}" for i in range(300)])
|
||||
long_text = start_words + " " + end_words
|
||||
|
||||
body_end = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "END",
|
||||
}
|
||||
body_start = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "START",
|
||||
}
|
||||
r_end = _cohere_embed_raw(server, body_end).json()
|
||||
r_start = _cohere_embed_raw(server, body_start).json()
|
||||
|
||||
emb_end = r_end["embeddings"]["float"][0]
|
||||
emb_start = r_start["embeddings"]["float"][0]
|
||||
cos = _cosine_sim(emb_end, emb_start)
|
||||
assert cos < 0.99, (
|
||||
f"START and END truncation should produce different embeddings "
|
||||
f"for long input, but cosine similarity was {cos}"
|
||||
)
|
||||
135
third_party/vllm/tests/entrypoints/pooling/embed/test_cohere_online_vision.py
vendored
Normal file
135
third_party/vllm/tests/entrypoints/pooling/embed/test_cohere_online_vision.py
vendored
Normal file
@@ -0,0 +1,135 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the Cohere /v2/embed API with a multimodal model (SigLIP).
|
||||
|
||||
Validates image embedding, batching, normalisation, and embedding type
|
||||
conversions through the /v2/embed endpoint.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
import zlib
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "google/siglip-so400m-patch14-384"
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"64",
|
||||
"--gpu-memory-utilization",
|
||||
"0.3",
|
||||
]
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _make_tiny_png(r: int, g: int, b: int, w: int = 2, h: int = 2) -> str:
|
||||
raw = b""
|
||||
for _ in range(h):
|
||||
raw += b"\x00" + bytes([r, g, b]) * w
|
||||
compressed = zlib.compress(raw)
|
||||
|
||||
def chunk(ctype: bytes, cdata: bytes) -> bytes:
|
||||
c = ctype + cdata
|
||||
return (
|
||||
struct.pack(">I", len(cdata))
|
||||
+ c
|
||||
+ struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
|
||||
)
|
||||
|
||||
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
|
||||
png = (
|
||||
b"\x89PNG\r\n\x1a\n"
|
||||
+ chunk(b"IHDR", ihdr)
|
||||
+ chunk(b"IDAT", compressed)
|
||||
+ chunk(b"IEND", b"")
|
||||
)
|
||||
return "data:image/png;base64," + base64.b64encode(png).decode()
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str] | None = None,
|
||||
images: list[str] | None = None,
|
||||
embedding_types: list[str] | None = None,
|
||||
) -> dict:
|
||||
body: dict = {"model": MODEL_NAME}
|
||||
if texts is not None:
|
||||
body["texts"] = texts
|
||||
if images is not None:
|
||||
body["images"] = images
|
||||
if embedding_types is not None:
|
||||
body["embedding_types"] = embedding_types
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def test_image_embed(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(255, 0, 0)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
assert r["meta"]["billed_units"]["image_tokens"] > 0
|
||||
assert r["meta"]["billed_units"]["input_tokens"] == 0
|
||||
|
||||
|
||||
def test_image_batch(server: RemoteOpenAIServer):
|
||||
red = _make_tiny_png(255, 0, 0)
|
||||
blue = _make_tiny_png(0, 0, 255)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[red, blue],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
assert len(r["embeddings"]["float"]) == 2
|
||||
|
||||
|
||||
def test_image_l2_normalized(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(0, 255, 0)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
emb = np.array(r["embeddings"]["float"][0])
|
||||
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
|
||||
|
||||
|
||||
def test_image_embedding_types(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(128, 128, 128)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float", "binary", "ubinary"],
|
||||
)
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
assert len(r["embeddings"]["binary"][0]) == dim // 8
|
||||
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
|
||||
|
||||
|
||||
def test_text_embed_on_multimodal(server: RemoteOpenAIServer):
|
||||
"""SigLIP also supports text-only embedding via /v2/embed."""
|
||||
r = _cohere_embed(server, texts=["hello world"], embedding_types=["float"])
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
102
third_party/vllm/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py
vendored
Normal file
102
third_party/vllm/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py
vendored
Normal file
@@ -0,0 +1,102 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Parity test between Cohere /v2/embed and OpenAI /v1/embeddings.
|
||||
|
||||
Verifies that both endpoints produce identical float embeddings when
|
||||
no prompt prefix is applied (input_type omitted for Cohere /v2/embed).
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "BAAI/bge-base-en-v1.5"
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--gpu-memory-utilization",
|
||||
"0.02",
|
||||
]
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str],
|
||||
) -> list[list[float]]:
|
||||
body = {
|
||||
"model": MODEL_NAME,
|
||||
"texts": texts,
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()["embeddings"]["float"]
|
||||
|
||||
|
||||
def _openai_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str],
|
||||
) -> list[list[float]]:
|
||||
body = {"model": MODEL_NAME, "input": texts, "encoding_format": "float"}
|
||||
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
|
||||
resp.raise_for_status()
|
||||
return [item["embedding"] for item in resp.json()["data"]]
|
||||
|
||||
|
||||
def test_single_text_parity(server: RemoteOpenAIServer):
|
||||
"""A single text should produce identical embeddings via both APIs."""
|
||||
texts = ["the quick brown fox jumps over the lazy dog"]
|
||||
v2 = _cohere_embed(server, texts)
|
||||
v1 = _openai_embed(server, texts)
|
||||
np.testing.assert_allclose(v2[0], v1[0], rtol=1e-5)
|
||||
|
||||
|
||||
def test_batch_parity(server: RemoteOpenAIServer):
|
||||
"""A batch of texts should produce identical embeddings via both APIs,
|
||||
in the same order."""
|
||||
texts = [
|
||||
"machine learning",
|
||||
"deep learning",
|
||||
"natural language processing",
|
||||
]
|
||||
v2 = _cohere_embed(server, texts)
|
||||
v1 = _openai_embed(server, texts)
|
||||
assert len(v2) == len(v1) == 3
|
||||
for i in range(3):
|
||||
np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
|
||||
|
||||
|
||||
def test_token_count_parity(server: RemoteOpenAIServer):
|
||||
"""Both APIs should report the same prompt token count."""
|
||||
texts = ["hello world"]
|
||||
v2_resp = requests.post(
|
||||
server.url_for("/v2/embed"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"texts": texts,
|
||||
"embedding_types": ["float"],
|
||||
},
|
||||
)
|
||||
v1_resp = requests.post(
|
||||
server.url_for("/v1/embeddings"),
|
||||
json={"model": MODEL_NAME, "input": texts, "encoding_format": "float"},
|
||||
)
|
||||
v2_resp.raise_for_status()
|
||||
v1_resp.raise_for_status()
|
||||
v2_tokens = v2_resp.json()["meta"]["billed_units"]["input_tokens"]
|
||||
v1_tokens = v1_resp.json()["usage"]["prompt_tokens"]
|
||||
assert v2_tokens == v1_tokens
|
||||
46
third_party/vllm/tests/entrypoints/pooling/embed/test_correctness_mteb.py
vendored
Normal file
46
third_party/vllm/tests/entrypoints/pooling/embed/test_correctness_mteb.py
vendored
Normal file
@@ -0,0 +1,46 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling_mteb_test.mteb_embed_utils import (
|
||||
MTEB_EMBED_TASKS,
|
||||
MTEB_EMBED_TOL,
|
||||
OpenAIClientMtebEncoder,
|
||||
run_mteb_embed_task,
|
||||
)
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
|
||||
|
||||
MODEL_NAME = "intfloat/e5-small"
|
||||
MAIN_SCORE = 0.7422994752439667
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = ["--runner", "pooling", "--enforce-eager", "--disable-uvicorn-access-log"]
|
||||
|
||||
# ROCm: Use Flex Attention to support encoder-only self-attention.
|
||||
if current_platform.is_rocm():
|
||||
args.extend(["--attention-backend", "FLEX_ATTENTION"])
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def test_mteb_embed(server):
|
||||
client = server.get_client()
|
||||
encoder = OpenAIClientMtebEncoder(MODEL_NAME, client)
|
||||
vllm_main_score = run_mteb_embed_task(encoder, MTEB_EMBED_TASKS)
|
||||
st_main_score = MAIN_SCORE
|
||||
|
||||
print("VLLM main score: ", vllm_main_score)
|
||||
print("SentenceTransformer main score: ", st_main_score)
|
||||
print("Difference: ", st_main_score - vllm_main_score)
|
||||
|
||||
# We are not concerned that the vllm mteb results are better
|
||||
# than SentenceTransformers, so we only perform one-sided testing.
|
||||
assert st_main_score - vllm_main_score < MTEB_EMBED_TOL
|
||||
208
third_party/vllm/tests/entrypoints/pooling/embed/test_io_processor.py
vendored
Normal file
208
third_party/vllm/tests/entrypoints/pooling/embed/test_io_processor.py
vendored
Normal file
@@ -0,0 +1,208 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for EmbedIOProcessor."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
|
||||
from vllm.entrypoints.pooling.embed.protocol import (
|
||||
CohereEmbedRequest,
|
||||
)
|
||||
|
||||
|
||||
class TestResolveTruncation:
|
||||
"""Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
|
||||
|
||||
@staticmethod
|
||||
def _make_request(**kwargs) -> CohereEmbedRequest:
|
||||
defaults = {
|
||||
"model": "test",
|
||||
"input_type": "search_document",
|
||||
"texts": ["hello"],
|
||||
}
|
||||
return CohereEmbedRequest(**(defaults | kwargs))
|
||||
|
||||
def test_truncate_end_default(self):
|
||||
req = self._make_request()
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side is None
|
||||
|
||||
def test_truncate_end_explicit(self):
|
||||
req = self._make_request(truncate="END")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side is None
|
||||
|
||||
def test_truncate_end_with_max_tokens(self):
|
||||
req = self._make_request(truncate="END", max_tokens=128)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == 128
|
||||
assert side is None
|
||||
|
||||
def test_truncate_none(self):
|
||||
req = self._make_request(truncate="NONE")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens is None
|
||||
assert side is None
|
||||
|
||||
def test_truncate_none_with_max_tokens(self):
|
||||
"""truncate=NONE should NOT set truncate_prompt_tokens; the
|
||||
max_tokens limit is enforced separately via _check_max_tokens."""
|
||||
req = self._make_request(truncate="NONE", max_tokens=10)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens is None
|
||||
assert side is None
|
||||
|
||||
def test_truncate_start(self):
|
||||
req = self._make_request(truncate="START")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side == "left"
|
||||
|
||||
def test_truncate_start_with_max_tokens(self):
|
||||
req = self._make_request(truncate="START", max_tokens=64)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == 64
|
||||
assert side == "left"
|
||||
|
||||
|
||||
class TestApplyStPrompt:
|
||||
"""Unit tests for EmbedIOProcessor._apply_task_instruction."""
|
||||
|
||||
@staticmethod
|
||||
def _make_handler(task_instructions: dict[str, str] | None):
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler.task_instructions = task_instructions
|
||||
return handler
|
||||
|
||||
def test_no_prompts_configured(self):
|
||||
handler = self._make_handler(None)
|
||||
texts = ["hello", "world"]
|
||||
assert handler._apply_task_instruction(texts, "query") is texts
|
||||
|
||||
def test_matching_input_type(self):
|
||||
handler = self._make_handler({"query": "search_query: "})
|
||||
result = handler._apply_task_instruction(["hello"], "query")
|
||||
assert result == ["search_query: hello"]
|
||||
|
||||
def test_non_matching_input_type(self):
|
||||
handler = self._make_handler({"query": "search_query: "})
|
||||
texts = ["hello"]
|
||||
assert handler._apply_task_instruction(texts, "document") is texts
|
||||
|
||||
def test_multiple_texts(self):
|
||||
handler = self._make_handler(
|
||||
{"query": "Represent this sentence for searching: "}
|
||||
)
|
||||
result = handler._apply_task_instruction(["a", "b", "c"], "query")
|
||||
assert result == [
|
||||
"Represent this sentence for searching: a",
|
||||
"Represent this sentence for searching: b",
|
||||
"Represent this sentence for searching: c",
|
||||
]
|
||||
|
||||
def test_empty_prefix_returns_unchanged(self):
|
||||
handler = self._make_handler({"passage": ""})
|
||||
texts = ["hello"]
|
||||
assert handler._apply_task_instruction(texts, "passage") is texts
|
||||
|
||||
|
||||
class TestLoadTaskInstructions:
|
||||
"""Unit tests for EmbedIOProcessor._load_task_instructions."""
|
||||
|
||||
def test_no_attribute(self):
|
||||
class FakeConfig:
|
||||
pass
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
def test_with_task_instructions(self):
|
||||
class FakeConfig:
|
||||
task_instructions = {
|
||||
"retrieval.query": "Represent the query: ",
|
||||
"retrieval.passage": "",
|
||||
}
|
||||
|
||||
result = EmbedIOProcessor._load_task_instructions(FakeConfig())
|
||||
assert result == {
|
||||
"retrieval.query": "Represent the query: ",
|
||||
"retrieval.passage": "",
|
||||
}
|
||||
|
||||
def test_empty_dict(self):
|
||||
class FakeConfig:
|
||||
task_instructions = {}
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
def test_non_dict(self):
|
||||
class FakeConfig:
|
||||
task_instructions = "not a dict"
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
|
||||
class TestCheckMaxTokens:
|
||||
"""Unit tests for EmbedIOProcessor._check_cohere_max_tokens."""
|
||||
|
||||
@staticmethod
|
||||
def _fake_output(n_tokens: int):
|
||||
class _Out:
|
||||
def __init__(self, n: int):
|
||||
self.prompt_token_ids = list(range(n))
|
||||
|
||||
return _Out(n_tokens)
|
||||
|
||||
def test_none_check_is_noop(self):
|
||||
outs = [self._fake_output(100)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, None)
|
||||
|
||||
def test_within_limit(self):
|
||||
outs = [self._fake_output(5), self._fake_output(3)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
def test_exceeds_limit(self):
|
||||
outs = [self._fake_output(3), self._fake_output(10)]
|
||||
with pytest.raises(ValueError, match="exceeds max_tokens=5"):
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
def test_exact_limit(self):
|
||||
outs = [self._fake_output(5)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
|
||||
class TestValidateInputType:
|
||||
"""Unit tests for EmbedIOProcessor._validate_input_type."""
|
||||
|
||||
@staticmethod
|
||||
def _make_handler(task_instructions: dict[str, str] | None):
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler.task_instructions = task_instructions
|
||||
return handler
|
||||
|
||||
def test_none_input_type_always_accepted(self):
|
||||
handler = self._make_handler(None)
|
||||
handler._validate_input_type(None)
|
||||
handler_with = self._make_handler({"query": "q: "})
|
||||
handler_with._validate_input_type(None)
|
||||
|
||||
def test_no_prompts_rejects(self):
|
||||
handler = self._make_handler(None)
|
||||
with pytest.raises(ValueError, match="does not define any input_type"):
|
||||
handler._validate_input_type("anything")
|
||||
|
||||
def test_known_type_accepted(self):
|
||||
handler = self._make_handler({"query": "q: ", "document": "d: "})
|
||||
handler._validate_input_type("query")
|
||||
handler._validate_input_type("document")
|
||||
|
||||
def test_unknown_type_rejected(self):
|
||||
handler = self._make_handler({"query": "q: ", "document": "d: "})
|
||||
with pytest.raises(ValueError, match="Unsupported input_type 'other'"):
|
||||
handler._validate_input_type("other")
|
||||
|
||||
def test_error_lists_supported(self):
|
||||
handler = self._make_handler({"a": "", "b": ""})
|
||||
with pytest.raises(ValueError, match="Supported values: a, b"):
|
||||
handler._validate_input_type("z")
|
||||
72
third_party/vllm/tests/entrypoints/pooling/embed/test_offline.py
vendored
Normal file
72
third_party/vllm/tests/entrypoints/pooling/embed/test_offline.py
vendored
Normal file
@@ -0,0 +1,72 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm import LLM, PoolingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "intfloat/multilingual-e5-small"
|
||||
|
||||
prompts = ["The chef prepared a delicious meal."]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
# ROCm: Use FLEX_ATTENTION backend as it's the only attention backend
|
||||
# that supports encoder-only models on ROCm.
|
||||
attention_config = None
|
||||
if current_platform.is_rocm():
|
||||
attention_config = {"backend": "FLEX_ATTENTION"}
|
||||
|
||||
# pytest caches the fixture so we use weakref.proxy to
|
||||
# enable garbage collection
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
max_num_batched_tokens=32768,
|
||||
tensor_parallel_size=1,
|
||||
gpu_memory_utilization=0.75,
|
||||
enforce_eager=True,
|
||||
seed=0,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_embed(llm: LLM):
|
||||
outputs = llm.encode(prompts, pooling_task="token_embed", use_tqdm=False)
|
||||
multi_vector = outputs[0].outputs.data
|
||||
assert multi_vector.shape == (11, 384)
|
||||
|
||||
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(normalize):
|
||||
outputs = llm.embed(
|
||||
prompts,
|
||||
pooling_params=PoolingParams(use_activation=normalize),
|
||||
use_tqdm=False,
|
||||
)
|
||||
return torch.tensor([x.outputs.embedding for x in outputs])
|
||||
|
||||
default = get_outputs(normalize=None)
|
||||
w_normal = get_outputs(normalize=True)
|
||||
wo_normal = get_outputs(normalize=False)
|
||||
|
||||
assert torch.allclose(default, w_normal, atol=1e-2), "Default should use normal."
|
||||
assert not torch.allclose(w_normal, wo_normal, atol=1e-2), (
|
||||
"wo_normal should not use normal."
|
||||
)
|
||||
assert torch.allclose(w_normal, F.normalize(wo_normal, p=2, dim=-1), atol=1e-2), (
|
||||
"w_normal should be close to normal(wo_normal)."
|
||||
)
|
||||
770
third_party/vllm/tests/entrypoints/pooling/embed/test_online.py
vendored
Normal file
770
third_party/vllm/tests/entrypoints/pooling/embed/test_online.py
vendored
Normal file
@@ -0,0 +1,770 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import requests
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
|
||||
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
|
||||
from vllm.entrypoints.pooling.utils import (
|
||||
MetadataItem,
|
||||
build_metadata_items,
|
||||
decode_pooling_output,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor
|
||||
|
||||
MODEL_NAME = "intfloat/multilingual-e5-small"
|
||||
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
|
||||
DTYPE = "bfloat16"
|
||||
input_text = "The best thing about vLLM is that it supports many different models"
|
||||
input_tokens = [
|
||||
0,
|
||||
581,
|
||||
2965,
|
||||
13580,
|
||||
1672,
|
||||
81,
|
||||
23708,
|
||||
594,
|
||||
83,
|
||||
450,
|
||||
442,
|
||||
8060,
|
||||
7,
|
||||
5941,
|
||||
12921,
|
||||
115774,
|
||||
2,
|
||||
]
|
||||
|
||||
if current_platform.is_rocm():
|
||||
# Disable Flash/MemEfficient SDP on ROCm to avoid HF Transformers
|
||||
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
|
||||
# TODO: Remove once ROCm SDP accuracy issues are resolved on HuggingFace
|
||||
torch.backends.cuda.enable_flash_sdp(False)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(False)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
|
||||
# On ROCm, floating-point reductions in attention and GEMM kernels are
|
||||
# non-associative and sensitive to batch geometry. Force LLM instances
|
||||
# into an identical, deterministic execution mode:
|
||||
ROCM_DETERMINISM_ARGS: list[str] = (
|
||||
["--max-num-seqs", "1"] if current_platform.is_rocm() else []
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--chat-template",
|
||||
DUMMY_CHAT_TEMPLATE,
|
||||
*ROCM_DETERMINISM_ARGS,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def hf_model(hf_runner):
|
||||
with hf_runner(MODEL_NAME, dtype=DTYPE, is_sentence_transformer=True) as hf_model:
|
||||
yield hf_model
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_basic(
|
||||
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
# test /v1/models
|
||||
response = requests.get(server.url_for("/v1/models"))
|
||||
model = response.json()["data"][0]["id"]
|
||||
assert model == MODEL_NAME
|
||||
|
||||
models = await client.models.list()
|
||||
models = models.data
|
||||
served_model = models[0]
|
||||
assert served_model.id == MODEL_NAME
|
||||
|
||||
# test /tokenize
|
||||
response = requests.post(
|
||||
server.url_for("/tokenize"),
|
||||
json={"model": model_name, "prompt": input_text},
|
||||
)
|
||||
assert response.json()["tokens"] == input_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_completion_request(
|
||||
client: openai.AsyncOpenAI, model_name: str, hf_model
|
||||
):
|
||||
# test input: str
|
||||
embedding_response = await client.embeddings.create(
|
||||
model=model_name,
|
||||
input=input_text,
|
||||
encoding_format="float",
|
||||
)
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert len(embeddings.data[0].embedding) == 384
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
assert embeddings.usage.prompt_tokens == len(input_tokens)
|
||||
assert embeddings.usage.total_tokens == len(input_tokens)
|
||||
|
||||
vllm_outputs = [d.embedding for d in embeddings.data]
|
||||
run_embedding_correctness_test(hf_model, [input_text], vllm_outputs)
|
||||
|
||||
# test input: list[int]
|
||||
embedding_response = await client.embeddings.create(
|
||||
model=model_name,
|
||||
input=input_tokens,
|
||||
encoding_format="float",
|
||||
)
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert len(embeddings.data[0].embedding) == 384
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
assert embeddings.usage.prompt_tokens == len(input_tokens)
|
||||
assert embeddings.usage.total_tokens == len(input_tokens)
|
||||
|
||||
vllm_outputs = [d.embedding for d in embeddings.data]
|
||||
run_embedding_correctness_test(hf_model, [input_text], vllm_outputs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_completion_request_batched(
|
||||
client: openai.AsyncOpenAI, model_name: str, hf_model
|
||||
):
|
||||
N = 10
|
||||
input_texts = [input_text] * N
|
||||
|
||||
# test input: list[str]
|
||||
embedding_response = await client.embeddings.create(
|
||||
model=model_name,
|
||||
input=input_texts,
|
||||
encoding_format="float",
|
||||
)
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == N
|
||||
assert len(embeddings.data[0].embedding) == 384
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
assert embeddings.usage.prompt_tokens == len(input_tokens) * N
|
||||
assert embeddings.usage.total_tokens == len(input_tokens) * N
|
||||
|
||||
vllm_outputs = [d.embedding for d in embeddings.data]
|
||||
run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
|
||||
|
||||
# test list[list[int]]
|
||||
embedding_response = await client.embeddings.create(
|
||||
model=model_name,
|
||||
input=[input_tokens] * N,
|
||||
encoding_format="float",
|
||||
)
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == N
|
||||
assert len(embeddings.data[0].embedding) == 384
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
assert embeddings.usage.prompt_tokens == len(input_tokens) * N
|
||||
assert embeddings.usage.total_tokens == len(input_tokens) * N
|
||||
|
||||
vllm_outputs = [d.embedding for d in embeddings.data]
|
||||
run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: str):
|
||||
input_texts = [
|
||||
"Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
|
||||
]
|
||||
|
||||
# test single embedding
|
||||
embedding_response = await client.embeddings.create(
|
||||
model=model_name, input=input_texts, extra_body={"truncate_prompt_tokens": 10}
|
||||
)
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert len(embeddings.data[0].embedding) == 384
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
assert embeddings.usage.prompt_tokens == 10
|
||||
assert embeddings.usage.total_tokens == 10
|
||||
|
||||
input_tokens = [
|
||||
1,
|
||||
24428,
|
||||
289,
|
||||
18341,
|
||||
26165,
|
||||
285,
|
||||
19323,
|
||||
283,
|
||||
289,
|
||||
26789,
|
||||
3871,
|
||||
28728,
|
||||
9901,
|
||||
340,
|
||||
2229,
|
||||
385,
|
||||
340,
|
||||
315,
|
||||
28741,
|
||||
28804,
|
||||
2,
|
||||
]
|
||||
embedding_response = await client.embeddings.create(
|
||||
model=model_name, input=input_tokens, extra_body={"truncate_prompt_tokens": 10}
|
||||
)
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert len(embeddings.data[0].embedding) == 384
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
assert embeddings.usage.prompt_tokens == 10
|
||||
assert embeddings.usage.total_tokens == 10
|
||||
|
||||
# invalid_truncate_prompt_tokens
|
||||
input_texts = [
|
||||
"Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
|
||||
]
|
||||
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
response = await client.embeddings.create(
|
||||
model=model_name,
|
||||
input=input_texts,
|
||||
extra_body={"truncate_prompt_tokens": 8193},
|
||||
)
|
||||
assert "error" in response.object
|
||||
assert (
|
||||
"truncate_prompt_tokens value is greater than max_model_len. "
|
||||
"Please, select a smaller truncation size." in response.message
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_chat_request(
|
||||
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "The cat sat on the mat.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "A feline was resting on a rug.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Stars twinkle brightly in the night sky.",
|
||||
},
|
||||
]
|
||||
|
||||
# test chat request basic usage
|
||||
chat_response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
chat_response.raise_for_status()
|
||||
chat_embeddings = EmbeddingResponse.model_validate(chat_response.json())
|
||||
|
||||
tokenizer = get_tokenizer(tokenizer_name=model_name)
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
chat_template=DUMMY_CHAT_TEMPLATE,
|
||||
add_generation_prompt=True,
|
||||
continue_final_message=False,
|
||||
tokenize=False,
|
||||
)
|
||||
completion_response = await client.embeddings.create(
|
||||
model=model_name,
|
||||
input=prompt,
|
||||
encoding_format="float",
|
||||
# To be consistent with chat
|
||||
extra_body={"add_special_tokens": False},
|
||||
)
|
||||
completion_embeddings = EmbeddingResponse.model_validate(
|
||||
completion_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert chat_embeddings.id is not None
|
||||
assert completion_embeddings.id is not None
|
||||
assert chat_embeddings.created <= completion_embeddings.created
|
||||
# Use tolerance-based comparison for embeddings
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=[d.embedding for d in chat_embeddings.data],
|
||||
embeddings_1_lst=[d.embedding for d in completion_embeddings.data],
|
||||
name_0="chat",
|
||||
name_1="completion",
|
||||
)
|
||||
assert chat_embeddings.model_dump(exclude={"id", "created", "data"}) == (
|
||||
completion_embeddings.model_dump(exclude={"id", "created", "data"})
|
||||
)
|
||||
|
||||
# test add_generation_prompt
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={"model": model_name, "messages": messages, "add_generation_prompt": True},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = EmbeddingResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert output.usage.prompt_tokens == 34
|
||||
|
||||
# test continue_final_message
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"continue_final_message": True,
|
||||
},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = EmbeddingResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert output.usage.prompt_tokens == 33
|
||||
|
||||
# test add_special_tokens
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={"model": model_name, "messages": messages, "add_special_tokens": True},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = EmbeddingResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert output.usage.prompt_tokens == 36
|
||||
|
||||
# test continue_final_message with add_generation_prompt
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"continue_final_message": True,
|
||||
"add_generation_prompt": True,
|
||||
},
|
||||
)
|
||||
assert (
|
||||
"Cannot set both `continue_final_message` and `add_generation_prompt` to True."
|
||||
in response.json()["error"]["message"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invocations_completion_request(
|
||||
server: RemoteOpenAIServer, client: openai.AsyncOpenAI
|
||||
):
|
||||
request_args = {
|
||||
"model": MODEL_NAME,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
|
||||
completion_response = await client.embeddings.create(**request_args)
|
||||
|
||||
invocation_response = requests.post(
|
||||
server.url_for("invocations"), json=request_args
|
||||
)
|
||||
invocation_response.raise_for_status()
|
||||
|
||||
completion_output = completion_response.model_dump()
|
||||
invocation_output = invocation_response.json()
|
||||
|
||||
assert completion_output.keys() == invocation_output.keys()
|
||||
for completion_data, invocation_data in zip(
|
||||
completion_output["data"], invocation_output["data"]
|
||||
):
|
||||
assert completion_data.keys() == invocation_data.keys()
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=[completion_data["embedding"]],
|
||||
embeddings_1_lst=[invocation_data["embedding"]],
|
||||
name_0="completion",
|
||||
name_1="invocation",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invocations_chat_request(server: RemoteOpenAIServer):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "The cat sat on the mat.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "A feline was resting on a rug.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Stars twinkle brightly in the night sky.",
|
||||
},
|
||||
]
|
||||
|
||||
request_args = {
|
||||
"model": MODEL_NAME,
|
||||
"messages": messages,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
|
||||
chat_response = requests.post(server.url_for("v1/embeddings"), json=request_args)
|
||||
chat_response.raise_for_status()
|
||||
|
||||
invocation_response = requests.post(
|
||||
server.url_for("invocations"), json=request_args
|
||||
)
|
||||
invocation_response.raise_for_status()
|
||||
|
||||
chat_output = chat_response.json()
|
||||
invocation_output = invocation_response.json()
|
||||
|
||||
assert chat_output.keys() == invocation_output.keys()
|
||||
for chat_data, invocation_data in zip(
|
||||
chat_output["data"], invocation_output["data"]
|
||||
):
|
||||
assert chat_data.keys() == invocation_data.keys()
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=[chat_data["embedding"]],
|
||||
embeddings_1_lst=[invocation_data["embedding"]],
|
||||
name_0="chat",
|
||||
name_1="invocation",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_base64_embedding(hf_model, client: openai.AsyncOpenAI, model_name: str):
|
||||
input_texts = [
|
||||
"Hello my name is",
|
||||
"The best thing about vLLM is that it supports many different models",
|
||||
]
|
||||
|
||||
responses_float = await client.embeddings.create(
|
||||
input=input_texts, model=model_name, encoding_format="float"
|
||||
)
|
||||
float_data = [d.embedding for d in responses_float.data]
|
||||
run_embedding_correctness_test(hf_model, input_texts, float_data)
|
||||
|
||||
responses_base64 = await client.embeddings.create(
|
||||
input=input_texts, model=model_name, encoding_format="base64"
|
||||
)
|
||||
base64_data = []
|
||||
for data in responses_base64.data:
|
||||
base64_data.append(
|
||||
np.frombuffer(base64.b64decode(data.embedding), dtype="float32").tolist()
|
||||
)
|
||||
|
||||
run_embedding_correctness_test(hf_model, input_texts, base64_data)
|
||||
|
||||
# Default response is float32 decoded from base64 by OpenAI Client
|
||||
responses_default = await client.embeddings.create(
|
||||
input=input_texts, model=model_name
|
||||
)
|
||||
default_data = [d.embedding for d in responses_default.data]
|
||||
run_embedding_correctness_test(hf_model, input_texts, default_data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_base64_embed_dtype_and_endianness(
|
||||
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
input_texts = [input_text] * 3
|
||||
responses_float = await client.embeddings.create(
|
||||
input=input_texts, model=model_name, encoding_format="float"
|
||||
)
|
||||
float_data = [d.embedding for d in responses_float.data]
|
||||
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
responses_base64 = requests.post(
|
||||
server.url_for("/v1/embeddings"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "base64",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
},
|
||||
)
|
||||
|
||||
base64_data = []
|
||||
for data in responses_base64.json()["data"]:
|
||||
binary = base64.b64decode(data["embedding"])
|
||||
tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
|
||||
base64_data.append(tensor.to(torch.float32).tolist())
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=float_data,
|
||||
embeddings_1_lst=base64_data,
|
||||
name_0="float_data",
|
||||
name_1="base64_data",
|
||||
tol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_bytes_embed_dtype_and_endianness(
|
||||
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
input_texts = [input_text] * 3
|
||||
responses_float = await client.embeddings.create(
|
||||
input=input_texts, model=model_name, encoding_format="float"
|
||||
)
|
||||
float_data = [d.embedding for d in responses_float.data]
|
||||
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
responses_bytes = requests.post(
|
||||
server.url_for("/v1/embeddings"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "bytes",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
},
|
||||
)
|
||||
|
||||
metadata = json.loads(responses_bytes.headers["metadata"])
|
||||
body = responses_bytes.content
|
||||
items = [MetadataItem(**x) for x in metadata["data"]]
|
||||
|
||||
bytes_data = decode_pooling_output(items=items, body=body)
|
||||
bytes_data = [x.to(torch.float32).tolist() for x in bytes_data]
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=float_data,
|
||||
embeddings_1_lst=bytes_data,
|
||||
name_0="float_data",
|
||||
name_1="bytes_data",
|
||||
tol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_bytes_only_embed_dtype_and_endianness(
|
||||
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
input_texts = [
|
||||
"The best thing about vLLM is that it supports many different models",
|
||||
] * 2
|
||||
|
||||
responses_float = await client.embeddings.create(
|
||||
input=input_texts, model=model_name, encoding_format="float"
|
||||
)
|
||||
float_data = [d.embedding for d in responses_float.data]
|
||||
embedding_size = len(float_data[0])
|
||||
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
responses_bytes = requests.post(
|
||||
server.url_for("/v1/embeddings"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "bytes_only",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
},
|
||||
)
|
||||
|
||||
assert "metadata" not in responses_bytes.headers
|
||||
body = responses_bytes.content
|
||||
items = build_metadata_items(
|
||||
embed_dtype=embed_dtype,
|
||||
endianness=endianness,
|
||||
shape=(embedding_size,),
|
||||
n_request=len(input_texts),
|
||||
)
|
||||
|
||||
bytes_data = decode_pooling_output(items=items, body=body)
|
||||
bytes_data = [x.to(torch.float32).tolist() for x in bytes_data]
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=float_data,
|
||||
embeddings_1_lst=bytes_data,
|
||||
name_0="float_data",
|
||||
name_1="bytes_data",
|
||||
tol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
|
||||
async def test_params_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, param_name: str
|
||||
):
|
||||
responses_base64 = requests.post(
|
||||
server.url_for("/v1/embeddings"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "base64",
|
||||
param_name: f"bad_{param_name}",
|
||||
},
|
||||
)
|
||||
|
||||
assert responses_base64.status_code == 400
|
||||
assert "literal_error" in responses_base64.json()["error"]["message"]
|
||||
assert f"bad_{param_name}" in responses_base64.json()["error"]["message"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
|
||||
async def get_outputs(use_activation):
|
||||
request_args = {
|
||||
"model": MODEL_NAME,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"use_activation": use_activation,
|
||||
}
|
||||
|
||||
response = requests.post(server.url_for("v1/embeddings"), json=request_args)
|
||||
outputs = response.json()
|
||||
|
||||
return torch.tensor([x["embedding"] for x in outputs["data"]])
|
||||
|
||||
default = await get_outputs(use_activation=None)
|
||||
w_normal = await get_outputs(use_activation=True)
|
||||
wo_normal = await get_outputs(use_activation=False)
|
||||
|
||||
assert torch.allclose(default, w_normal, atol=1e-2), "Default should use normal."
|
||||
assert not torch.allclose(w_normal, wo_normal, atol=1e-2), (
|
||||
"wo_normal should not use normal."
|
||||
)
|
||||
assert torch.allclose(w_normal, F.normalize(wo_normal, p=2, dim=-1), atol=1e-2), (
|
||||
"w_normal should be close to normal(wo_normal)."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "embed"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == 384
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "token_embed"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert len(poolings.data[0].data[0]) == 384
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": "test",
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")
|
||||
131
third_party/vllm/tests/entrypoints/pooling/embed/test_online_dimensions.py
vendored
Normal file
131
third_party/vllm/tests/entrypoints/pooling/embed/test_online_dimensions.py
vendored
Normal file
@@ -0,0 +1,131 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`.
|
||||
"""
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
|
||||
from tests.conftest import HfRunner
|
||||
from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
|
||||
from tests.models.utils import EmbedModelInfo
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODELS = [
|
||||
EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
is_matryoshka=True,
|
||||
matryoshka_dimensions=[256],
|
||||
),
|
||||
]
|
||||
|
||||
input_texts = [
|
||||
"The chef prepared a delicious meal.",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=MODELS)
|
||||
def model_info(request):
|
||||
return request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=["bfloat16"])
|
||||
def dtype(request):
|
||||
return request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(model_info, dtype: str):
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
dtype,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
]
|
||||
|
||||
if model_info.name == "Snowflake/snowflake-arctic-embed-m-v1.5":
|
||||
# Manually enable Matryoshka Embeddings
|
||||
args.extend(
|
||||
["--trust_remote_code", "--hf_overrides", '{"matryoshka_dimensions":[256]}']
|
||||
)
|
||||
|
||||
# ROCm: Use Flex Attention to support encoder-only self-attention.
|
||||
if current_platform.is_rocm():
|
||||
args.extend(["--attention-backend", "FLEX_ATTENTION"])
|
||||
|
||||
with RemoteOpenAIServer(model_info.name, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def hf_model(hf_runner, model_info, dtype: str):
|
||||
with hf_runner(
|
||||
model_info.name, dtype=dtype, is_sentence_transformer=True
|
||||
) as hf_model:
|
||||
yield hf_model
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_matryoshka(
|
||||
model_info: EmbedModelInfo, server: RemoteOpenAIServer, hf_model: HfRunner
|
||||
):
|
||||
client = server.get_async_client()
|
||||
|
||||
async def make_request_and_correctness_test(dimensions):
|
||||
prompts = input_texts * 3
|
||||
|
||||
embedding_response = await client.embeddings.create(
|
||||
model=model_info.name,
|
||||
input=prompts,
|
||||
dimensions=dimensions,
|
||||
encoding_format="float",
|
||||
)
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 3
|
||||
assert len(embeddings.data[0].embedding) > 0
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
assert embeddings.usage.prompt_tokens > 0
|
||||
assert embeddings.usage.total_tokens > 0
|
||||
|
||||
if dimensions is not None:
|
||||
assert len(embeddings.data[0].embedding) == dimensions
|
||||
|
||||
vllm_outputs = [d.embedding for d in embeddings.data]
|
||||
run_embedding_correctness_test(hf_model, prompts, vllm_outputs, dimensions)
|
||||
|
||||
if model_info.is_matryoshka:
|
||||
valid_dimensions: list[int | None] = [None]
|
||||
if model_info.matryoshka_dimensions is not None:
|
||||
valid_dimensions += model_info.matryoshka_dimensions[:2]
|
||||
|
||||
for dimensions in valid_dimensions:
|
||||
await make_request_and_correctness_test(dimensions)
|
||||
|
||||
invalid_dimensions: list[int | None] = [-1]
|
||||
if model_info.matryoshka_dimensions is not None:
|
||||
assert 5 not in model_info.matryoshka_dimensions
|
||||
invalid_dimensions.append(5)
|
||||
|
||||
for dimensions in invalid_dimensions:
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await make_request_and_correctness_test(dimensions)
|
||||
|
||||
else:
|
||||
for dimensions in [None]:
|
||||
await make_request_and_correctness_test(dimensions)
|
||||
|
||||
for dimensions in [-1, 16]:
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await make_request_and_correctness_test(dimensions)
|
||||
457
third_party/vllm/tests/entrypoints/pooling/embed/test_online_long_text.py
vendored
Normal file
457
third_party/vllm/tests/entrypoints/pooling/embed/test_online_long_text.py
vendored
Normal file
@@ -0,0 +1,457 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test cases for long text embedding with automatic chunking mechanism.
|
||||
|
||||
This test suite validates vLLM's automatic chunking functionality for handling
|
||||
text inputs that exceed the model's maximum token length, specifically targeting
|
||||
the intfloat/multilingual-e5-small model (max token length: 512).
|
||||
"""
|
||||
|
||||
import random
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def _generate_random_text(word_count: int) -> str:
|
||||
"""Generate random text with approximately the specified word count."""
|
||||
# Common English words with focus on verbs and nouns for realistic text
|
||||
common_words = [
|
||||
# Essential articles and pronouns (minimal)
|
||||
"the",
|
||||
"and",
|
||||
"you",
|
||||
"they",
|
||||
"this",
|
||||
"that",
|
||||
"these",
|
||||
"those",
|
||||
# Action verbs
|
||||
"create",
|
||||
"build",
|
||||
"develop",
|
||||
"design",
|
||||
"implement",
|
||||
"execute",
|
||||
"analyze",
|
||||
"process",
|
||||
"generate",
|
||||
"calculate",
|
||||
"evaluate",
|
||||
"optimize",
|
||||
"transform",
|
||||
"integrate",
|
||||
"configure",
|
||||
"deploy",
|
||||
"monitor",
|
||||
"manage",
|
||||
"discover",
|
||||
"explore",
|
||||
"investigate",
|
||||
"research",
|
||||
"study",
|
||||
"examine",
|
||||
"improve",
|
||||
"enhance",
|
||||
"upgrade",
|
||||
"modify",
|
||||
"update",
|
||||
"maintain",
|
||||
"solve",
|
||||
"resolve",
|
||||
"handle",
|
||||
"address",
|
||||
"tackle",
|
||||
"overcome",
|
||||
"communicate",
|
||||
"collaborate",
|
||||
"coordinate",
|
||||
"organize",
|
||||
"plan",
|
||||
"achieve",
|
||||
"accomplish",
|
||||
"complete",
|
||||
"finish",
|
||||
"deliver",
|
||||
"provide",
|
||||
# Technology and science nouns
|
||||
"system",
|
||||
"application",
|
||||
"software",
|
||||
"hardware",
|
||||
"network",
|
||||
"database",
|
||||
"algorithm",
|
||||
"model",
|
||||
"framework",
|
||||
"platform",
|
||||
"interface",
|
||||
"protocol",
|
||||
"architecture",
|
||||
"infrastructure",
|
||||
"component",
|
||||
"module",
|
||||
"service",
|
||||
"technology",
|
||||
"innovation",
|
||||
"solution",
|
||||
"methodology",
|
||||
"approach",
|
||||
"artificial",
|
||||
"intelligence",
|
||||
"machine",
|
||||
"learning",
|
||||
"neural",
|
||||
"network",
|
||||
"computer",
|
||||
"processor",
|
||||
"memory",
|
||||
"storage",
|
||||
"computation",
|
||||
"data",
|
||||
"information",
|
||||
"knowledge",
|
||||
"insight",
|
||||
"pattern",
|
||||
"trend",
|
||||
"analysis",
|
||||
"research",
|
||||
"development",
|
||||
"engineering",
|
||||
"science",
|
||||
"mathematics",
|
||||
"statistics",
|
||||
"probability",
|
||||
"optimization",
|
||||
"performance",
|
||||
"efficiency",
|
||||
# General nouns
|
||||
"project",
|
||||
"team",
|
||||
"organization",
|
||||
"company",
|
||||
"business",
|
||||
"industry",
|
||||
"market",
|
||||
"customer",
|
||||
"user",
|
||||
"client",
|
||||
"product",
|
||||
"feature",
|
||||
"function",
|
||||
"requirement",
|
||||
"specification",
|
||||
"documentation",
|
||||
"report",
|
||||
"result",
|
||||
"outcome",
|
||||
"impact",
|
||||
"benefit",
|
||||
"advantage",
|
||||
"challenge",
|
||||
"problem",
|
||||
"opportunity",
|
||||
"strategy",
|
||||
"goal",
|
||||
"objective",
|
||||
"target",
|
||||
"milestone",
|
||||
"process",
|
||||
"procedure",
|
||||
"workflow",
|
||||
"pipeline",
|
||||
"operation",
|
||||
"task",
|
||||
"activity",
|
||||
"event",
|
||||
"session",
|
||||
"meeting",
|
||||
"discussion",
|
||||
"decision",
|
||||
]
|
||||
|
||||
words = []
|
||||
for _ in range(word_count):
|
||||
words.append(random.choice(common_words))
|
||||
|
||||
# Add some punctuation for more realistic text
|
||||
text = " ".join(words)
|
||||
# Add periods every 10-20 words
|
||||
words_list = text.split()
|
||||
result = []
|
||||
for i, word in enumerate(words_list):
|
||||
result.append(word)
|
||||
if (i + 1) % random.randint(10, 20) == 0 and i < len(words_list) - 1:
|
||||
result[-1] += "."
|
||||
|
||||
return " ".join(result)
|
||||
|
||||
|
||||
MODEL_NAME = "intfloat/multilingual-e5-small"
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
# Test text: Generate text with approximately 1500 words to exceed 1024 tokens
|
||||
LONG_TEXT_1500_WORDS = _generate_random_text(1500)
|
||||
|
||||
# Test text: Generate text with approximately 2500 words to exceed 2048 tokens
|
||||
LONG_TEXT_2500_WORDS = _generate_random_text(2500)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server_with_chunked_processing():
|
||||
"""Start server with automatic chunking processing enabled."""
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512", # Set smaller max_model_len to trigger chunking mechanism
|
||||
"--pooler-config",
|
||||
(
|
||||
'{"pooling_type": "MEAN", "use_activation": true, '
|
||||
'"enable_chunked_processing": true, "max_embed_len": 10000}'
|
||||
),
|
||||
"--gpu-memory-utilization",
|
||||
"0.8",
|
||||
]
|
||||
|
||||
# ROCm: Use Flex Attention to support encoder-only self-attention.
|
||||
if current_platform.is_rocm():
|
||||
args.extend(["--attention-backend", "FLEX_ATTENTION"])
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client_with_chunked_processing(server_with_chunked_processing):
|
||||
"""Create async client with chunking processing support."""
|
||||
async with server_with_chunked_processing.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_long_text_embedding_1500_chars(
|
||||
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test embedding processing for ~1500 character long text
|
||||
(~1028 tokens, exceeding 512 token limit)."""
|
||||
|
||||
# Verify text length
|
||||
# Verify text has sufficient word count (approximately 1500 words)
|
||||
word_count = len(LONG_TEXT_1500_WORDS.split())
|
||||
assert word_count >= 1400, f"Test text word count insufficient: {word_count} words"
|
||||
|
||||
# Send embedding request
|
||||
embedding_response = await client_with_chunked_processing.embeddings.create(
|
||||
model=model_name,
|
||||
input=[LONG_TEXT_1500_WORDS],
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
# Verify response structure
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert (
|
||||
len(embeddings.data[0].embedding) == 384
|
||||
) # multilingual-e5-small embedding dimension
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
# Due to chunked processing, token count should
|
||||
# reflect actual processed tokens
|
||||
# With ~1500 words, we expect roughly
|
||||
# 1024+ tokens (exceeding 512 token limit)
|
||||
# Should exceed single chunk limit of 512
|
||||
assert embeddings.usage.prompt_tokens > 800
|
||||
assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
|
||||
|
||||
# Verify embedding vector validity
|
||||
embedding_vector = embeddings.data[0].embedding
|
||||
assert all(isinstance(x, float) for x in embedding_vector), (
|
||||
"Embedding vector should contain floats"
|
||||
)
|
||||
assert not all(x == 0 for x in embedding_vector), (
|
||||
"Embedding vector should not be all zeros"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_long_text_embedding_2500_chars(
|
||||
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test embedding processing for ~2500 character long text
|
||||
(~2048 tokens, requiring multiple chunks)."""
|
||||
|
||||
# Verify text length
|
||||
# Verify text has sufficient word count (approximately 2500 words)
|
||||
word_count = len(LONG_TEXT_2500_WORDS.split())
|
||||
assert word_count >= 2300, f"Test text word count insufficient: {word_count} words"
|
||||
|
||||
# Send embedding request
|
||||
embedding_response = await client_with_chunked_processing.embeddings.create(
|
||||
model=model_name,
|
||||
input=[LONG_TEXT_2500_WORDS],
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
# Verify response structure
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert (
|
||||
len(embeddings.data[0].embedding) == 384
|
||||
) # multilingual-e5-small embedding dimension
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
# Due to chunked processing, token count should
|
||||
# reflect actual processed tokens
|
||||
# With ~2500 words, we expect
|
||||
# roughly 2048+ tokens (requiring multiple chunks)
|
||||
# Should require multiple chunks for processing
|
||||
assert embeddings.usage.prompt_tokens > 1500
|
||||
assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
|
||||
|
||||
# Verify embedding vector validity
|
||||
embedding_vector = embeddings.data[0].embedding
|
||||
assert all(isinstance(x, float) for x in embedding_vector), (
|
||||
"Embedding vector should contain floats"
|
||||
)
|
||||
assert not all(x == 0 for x in embedding_vector), (
|
||||
"Embedding vector should not be all zeros"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_batch_long_text_embedding(
|
||||
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test batch long text embedding processing."""
|
||||
|
||||
input_texts = [
|
||||
LONG_TEXT_1500_WORDS,
|
||||
LONG_TEXT_2500_WORDS,
|
||||
"This is a short text test.", # Short text for comparison
|
||||
]
|
||||
|
||||
# Send batch embedding request
|
||||
embedding_response = await client_with_chunked_processing.embeddings.create(
|
||||
model=model_name,
|
||||
input=input_texts,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
# Verify response structure
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 3 # Three input texts
|
||||
|
||||
# Verify each embedding dimension
|
||||
for i, embedding_data in enumerate(embeddings.data):
|
||||
assert len(embedding_data.embedding) == 384
|
||||
assert embedding_data.index == i
|
||||
|
||||
# Verify embedding vector validity
|
||||
embedding_vector = embedding_data.embedding
|
||||
assert all(isinstance(x, float) for x in embedding_vector)
|
||||
assert not all(x == 0 for x in embedding_vector)
|
||||
|
||||
# Verify token usage
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
# Total token count should be very substantial
|
||||
assert embeddings.usage.prompt_tokens > 1000
|
||||
assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_chunked_vs_normal_consistency(
|
||||
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test consistency between chunked and
|
||||
normal processing (using short text)."""
|
||||
|
||||
# Use a short text within the 512 token limit
|
||||
short_text = (
|
||||
"Artificial intelligence technology is changing our world, "
|
||||
"bringing unprecedented opportunities and challenges."
|
||||
)
|
||||
|
||||
# Send embedding request
|
||||
embedding_response = await client_with_chunked_processing.embeddings.create(
|
||||
model=model_name,
|
||||
input=[short_text],
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
# Verify response structure
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert len(embeddings.data[0].embedding) == 384
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
# Short text should not require chunked processing
|
||||
assert embeddings.usage.prompt_tokens < 512
|
||||
assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
|
||||
|
||||
# 验证embedding向量的有效性
|
||||
embedding_vector = embeddings.data[0].embedding
|
||||
assert all(isinstance(x, float) for x in embedding_vector)
|
||||
assert not all(x == 0 for x in embedding_vector)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_chunked_processing_response_format(
|
||||
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test response format and structure during chunked processing."""
|
||||
|
||||
# Test with long text to trigger chunking
|
||||
embedding_response = await client_with_chunked_processing.embeddings.create(
|
||||
model=model_name,
|
||||
input=[LONG_TEXT_1500_WORDS],
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
# Verify response structure
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert embeddings.data[0].object == "embedding"
|
||||
assert embeddings.data[0].index == 0
|
||||
|
||||
# Verify embedding vector properties
|
||||
embedding_vector = embeddings.data[0].embedding
|
||||
import math
|
||||
|
||||
vector_norm = math.sqrt(sum(x * x for x in embedding_vector))
|
||||
# Check that the vector is normalized
|
||||
# (default behavior for most embedding models)
|
||||
assert 0.8 < vector_norm < 1.2, (
|
||||
f"Vector norm should be reasonable, actual: {vector_norm}"
|
||||
)
|
||||
210
third_party/vllm/tests/entrypoints/pooling/embed/test_online_vision.py
vendored
Normal file
210
third_party/vllm/tests/entrypoints/pooling/embed/test_online_vision.py
vendored
Normal file
@@ -0,0 +1,210 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from tests.utils import VLLM_PATH, RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
|
||||
from vllm.multimodal.media import MediaWithBytes
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
MODEL_NAME = "TIGER-Lab/VLM2Vec-Full"
|
||||
MAXIMUM_IMAGES = 2
|
||||
|
||||
vlm2vec_jinja_path = VLLM_PATH / "examples/pooling/embed/template/vlm2vec_phi3v.jinja"
|
||||
assert vlm2vec_jinja_path.exists()
|
||||
|
||||
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
|
||||
TEST_IMAGE_ASSETS = [
|
||||
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
|
||||
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
|
||||
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
|
||||
]
|
||||
|
||||
input_text = "The best thing about vLLM is that it supports many different models"
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
image_base64 = {"url": encode_image_url(fetch_image(image_url))}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"5",
|
||||
"--enforce-eager",
|
||||
"--trust-remote-code",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"image": MAXIMUM_IMAGES}),
|
||||
"--chat-template",
|
||||
str(vlm2vec_jinja_path),
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_chat_text_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": input_text,
|
||||
},
|
||||
]
|
||||
|
||||
# note: vlm2vec_phi3v.jinja
|
||||
# Embedding models should only embed one message at a time.
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output = EmbeddingResponse.model_validate(response.json())
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data[0].embedding) == 3072
|
||||
assert output.usage.prompt_tokens == 14
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_chat_image_url_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Represent the user's input."},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output = EmbeddingResponse.model_validate(response.json())
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data[0].embedding) == 3072
|
||||
assert output.usage.prompt_tokens == 767
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_chat_image_base64_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Represent the user's input."},
|
||||
{"type": "image_url", "image_url": image_base64},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output = EmbeddingResponse.model_validate(response.json())
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert len(output.data[0].embedding) == 3072
|
||||
assert output.usage.prompt_tokens == 767
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_chat_image_with_media_io_kwargs(server: RemoteOpenAIServer, model_name: str):
|
||||
rgba_image_url = (
|
||||
"https://vllm-public-assets.s3.us-west-2.amazonaws.com"
|
||||
"/vision_model_images/RGBA_comp.png"
|
||||
)
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Represent the user's input."},
|
||||
{"type": "image_url", "image_url": {"url": rgba_image_url}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"media_io_kwargs": {
|
||||
"image": {"rgba_background_color": [0, 0, 0]},
|
||||
},
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output = EmbeddingResponse.model_validate(response.json())
|
||||
assert len(output.data) == 1
|
||||
assert len(output.data[0].embedding) == 3072
|
||||
|
||||
|
||||
def get_hf_prompt_tokens(model_name, content, image_url):
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
model_name, trust_remote_code=True, num_crops=4
|
||||
)
|
||||
|
||||
placeholder = "<|image_1|> "
|
||||
prompt = f"{placeholder}{content}"
|
||||
image = fetch_image(image_url)
|
||||
# Unwrap MediaWithBytes if present
|
||||
if isinstance(image, MediaWithBytes):
|
||||
image = image.media
|
||||
images = [image]
|
||||
inputs = processor(prompt, images, return_tensors="pt")
|
||||
return inputs.input_ids.shape[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
|
||||
async def test_image_embedding(
|
||||
server: RemoteOpenAIServer, model_name: str, image_url: str
|
||||
):
|
||||
content_text = "Represent the given image."
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
{"type": "text", "text": content_text},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("v1/embeddings"),
|
||||
json={"model": model_name, "messages": messages, "encoding_format": "float"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
embeddings = EmbeddingResponse.model_validate(response.json())
|
||||
|
||||
hf_prompt_tokens = get_hf_prompt_tokens(model_name, content_text, image_url)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 1
|
||||
assert len(embeddings.data[0].embedding) == 3072
|
||||
assert embeddings.usage.completion_tokens == 0
|
||||
assert embeddings.usage.prompt_tokens == hf_prompt_tokens
|
||||
assert embeddings.usage.total_tokens == hf_prompt_tokens
|
||||
129
third_party/vllm/tests/entrypoints/pooling/embed/test_protocol.py
vendored
Normal file
129
third_party/vllm/tests/entrypoints/pooling/embed/test_protocol.py
vendored
Normal file
@@ -0,0 +1,129 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for Cohere embed protocol: build_typed_embeddings and its
|
||||
underlying packing helpers, plus Cohere-specific serving helpers."""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.pooling.embed.protocol import (
|
||||
build_typed_embeddings,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_embeddings() -> list[list[float]]:
|
||||
return [
|
||||
[0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8],
|
||||
[-0.05, 0.15, -0.25, 0.35, -0.45, 0.55, -0.65, 0.75],
|
||||
]
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsFloat:
|
||||
def test_float_passthrough(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float"])
|
||||
assert result.float == sample_embeddings
|
||||
assert result.binary is None
|
||||
|
||||
def test_empty_input(self):
|
||||
result = build_typed_embeddings([], ["float"])
|
||||
assert result.float == []
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsBinary:
|
||||
def test_binary_packing(self):
|
||||
# 8 values: positive->1, negative->0 => bits: 10101010 = 0xAA = 170
|
||||
# signed: 170 - 128 = 42
|
||||
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
assert result.binary[0] == [42]
|
||||
|
||||
def test_ubinary_packing(self):
|
||||
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
|
||||
result = build_typed_embeddings(embs, ["ubinary"])
|
||||
assert result.ubinary is not None
|
||||
assert result.ubinary[0] == [170] # 0b10101010
|
||||
|
||||
def test_binary_all_positive(self):
|
||||
embs = [[0.1] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# all bits = 1 => 0xFF = 255, signed: 255 - 128 = 127
|
||||
assert result.binary[0] == [127]
|
||||
|
||||
def test_binary_all_negative(self):
|
||||
embs = [[-0.1] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# all bits = 0, signed: 0 - 128 = -128
|
||||
assert result.binary[0] == [-128]
|
||||
|
||||
def test_binary_dimension_is_eighth(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["binary"])
|
||||
assert result.binary is not None
|
||||
for orig, packed in zip(sample_embeddings, result.binary):
|
||||
assert len(packed) == len(orig) // 8
|
||||
|
||||
def test_zero_treated_as_positive(self):
|
||||
embs = [[0.0] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# 0.0 >= 0 is True, so bit=1 for all => 127 (signed)
|
||||
assert result.binary[0] == [127]
|
||||
|
||||
def test_non_multiple_of_8_raises(self):
|
||||
embs = [[0.1] * 7]
|
||||
with pytest.raises(ValueError, match="multiple of 8"):
|
||||
build_typed_embeddings(embs, ["binary"])
|
||||
|
||||
def test_ubinary_non_multiple_of_8_raises(self):
|
||||
embs = [[0.1] * 10]
|
||||
with pytest.raises(ValueError, match="multiple of 8"):
|
||||
build_typed_embeddings(embs, ["ubinary"])
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsBase64:
|
||||
def test_base64_roundtrip(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["base64"])
|
||||
assert result.base64 is not None
|
||||
assert len(result.base64) == 2
|
||||
|
||||
for orig, b64_str in zip(sample_embeddings, result.base64):
|
||||
decoded = base64.b64decode(b64_str)
|
||||
n = len(orig)
|
||||
values = struct.unpack(f"<{n}f", decoded)
|
||||
np.testing.assert_allclose(orig, values, rtol=1e-5)
|
||||
|
||||
def test_base64_byte_length(self):
|
||||
embs = [[0.1, 0.2, 0.3]]
|
||||
result = build_typed_embeddings(embs, ["base64"])
|
||||
assert result.base64 is not None
|
||||
raw = base64.b64decode(result.base64[0])
|
||||
assert len(raw) == 3 * 4 # 3 floats * 4 bytes each
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsMultiple:
|
||||
def test_all_types_at_once(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(
|
||||
sample_embeddings,
|
||||
["float", "binary", "ubinary", "base64"],
|
||||
)
|
||||
assert result.float is not None
|
||||
assert result.binary is not None
|
||||
assert result.ubinary is not None
|
||||
assert result.base64 is not None
|
||||
|
||||
def test_subset_types(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float", "binary"])
|
||||
assert result.float is not None
|
||||
assert result.binary is not None
|
||||
assert result.ubinary is None
|
||||
assert result.base64 is None
|
||||
|
||||
def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"])
|
||||
assert result.float is not None
|
||||
0
third_party/vllm/tests/entrypoints/pooling/pooling/__init__.py
vendored
Normal file
0
third_party/vllm/tests/entrypoints/pooling/pooling/__init__.py
vendored
Normal file
569
third_party/vllm/tests/entrypoints/pooling/pooling/test_online.py
vendored
Normal file
569
third_party/vllm/tests/entrypoints/pooling/pooling/test_online.py
vendored
Normal file
@@ -0,0 +1,569 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
import torch
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
|
||||
from vllm.entrypoints.pooling.utils import (
|
||||
MetadataItem,
|
||||
build_metadata_items,
|
||||
decode_pooling_output,
|
||||
)
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor
|
||||
|
||||
MODEL_NAME = "internlm/internlm2-1_8b-reward"
|
||||
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
|
||||
input_text = "The chef prepared a delicious meal."
|
||||
input_tokens = [1, 918, 29981, 10166, 395, 18067, 15265, 281]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--chat-template",
|
||||
DUMMY_CHAT_TEMPLATE,
|
||||
"--trust-remote-code",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_basic(server: RemoteOpenAIServer, model_name: str):
|
||||
# test /v1/models
|
||||
response = requests.get(server.url_for("/v1/models"))
|
||||
served_model = response.json()["data"][0]["id"]
|
||||
assert served_model == MODEL_NAME
|
||||
|
||||
# test /tokenize
|
||||
response = requests.post(
|
||||
server.url_for("/tokenize"),
|
||||
json={"model": model_name, "prompt": input_text},
|
||||
)
|
||||
assert response.json()["tokens"] == input_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_completion_request(server: RemoteOpenAIServer, model_name: str):
|
||||
# test input: str
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={"model": model_name, "input": input_text, "encoding_format": "float"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert poolings.id is not None
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert poolings.usage.completion_tokens == 0
|
||||
assert poolings.usage.prompt_tokens == len(input_tokens)
|
||||
assert poolings.usage.total_tokens == len(input_tokens)
|
||||
|
||||
# test input: list[int]
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={"model": model_name, "input": input_tokens, "encoding_format": "float"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert poolings.id is not None
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert poolings.usage.completion_tokens == 0
|
||||
assert poolings.usage.prompt_tokens == len(input_tokens)
|
||||
assert poolings.usage.total_tokens == len(input_tokens)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str):
|
||||
N = 10
|
||||
input_texts = [input_text] * N
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={"model": model_name, "input": input_texts, "encoding_format": "float"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert poolings.id is not None
|
||||
assert len(poolings.data) == N
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert poolings.usage.completion_tokens == 0
|
||||
assert poolings.usage.prompt_tokens == len(input_tokens) * N
|
||||
assert poolings.usage.total_tokens == len(input_tokens) * N
|
||||
|
||||
# test list[list[int]]
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": [input_tokens] * N,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert poolings.id is not None
|
||||
assert len(poolings.data) == N
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert poolings.usage.completion_tokens == 0
|
||||
assert poolings.usage.prompt_tokens == len(input_tokens) * N
|
||||
assert poolings.usage.total_tokens == len(input_tokens) * N
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_chat_request(server: RemoteOpenAIServer, model_name: str):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "The cat sat on the mat.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "A feline was resting on a rug.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Stars twinkle brightly in the night sky.",
|
||||
},
|
||||
]
|
||||
|
||||
# test chat request basic usage
|
||||
chat_response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
chat_response.raise_for_status()
|
||||
chat_poolings = PoolingResponse.model_validate(chat_response.json())
|
||||
|
||||
tokenizer = get_tokenizer(tokenizer_name=model_name, trust_remote_code=True)
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
chat_template=DUMMY_CHAT_TEMPLATE,
|
||||
add_generation_prompt=True,
|
||||
continue_final_message=False,
|
||||
tokenize=False,
|
||||
)
|
||||
completions_response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": prompt,
|
||||
"encoding_format": "float",
|
||||
# To be consistent with chat
|
||||
"add_special_tokens": False,
|
||||
},
|
||||
)
|
||||
completions_response.raise_for_status()
|
||||
completion_poolings = PoolingResponse.model_validate(completions_response.json())
|
||||
|
||||
assert chat_poolings.id is not None
|
||||
assert completion_poolings.id is not None
|
||||
assert chat_poolings.created <= completion_poolings.created
|
||||
assert chat_poolings.model_dump(exclude={"id", "created"}) == (
|
||||
completion_poolings.model_dump(exclude={"id", "created"})
|
||||
)
|
||||
|
||||
# test add_generation_prompt
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={"model": model_name, "messages": messages, "add_generation_prompt": True},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert output.usage.prompt_tokens == 33
|
||||
|
||||
# test continue_final_message
|
||||
# The continue_final_message parameter doesn't seem to be working with this model.
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"continue_final_message": True,
|
||||
},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert output.usage.prompt_tokens == 33
|
||||
|
||||
# test add_special_tokens
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={"model": model_name, "messages": messages, "add_special_tokens": True},
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
output = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert output.object == "list"
|
||||
assert len(output.data) == 1
|
||||
assert output.model == MODEL_NAME
|
||||
assert output.usage.prompt_tokens == 34
|
||||
|
||||
# test continue_final_message with add_generation_prompt
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": messages,
|
||||
"continue_final_message": True,
|
||||
"add_generation_prompt": True,
|
||||
},
|
||||
)
|
||||
assert (
|
||||
"Cannot set both `continue_final_message` and `add_generation_prompt` to True."
|
||||
in response.json()["error"]["message"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_batch_base64_pooling(server: RemoteOpenAIServer, model_name: str):
|
||||
input_texts = [
|
||||
"Hello my name is",
|
||||
"The best thing about vLLM is that it supports many different models",
|
||||
]
|
||||
|
||||
float_response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"input": input_texts,
|
||||
"model": model_name,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
float_response.raise_for_status()
|
||||
responses_float = PoolingResponse.model_validate(float_response.json())
|
||||
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
|
||||
|
||||
base64_response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"input": input_texts,
|
||||
"model": model_name,
|
||||
"encoding_format": "base64",
|
||||
},
|
||||
)
|
||||
base64_response.raise_for_status()
|
||||
responses_base64 = PoolingResponse.model_validate(base64_response.json())
|
||||
|
||||
decoded_responses_base64_data = []
|
||||
for data in responses_base64.data:
|
||||
decoded_responses_base64_data.append(
|
||||
np.frombuffer(base64.b64decode(data.data), dtype="float32").tolist()
|
||||
)
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=float_data,
|
||||
embeddings_1_lst=decoded_responses_base64_data,
|
||||
name_0="float32",
|
||||
name_1="base64",
|
||||
)
|
||||
|
||||
# Default response is float32 decoded from base64 by OpenAI Client
|
||||
default_response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"input": input_texts,
|
||||
"model": model_name,
|
||||
},
|
||||
)
|
||||
default_response.raise_for_status()
|
||||
responses_default = PoolingResponse.model_validate(default_response.json())
|
||||
default_data = [
|
||||
np.array(d.data).squeeze(-1).tolist() for d in responses_default.data
|
||||
]
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=float_data,
|
||||
embeddings_1_lst=default_data,
|
||||
name_0="float32",
|
||||
name_1="default",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_base64_embed_dtype_and_endianness(
|
||||
server: RemoteOpenAIServer, model_name: str
|
||||
):
|
||||
input_texts = [input_text] * 3
|
||||
|
||||
url = server.url_for("pooling")
|
||||
float_response = requests.post(
|
||||
url,
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
responses_float = PoolingResponse.model_validate(float_response.json())
|
||||
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
|
||||
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
responses_base64 = requests.post(
|
||||
url,
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "base64",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
},
|
||||
)
|
||||
|
||||
base64_data = []
|
||||
for data in responses_base64.json()["data"]:
|
||||
binary = base64.b64decode(data["data"])
|
||||
tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
|
||||
base64_data.append(tensor.to(torch.float32).tolist())
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=float_data,
|
||||
embeddings_1_lst=base64_data,
|
||||
name_0="float_data",
|
||||
name_1="base64_data",
|
||||
tol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_bytes_embed_dtype_and_endianness(
|
||||
server: RemoteOpenAIServer, model_name: str
|
||||
):
|
||||
input_texts = [input_text] * 3
|
||||
|
||||
url = server.url_for("pooling")
|
||||
float_response = requests.post(
|
||||
url,
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
responses_float = PoolingResponse.model_validate(float_response.json())
|
||||
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
|
||||
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
responses_bytes = requests.post(
|
||||
url,
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "bytes",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
},
|
||||
)
|
||||
|
||||
metadata = json.loads(responses_bytes.headers["metadata"])
|
||||
body = responses_bytes.content
|
||||
items = [MetadataItem(**x) for x in metadata["data"]]
|
||||
|
||||
bytes_data = decode_pooling_output(items=items, body=body)
|
||||
bytes_data = [x.to(torch.float32).view(-1).tolist() for x in bytes_data]
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=float_data,
|
||||
embeddings_1_lst=bytes_data,
|
||||
name_0="float_data",
|
||||
name_1="bytes_data",
|
||||
tol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_bytes_only_embed_dtype_and_endianness(
|
||||
server: RemoteOpenAIServer, model_name: str
|
||||
):
|
||||
input_texts = [input_text] * 3
|
||||
|
||||
url = server.url_for("pooling")
|
||||
float_response = requests.post(
|
||||
url,
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
responses_float = PoolingResponse.model_validate(float_response.json())
|
||||
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
|
||||
n_tokens = responses_float.usage.prompt_tokens // len(input_texts)
|
||||
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
responses_bytes = requests.post(
|
||||
url,
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_texts,
|
||||
"encoding_format": "bytes_only",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
},
|
||||
)
|
||||
|
||||
assert "metadata" not in responses_bytes.headers
|
||||
body = responses_bytes.content
|
||||
items = build_metadata_items(
|
||||
embed_dtype=embed_dtype,
|
||||
endianness=endianness,
|
||||
shape=(n_tokens, 1),
|
||||
n_request=len(input_texts),
|
||||
)
|
||||
bytes_data = decode_pooling_output(items=items, body=body)
|
||||
bytes_data = [x.to(torch.float32).view(-1).tolist() for x in bytes_data]
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=float_data,
|
||||
embeddings_1_lst=bytes_data,
|
||||
name_0="float_data",
|
||||
name_1="bytes_data",
|
||||
tol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
|
||||
async def test_params_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, param_name: str
|
||||
):
|
||||
responses_base64 = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "base64",
|
||||
param_name: f"bad_{param_name}",
|
||||
},
|
||||
)
|
||||
|
||||
assert responses_base64.status_code == 400
|
||||
assert "literal_error" in responses_base64.json()["error"]["message"]
|
||||
assert f"bad_{param_name}" in responses_base64.json()["error"]["message"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invocations_chat_request(server: RemoteOpenAIServer):
|
||||
request_args = {
|
||||
"model": MODEL_NAME,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
|
||||
completion_response = requests.post(server.url_for("pooling"), json=request_args)
|
||||
completion_response.raise_for_status()
|
||||
|
||||
invocation_response = requests.post(
|
||||
server.url_for("invocations"), json=request_args
|
||||
)
|
||||
invocation_response.raise_for_status()
|
||||
|
||||
completion_output = completion_response.json()
|
||||
invocation_output = invocation_response.json()
|
||||
|
||||
assert completion_output.keys() == invocation_output.keys()
|
||||
for completion_data, invocation_data in zip(
|
||||
completion_output["data"], invocation_output["data"]
|
||||
):
|
||||
assert completion_data.keys() == invocation_data.keys()
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=completion_data["data"],
|
||||
embeddings_1_lst=invocation_data["data"],
|
||||
name_0="completion",
|
||||
name_1="invocation",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invocations_conversation_chat_request(server: RemoteOpenAIServer):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "The cat sat on the mat.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "A feline was resting on a rug.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Stars twinkle brightly in the night sky.",
|
||||
},
|
||||
]
|
||||
|
||||
request_args = {
|
||||
"model": MODEL_NAME,
|
||||
"messages": messages,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
|
||||
chat_response = requests.post(server.url_for("pooling"), json=request_args)
|
||||
chat_response.raise_for_status()
|
||||
|
||||
invocation_response = requests.post(
|
||||
server.url_for("invocations"), json=request_args
|
||||
)
|
||||
invocation_response.raise_for_status()
|
||||
|
||||
chat_output = chat_response.json()
|
||||
invocation_output = invocation_response.json()
|
||||
|
||||
assert chat_output.keys() == invocation_output.keys()
|
||||
for chat_data, invocation_data in zip(
|
||||
chat_output["data"], invocation_output["data"]
|
||||
):
|
||||
assert chat_data.keys() == invocation_data.keys()
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=chat_data["data"],
|
||||
embeddings_1_lst=invocation_data["data"],
|
||||
name_0="chat",
|
||||
name_1="invocation",
|
||||
)
|
||||
0
third_party/vllm/tests/entrypoints/pooling/reward/__init__.py
vendored
Normal file
0
third_party/vllm/tests/entrypoints/pooling/reward/__init__.py
vendored
Normal file
67
third_party/vllm/tests/entrypoints/pooling/reward/test_offline.py
vendored
Normal file
67
third_party/vllm/tests/entrypoints/pooling/reward/test_offline.py
vendored
Normal file
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.models.utils import softmax
|
||||
from vllm import LLM, PoolingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
|
||||
MODEL_NAME = "internlm/internlm2-1_8b-reward"
|
||||
|
||||
prompts = ["The chef prepared a delicious meal."]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
# pytest caches the fixture so we use weakref.proxy to
|
||||
# enable garbage collection
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
max_num_batched_tokens=32768,
|
||||
tensor_parallel_size=1,
|
||||
gpu_memory_utilization=0.75,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
seed=0,
|
||||
)
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_config(llm: LLM):
|
||||
vllm_config = llm.llm_engine.vllm_config
|
||||
assert vllm_config.cache_config.enable_prefix_caching
|
||||
assert vllm_config.scheduler_config.enable_chunked_prefill
|
||||
|
||||
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(use_activation):
|
||||
outputs = llm.reward(
|
||||
prompts,
|
||||
pooling_params=PoolingParams(use_activation=use_activation),
|
||||
use_tqdm=False,
|
||||
)
|
||||
return torch.cat([x.outputs.data for x in outputs])
|
||||
|
||||
default = get_outputs(use_activation=None)
|
||||
w_activation = get_outputs(use_activation=True)
|
||||
wo_activation = get_outputs(use_activation=False)
|
||||
|
||||
assert torch.allclose(default, w_activation, atol=1e-2), (
|
||||
"Default should use activation."
|
||||
)
|
||||
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
|
||||
"wo_activation should not use activation."
|
||||
)
|
||||
assert torch.allclose(softmax(wo_activation), w_activation, atol=1e-2), (
|
||||
"w_activation should be close to activation(wo_activation)."
|
||||
)
|
||||
0
third_party/vllm/tests/entrypoints/pooling/score/__init__.py
vendored
Normal file
0
third_party/vllm/tests/entrypoints/pooling/score/__init__.py
vendored
Normal file
61
third_party/vllm/tests/entrypoints/pooling/score/test_correctness_mteb.py
vendored
Normal file
61
third_party/vllm/tests/entrypoints/pooling/score/test_correctness_mteb.py
vendored
Normal file
@@ -0,0 +1,61 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling_mteb_test.mteb_score_utils import (
|
||||
MTEB_RERANK_LANGS,
|
||||
MTEB_RERANK_TASKS,
|
||||
MTEB_RERANK_TOL,
|
||||
RerankClientMtebEncoder,
|
||||
ScoreClientMtebEncoder,
|
||||
run_mteb_rerank,
|
||||
)
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
|
||||
|
||||
MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
st_main_score = 0.33457
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = ["--runner", "pooling", "--enforce-eager", "--disable-uvicorn-access-log"]
|
||||
|
||||
# ROCm: Use Flex Attention to support encoder-only self-attention.
|
||||
if current_platform.is_rocm():
|
||||
args.extend(["--attention-backend", "FLEX_ATTENTION"])
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def test_mteb_score(server):
|
||||
url = server.url_for("score")
|
||||
encoder = ScoreClientMtebEncoder(MODEL_NAME, url)
|
||||
vllm_main_score = run_mteb_rerank(encoder, MTEB_RERANK_TASKS, MTEB_RERANK_LANGS)
|
||||
|
||||
print("VLLM main score: ", vllm_main_score)
|
||||
print("SentenceTransformer main score: ", st_main_score)
|
||||
print("Difference: ", st_main_score - vllm_main_score)
|
||||
|
||||
# We are not concerned that the vllm mteb results are better
|
||||
# than SentenceTransformers, so we only perform one-sided testing.
|
||||
assert st_main_score - vllm_main_score < MTEB_RERANK_TOL
|
||||
|
||||
|
||||
def test_mteb_rerank(server):
|
||||
url = server.url_for("rerank")
|
||||
encoder = RerankClientMtebEncoder(MODEL_NAME, url)
|
||||
vllm_main_score = run_mteb_rerank(encoder, MTEB_RERANK_TASKS, MTEB_RERANK_LANGS)
|
||||
|
||||
print("VLLM main score: ", vllm_main_score)
|
||||
print("SentenceTransformer main score: ", st_main_score)
|
||||
print("Difference: ", st_main_score - vllm_main_score)
|
||||
|
||||
# We are not concerned that the vllm mteb results are better
|
||||
# than SentenceTransformers, so we only perform one-sided testing.
|
||||
assert st_main_score - vllm_main_score < MTEB_RERANK_TOL
|
||||
69
third_party/vllm/tests/entrypoints/pooling/score/test_offline.py
vendored
Normal file
69
third_party/vllm/tests/entrypoints/pooling/score/test_offline.py
vendored
Normal file
@@ -0,0 +1,69 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.models.utils import softmax
|
||||
from vllm import LLM, PoolingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
# ROCm: Use FLEX_ATTENTION backend as it's the only attention backend
|
||||
# that supports encoder-only models on ROCm.
|
||||
attention_config = None
|
||||
if current_platform.is_rocm():
|
||||
attention_config = {"backend": "FLEX_ATTENTION"}
|
||||
|
||||
# pytest caches the fixture so we use weakref.proxy to
|
||||
# enable garbage collection
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
max_num_batched_tokens=32768,
|
||||
tensor_parallel_size=1,
|
||||
gpu_memory_utilization=0.75,
|
||||
enforce_eager=True,
|
||||
seed=0,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(use_activation):
|
||||
queries = "What is the capital of France?"
|
||||
documents = "The capital of France is Paris."
|
||||
|
||||
outputs = llm.score(
|
||||
queries,
|
||||
documents,
|
||||
pooling_params=PoolingParams(use_activation=use_activation),
|
||||
use_tqdm=False,
|
||||
)
|
||||
return torch.tensor([x.outputs.score for x in outputs])
|
||||
|
||||
default = get_outputs(use_activation=None)
|
||||
w_activation = get_outputs(use_activation=True)
|
||||
wo_activation = get_outputs(use_activation=False)
|
||||
|
||||
assert torch.allclose(default, w_activation, atol=1e-2), (
|
||||
"Default should use activation."
|
||||
)
|
||||
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
|
||||
"wo_activation should not use activation."
|
||||
)
|
||||
assert torch.allclose(softmax(wo_activation), w_activation, atol=1e-2), (
|
||||
"w_activation should be close to activation(wo_activation)."
|
||||
)
|
||||
142
third_party/vllm/tests/entrypoints/pooling/score/test_online_colbert.py
vendored
Normal file
142
third_party/vllm/tests/entrypoints/pooling/score/test_online_colbert.py
vendored
Normal file
@@ -0,0 +1,142 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Online API tests for ColBERT late interaction scoring."""
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
|
||||
|
||||
MODEL_NAME = "answerdotai/answerai-colbert-small-v1"
|
||||
COLBERT_DIM = 96
|
||||
MAX_MODEL_LEN = 512
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--max-model-len",
|
||||
str(MAX_MODEL_LEN),
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
class TestColBERTOnline:
|
||||
def test_rerank(self, server: RemoteOpenAIServer):
|
||||
"""Test ColBERT rerank endpoint."""
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
rerank_response = requests.post(
|
||||
server.url_for("rerank"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
},
|
||||
)
|
||||
rerank_response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(rerank_response.json())
|
||||
|
||||
assert rerank.id is not None
|
||||
assert rerank.results is not None
|
||||
assert len(rerank.results) == 2
|
||||
|
||||
paris_result = next(r for r in rerank.results if r.index == 1)
|
||||
brazil_result = next(r for r in rerank.results if r.index == 0)
|
||||
|
||||
assert paris_result.relevance_score > brazil_result.relevance_score
|
||||
|
||||
def test_rerank_top_n(self, server: RemoteOpenAIServer):
|
||||
"""Test ColBERT rerank with top_n parameter."""
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
"Machine learning is a field of AI.",
|
||||
]
|
||||
|
||||
rerank_response = requests.post(
|
||||
server.url_for("rerank"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
"top_n": 2,
|
||||
},
|
||||
)
|
||||
rerank_response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(rerank_response.json())
|
||||
|
||||
assert len(rerank.results) == 2
|
||||
assert rerank.results[0].index == 1
|
||||
|
||||
def test_score(self, server: RemoteOpenAIServer):
|
||||
"""Test ColBERT score endpoint."""
|
||||
text_1 = "What is the capital of France?"
|
||||
text_2 = ["The capital of France is Paris.", "Python is a language."]
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"text_1": text_1,
|
||||
"text_2": text_2,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 2
|
||||
|
||||
assert score.data[0].score > score.data[1].score
|
||||
|
||||
def test_token_embed(self, server: RemoteOpenAIServer):
|
||||
"""Test ColBERT token_embed task via pooling endpoint."""
|
||||
text = "What is the capital of France?"
|
||||
|
||||
pooling_response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"input": text,
|
||||
"task": "token_embed",
|
||||
},
|
||||
)
|
||||
pooling_response.raise_for_status()
|
||||
pooling = pooling_response.json()
|
||||
|
||||
assert "data" in pooling
|
||||
assert len(pooling["data"]) == 1
|
||||
|
||||
embeddings = pooling["data"][0]["data"]
|
||||
assert isinstance(embeddings, list)
|
||||
assert len(embeddings) > 0
|
||||
assert len(embeddings[0]) == COLBERT_DIM
|
||||
|
||||
def test_embed_not_supported(self, server: RemoteOpenAIServer):
|
||||
"""Test that ColBERT model does not support 'embed' task."""
|
||||
task = "embed"
|
||||
text = "What is the capital of France?"
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"input": text,
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
assert response.json()["error"]["message"].startswith(
|
||||
f"Unsupported task: {task!r}"
|
||||
)
|
||||
235
third_party/vllm/tests/entrypoints/pooling/score/test_online_rerank.py
vendored
Normal file
235
third_party/vllm/tests/entrypoints/pooling/score/test_online_rerank.py
vendored
Normal file
@@ -0,0 +1,235 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
|
||||
from vllm.entrypoints.pooling.score.protocol import RerankResponse
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "BAAI/bge-reranker-base"
|
||||
DTYPE = "bfloat16"
|
||||
input_text = "This product was excellent and exceeded my expectations"
|
||||
input_tokens = [0, 3293, 12996, 509, 40881, 136, 204839, 297, 759, 202702, 2]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
|
||||
|
||||
# ROCm: Use Flex Attention to support encoder-only self-attention.
|
||||
if current_platform.is_rocm():
|
||||
args.extend(["--attention-backend", "FLEX_ATTENTION"])
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_basic(server: RemoteOpenAIServer, model_name: str):
|
||||
# test /v1/models
|
||||
response = requests.get(server.url_for("/v1/models"))
|
||||
served_model = response.json()["data"][0]["id"]
|
||||
assert served_model == MODEL_NAME
|
||||
|
||||
# test /tokenize
|
||||
response = requests.post(
|
||||
server.url_for("/tokenize"),
|
||||
json={"model": model_name, "prompt": input_text},
|
||||
)
|
||||
assert response.json()["tokens"] == input_tokens
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_rerank_texts(server: RemoteOpenAIServer, model_name: str):
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
rerank_response = requests.post(
|
||||
server.url_for("rerank"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
},
|
||||
)
|
||||
rerank_response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(rerank_response.json())
|
||||
|
||||
assert rerank.id is not None
|
||||
assert rerank.results is not None
|
||||
assert len(rerank.results) == 2
|
||||
assert rerank.results[0].relevance_score >= 0.9
|
||||
assert rerank.results[1].relevance_score <= 0.01
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_top_n(server: RemoteOpenAIServer, model_name: str):
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
"Cross-encoder models are neat",
|
||||
]
|
||||
|
||||
rerank_response = requests.post(
|
||||
server.url_for("rerank"),
|
||||
json={"model": model_name, "query": query, "documents": documents, "top_n": 2},
|
||||
)
|
||||
rerank_response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(rerank_response.json())
|
||||
|
||||
assert rerank.id is not None
|
||||
assert rerank.results is not None
|
||||
assert len(rerank.results) == 2
|
||||
assert rerank.results[0].relevance_score >= 0.9
|
||||
assert rerank.results[1].relevance_score <= 0.01
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
def test_rerank_max_model_len(server: RemoteOpenAIServer, model_name: str):
|
||||
query = "What is the capital of France?" * 100
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
rerank_response = requests.post(
|
||||
server.url_for("rerank"),
|
||||
json={"model": model_name, "query": query, "documents": documents},
|
||||
)
|
||||
assert rerank_response.status_code == 400
|
||||
# Assert just a small fragments of the response
|
||||
assert "Please reduce the length of the input." in rerank_response.text
|
||||
|
||||
|
||||
def test_invocations(server: RemoteOpenAIServer):
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
request_args = {
|
||||
"model": MODEL_NAME,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
}
|
||||
|
||||
rerank_response = requests.post(server.url_for("rerank"), json=request_args)
|
||||
rerank_response.raise_for_status()
|
||||
|
||||
invocation_response = requests.post(
|
||||
server.url_for("invocations"), json=request_args
|
||||
)
|
||||
invocation_response.raise_for_status()
|
||||
|
||||
rerank_output = rerank_response.json()
|
||||
invocation_output = invocation_response.json()
|
||||
|
||||
assert rerank_output.keys() == invocation_output.keys()
|
||||
for rerank_result, invocations_result in zip(
|
||||
rerank_output["results"], invocation_output["results"]
|
||||
):
|
||||
assert rerank_result.keys() == invocations_result.keys()
|
||||
assert rerank_result["relevance_score"] == pytest.approx(
|
||||
invocations_result["relevance_score"], rel=0.05
|
||||
)
|
||||
# TODO: reset this tolerance to 0.01 once we find
|
||||
# an alternative to flash_attn with bfloat16
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
|
||||
async def get_outputs(use_activation):
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("rerank"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
"use_activation": use_activation,
|
||||
},
|
||||
)
|
||||
outputs = response.json()
|
||||
|
||||
return torch.tensor([x["relevance_score"] for x in outputs["results"]])
|
||||
|
||||
default = await get_outputs(use_activation=None)
|
||||
w_activation = await get_outputs(use_activation=True)
|
||||
wo_activation = await get_outputs(use_activation=False)
|
||||
|
||||
assert torch.allclose(default, w_activation, atol=1e-2), (
|
||||
"Default should use activation."
|
||||
)
|
||||
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
|
||||
"wo_activation should not use activation."
|
||||
)
|
||||
assert torch.allclose(F.sigmoid(wo_activation), w_activation, atol=1e-2), (
|
||||
"w_activation should be close to activation(wo_activation)."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": "classify",
|
||||
},
|
||||
)
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={"model": model_name, "input": input_text, "encoding_format": "float"},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert len(poolings.data[0].data[0]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")
|
||||
342
third_party/vllm/tests/entrypoints/pooling/score/test_online_score.py
vendored
Normal file
342
third_party/vllm/tests/entrypoints/pooling/score/test_online_score.py
vendored
Normal file
@@ -0,0 +1,342 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import tensor
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.score.protocol import ScoreResponse
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODELS = [
|
||||
{"name": "BAAI/bge-reranker-v2-m3", "is_cross_encoder": True},
|
||||
{"name": "BAAI/bge-base-en-v1.5", "is_cross_encoder": False},
|
||||
]
|
||||
DTYPE = "half"
|
||||
|
||||
|
||||
def run_transformers(hf_model, model, text_pairs):
|
||||
if model["is_cross_encoder"]:
|
||||
return hf_model.predict(text_pairs).tolist()
|
||||
else:
|
||||
hf_embeddings = [hf_model.encode(text_pair) for text_pair in text_pairs]
|
||||
return [
|
||||
F.cosine_similarity(tensor(pair[0]), tensor(pair[1]), dim=0)
|
||||
for pair in hf_embeddings
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="class", params=MODELS)
|
||||
def model(request):
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def server(model: dict[str, Any]):
|
||||
args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
|
||||
|
||||
# ROCm: Use Flex Attention to support encoder-only self-attention.
|
||||
if current_platform.is_rocm():
|
||||
args.extend(["--attention-backend", "FLEX_ATTENTION"])
|
||||
|
||||
with RemoteOpenAIServer(model["name"], args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def runner(model: dict[str, Any], hf_runner):
|
||||
kwargs = {
|
||||
"dtype": DTYPE,
|
||||
"is_cross_encoder"
|
||||
if model["is_cross_encoder"]
|
||||
else "is_sentence_transformer": True,
|
||||
}
|
||||
|
||||
with hf_runner(model["name"], **kwargs) as hf_model:
|
||||
yield hf_model
|
||||
|
||||
|
||||
class TestModel:
|
||||
def test_queries_str_documents_str(
|
||||
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
|
||||
):
|
||||
queries = "What is the capital of France?"
|
||||
documents = "The capital of France is Paris."
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
vllm_outputs = [d.score for d in score.data]
|
||||
|
||||
text_pairs = [[queries, documents]]
|
||||
hf_outputs = run_transformers(runner, model, text_pairs)
|
||||
|
||||
for i in range(len(vllm_outputs)):
|
||||
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
|
||||
|
||||
def test_queries_str_items_str(
|
||||
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
|
||||
):
|
||||
queries = "What is the capital of France?"
|
||||
items = "The capital of France is Paris."
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"queries": queries,
|
||||
"items": items,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
vllm_outputs = [d.score for d in score.data]
|
||||
|
||||
text_pairs = [[queries, items]]
|
||||
hf_outputs = run_transformers(runner, model, text_pairs)
|
||||
|
||||
for i in range(len(vllm_outputs)):
|
||||
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
|
||||
|
||||
def test_text_1_str_text_2_str(
|
||||
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
|
||||
):
|
||||
text_1 = "What is the capital of France?"
|
||||
text_2 = "The capital of France is Paris."
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"text_1": text_1,
|
||||
"text_2": text_2,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
vllm_outputs = [d.score for d in score.data]
|
||||
|
||||
text_pairs = [[text_1, text_2]]
|
||||
hf_outputs = run_transformers(runner, model, text_pairs)
|
||||
|
||||
for i in range(len(vllm_outputs)):
|
||||
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
|
||||
|
||||
def test_data_1_str_data_2_str(
|
||||
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
|
||||
):
|
||||
data_1 = "What is the capital of France?"
|
||||
data_2 = "The capital of France is Paris."
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"data_1": data_1,
|
||||
"data_2": data_2,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
vllm_outputs = [d.score for d in score.data]
|
||||
|
||||
text_pairs = [[data_1, data_2]]
|
||||
hf_outputs = run_transformers(runner, model, text_pairs)
|
||||
|
||||
for i in range(len(vllm_outputs)):
|
||||
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
|
||||
|
||||
def test_queries_str_documents_list(
|
||||
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
|
||||
):
|
||||
queries = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 2
|
||||
|
||||
vllm_outputs = [d.score for d in score.data]
|
||||
|
||||
text_pairs = [[queries, documents[0]], [queries, documents[1]]]
|
||||
hf_outputs = run_transformers(runner, model, text_pairs)
|
||||
|
||||
for i in range(len(vllm_outputs)):
|
||||
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
|
||||
|
||||
def test_queries_list_documents_list(
|
||||
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
|
||||
):
|
||||
queries = [
|
||||
"What is the capital of the United States?",
|
||||
"What is the capital of France?",
|
||||
]
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 2
|
||||
|
||||
vllm_outputs = [d.score for d in score.data]
|
||||
|
||||
text_pairs = [[queries[0], documents[0]], [queries[1], documents[1]]]
|
||||
hf_outputs = run_transformers(runner, model, text_pairs)
|
||||
|
||||
for i in range(len(vllm_outputs)):
|
||||
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
|
||||
|
||||
def test_score_max_model_len(
|
||||
self, server: RemoteOpenAIServer, model: dict[str, Any]
|
||||
):
|
||||
queries = "What is the capital of France?" * 20
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
},
|
||||
)
|
||||
assert score_response.status_code == 400
|
||||
# Assert just a small fragments of the response
|
||||
assert "Please reduce the length of the input." in score_response.text
|
||||
|
||||
# Test truncation
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
"truncate_prompt_tokens": 101,
|
||||
},
|
||||
)
|
||||
assert score_response.status_code == 400
|
||||
assert "Please request a smaller truncation size." in score_response.text
|
||||
|
||||
def test_invocations(self, server: RemoteOpenAIServer, model: dict[str, Any]):
|
||||
queries = "What is the capital of France?"
|
||||
documents = "The capital of France is Paris."
|
||||
|
||||
request_args = {
|
||||
"model": model["name"],
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
}
|
||||
|
||||
score_response = requests.post(server.url_for("score"), json=request_args)
|
||||
score_response.raise_for_status()
|
||||
|
||||
invocation_response = requests.post(
|
||||
server.url_for("invocations"), json=request_args
|
||||
)
|
||||
invocation_response.raise_for_status()
|
||||
|
||||
score_output = score_response.json()
|
||||
invocation_output = invocation_response.json()
|
||||
|
||||
assert score_output.keys() == invocation_output.keys()
|
||||
for score_data, invocation_data in zip(
|
||||
score_output["data"], invocation_output["data"]
|
||||
):
|
||||
assert score_data.keys() == invocation_data.keys()
|
||||
assert score_data["score"] == pytest.approx(
|
||||
invocation_data["score"], rel=0.05
|
||||
)
|
||||
# TODO: reset this tolerance to 0.01 once we find
|
||||
# an alternative to flash_attn with bfloat16
|
||||
|
||||
def test_use_activation(self, server: RemoteOpenAIServer, model: dict[str, Any]):
|
||||
def get_outputs(use_activation):
|
||||
queries = "What is the capital of France?"
|
||||
documents = "The capital of France is Paris."
|
||||
response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": model["name"],
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
"use_activation": use_activation,
|
||||
},
|
||||
)
|
||||
outputs = response.json()
|
||||
return torch.tensor([x["score"] for x in outputs["data"]])
|
||||
|
||||
default = get_outputs(use_activation=None)
|
||||
w_activation = get_outputs(use_activation=True)
|
||||
wo_activation = get_outputs(use_activation=False)
|
||||
|
||||
if model["is_cross_encoder"]:
|
||||
assert torch.allclose(default, w_activation, atol=1e-2), (
|
||||
"Default should use activation."
|
||||
)
|
||||
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
|
||||
"wo_activation should not use activation."
|
||||
)
|
||||
assert torch.allclose(F.sigmoid(wo_activation), w_activation, atol=1e-2), (
|
||||
"w_activation should be close to activation(wo_activation)."
|
||||
)
|
||||
355
third_party/vllm/tests/entrypoints/pooling/score/test_online_score_vision.py
vendored
Normal file
355
third_party/vllm/tests/entrypoints/pooling/score/test_online_score_vision.py
vendored
Normal file
@@ -0,0 +1,355 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import VLLM_PATH, RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-VL-Reranker-2B"
|
||||
HF_OVERRIDES = {
|
||||
"architectures": ["Qwen3VLForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
}
|
||||
|
||||
ROCM_ATTN_BACKENDS = [
|
||||
"ROCM_ATTN",
|
||||
"ROCM_AITER_FA",
|
||||
"TRITON_ATTN",
|
||||
"FLEX_ATTENTION",
|
||||
]
|
||||
|
||||
ATTN_BACKENDS = ROCM_ATTN_BACKENDS if current_platform.is_rocm() else ["auto"]
|
||||
|
||||
# Per-backend tolerance with explicit entries; "default" is the fallback
|
||||
BACKEND_TOL: dict[str, float] = {
|
||||
"default": 0.05, # 5% tolerance for other backends (e.g. FLASH_ATTN)
|
||||
# Relaxed tolerances for ROCm attn
|
||||
# See: https://github.com/vllm-project/vllm/issues/35569
|
||||
"ROCM_ATTN": 0.09, # gfx950:~8.45%, gfx942:~3.70%
|
||||
"ROCM_AITER_FA": 0.045, # gfx950:~2.00%, gfx942:~0.80%
|
||||
"TRITON_ATTN": 0.045, # gfx950:~3.00%, gfx942:~2.20%
|
||||
"FLEX_ATTENTION": 0.045, # gfx950:~3.25%, gfx942:~1.10%
|
||||
}
|
||||
|
||||
# ROCm: disable skinny GEMM to avoid non-deterministic results from
|
||||
# atomic reductions in wvSplitKrc kernel.
|
||||
# See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
|
||||
ROCM_ENV_OVERRIDES = (
|
||||
{"VLLM_ROCM_USE_SKINNY_GEMM": "0"} if current_platform.is_rocm() else {}
|
||||
)
|
||||
# ROCm: disable prefix caching and eliminate batch variance to reduce
|
||||
# test flakiness.
|
||||
ROCM_EXTRA_ARGS = (
|
||||
["--no-enable-prefix-caching", "--max-num-seqs", "1"]
|
||||
if current_platform.is_rocm()
|
||||
else []
|
||||
)
|
||||
|
||||
|
||||
def get_tol(backend: str) -> float:
|
||||
return BACKEND_TOL.get(backend, BACKEND_TOL["default"])
|
||||
|
||||
|
||||
def assert_score(actual: float, expected: float, backend: str, label: str):
|
||||
tol = get_tol(backend)
|
||||
diff = abs(actual - expected)
|
||||
rel_diff = diff / abs(expected) if expected != 0 else diff
|
||||
print(
|
||||
f"[{backend}] {label}: actual={actual:.6f} expected={expected:.6f} "
|
||||
f"diff={diff:.6f} rel_diff={rel_diff:.4f} tol={tol}"
|
||||
)
|
||||
assert actual == pytest.approx(expected, rel=tol), (
|
||||
f"[{backend}] {label}: score mismatch — "
|
||||
f"actual={actual:.6f}, expected={expected:.6f}, "
|
||||
f"rel_diff={rel_diff:.4f}, tol={tol}"
|
||||
)
|
||||
|
||||
|
||||
query = "A cat standing in the snow."
|
||||
document = "This product was excellent and exceeded my expectations."
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": document,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
]
|
||||
|
||||
TEXT_VS_TEXT = 0.10040374100208282
|
||||
TEXT_VS_IMAGE = 0.7423753142356873
|
||||
TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=ATTN_BACKENDS)
|
||||
def server(request):
|
||||
backend = request.param
|
||||
print(f"\n=== Starting server with attention backend: {backend} ===")
|
||||
args = [
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"8192",
|
||||
"--chat-template",
|
||||
str(VLLM_PATH / "examples/pooling/score/template/qwen3_vl_reranker.jinja"),
|
||||
]
|
||||
|
||||
env = dict()
|
||||
if backend != "auto":
|
||||
args += ["--attention-config", json.dumps({"backend": backend})]
|
||||
args += ROCM_EXTRA_ARGS
|
||||
|
||||
env = dict(ROCM_ENV_OVERRIDES)
|
||||
if backend != "ROCM_AITER_FA":
|
||||
env["VLLM_ROCM_USE_AITER"] = "0"
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, args, override_hf_configs=HF_OVERRIDES, env_dict=env
|
||||
) as remote_server:
|
||||
print(f"=== Server ready with backend: {backend} ===")
|
||||
yield remote_server, backend
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_str(server: tuple[RemoteOpenAIServer, str]):
|
||||
remote_server, backend = server
|
||||
score_response = requests.post(
|
||||
remote_server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": document,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
assert score.usage.prompt_tokens == 81
|
||||
assert_score(score.data[0].score, TEXT_VS_TEXT, backend, "text_vs_text")
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_text_content(
|
||||
server: tuple[RemoteOpenAIServer, str],
|
||||
):
|
||||
remote_server, backend = server
|
||||
score_response = requests.post(
|
||||
remote_server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[0]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
assert score.usage.prompt_tokens == 81
|
||||
assert_score(score.data[0].score, TEXT_VS_TEXT, backend, "text_vs_text")
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_image_url_content(
|
||||
server: tuple[RemoteOpenAIServer, str],
|
||||
):
|
||||
remote_server, backend = server
|
||||
score_response = requests.post(
|
||||
remote_server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[1]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
assert score.usage.prompt_tokens == 98
|
||||
assert_score(score.data[0].score, TEXT_VS_IMAGE, backend, "text_vs_image")
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_image_base64_content(
|
||||
server: tuple[RemoteOpenAIServer, str],
|
||||
):
|
||||
remote_server, backend = server
|
||||
score_response = requests.post(
|
||||
remote_server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
assert score.usage.prompt_tokens == 98
|
||||
assert_score(score.data[0].score, TEXT_VS_IMAGE, backend, "text_vs_image_base64")
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_image_url_plus_text_content(
|
||||
server: tuple[RemoteOpenAIServer, str],
|
||||
):
|
||||
remote_server, backend = server
|
||||
score_response = requests.post(
|
||||
remote_server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[0], documents[1]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
assert score.usage.prompt_tokens == 108
|
||||
assert_score(
|
||||
score.data[0].score, TEXT_VS_TEXT_PLUS_IMAGE, backend, "text_vs_text_plus_image"
|
||||
)
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_list(
|
||||
server: tuple[RemoteOpenAIServer, str],
|
||||
):
|
||||
remote_server, backend = server
|
||||
score_response = requests.post(
|
||||
remote_server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": [
|
||||
document,
|
||||
{"content": [documents[0]]},
|
||||
{"content": [documents[1]]},
|
||||
{"content": [documents[0], documents[1]]},
|
||||
],
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 4
|
||||
assert score.usage.prompt_tokens == 368
|
||||
assert_score(score.data[0].score, TEXT_VS_TEXT, backend, "list[0]_text_vs_text")
|
||||
assert_score(score.data[1].score, TEXT_VS_TEXT, backend, "list[1]_text_vs_text")
|
||||
assert_score(score.data[2].score, TEXT_VS_IMAGE, backend, "list[2]_text_vs_image")
|
||||
assert_score(
|
||||
score.data[3].score,
|
||||
TEXT_VS_TEXT_PLUS_IMAGE,
|
||||
backend,
|
||||
"list[3]_text_vs_text_plus_image",
|
||||
)
|
||||
|
||||
|
||||
def test_rerank_api_queries_str_documents_list(
|
||||
server: tuple[RemoteOpenAIServer, str],
|
||||
):
|
||||
remote_server, backend = server
|
||||
rerank_response = requests.post(
|
||||
remote_server.url_for("rerank"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"query": query,
|
||||
"documents": [
|
||||
document,
|
||||
{"content": [documents[0]]},
|
||||
{"content": [documents[1]]},
|
||||
{"content": [documents[0], documents[1]]},
|
||||
],
|
||||
},
|
||||
)
|
||||
rerank_response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(rerank_response.json())
|
||||
|
||||
assert rerank.id is not None
|
||||
assert rerank.model is not None
|
||||
assert rerank.usage is not None
|
||||
assert len(rerank.results) == 4
|
||||
|
||||
rerank.results.sort(key=lambda x: x.index)
|
||||
assert_score(
|
||||
rerank.results[0].relevance_score,
|
||||
TEXT_VS_TEXT,
|
||||
backend,
|
||||
"rerank[0]_text_vs_text",
|
||||
)
|
||||
assert_score(
|
||||
rerank.results[1].relevance_score,
|
||||
TEXT_VS_TEXT,
|
||||
backend,
|
||||
"rerank[1]_text_vs_text",
|
||||
)
|
||||
assert_score(
|
||||
rerank.results[2].relevance_score,
|
||||
TEXT_VS_IMAGE,
|
||||
backend,
|
||||
"rerank[2]_text_vs_image",
|
||||
)
|
||||
assert_score(
|
||||
rerank.results[3].relevance_score,
|
||||
TEXT_VS_TEXT_PLUS_IMAGE,
|
||||
backend,
|
||||
"rerank[3]_text_vs_text_plus_image",
|
||||
)
|
||||
|
||||
|
||||
def test_score_api_queries_list_documents_list(
|
||||
server: tuple[RemoteOpenAIServer, str],
|
||||
):
|
||||
remote_server, backend = server
|
||||
score_response = requests.post(
|
||||
remote_server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": [query] * 4,
|
||||
"documents": [
|
||||
document,
|
||||
{"content": [documents[0]]},
|
||||
{"content": [documents[1]]},
|
||||
{"content": [documents[0], documents[1]]},
|
||||
],
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 4
|
||||
assert score.usage.prompt_tokens == 368
|
||||
assert_score(score.data[0].score, TEXT_VS_TEXT, backend, "paired[0]_text_vs_text")
|
||||
assert_score(score.data[1].score, TEXT_VS_TEXT, backend, "paired[1]_text_vs_text")
|
||||
assert_score(score.data[2].score, TEXT_VS_IMAGE, backend, "paired[2]_text_vs_image")
|
||||
assert_score(
|
||||
score.data[3].score,
|
||||
TEXT_VS_TEXT_PLUS_IMAGE,
|
||||
backend,
|
||||
"paired[3]_text_vs_text_plus_image",
|
||||
)
|
||||
353
third_party/vllm/tests/entrypoints/pooling/score/test_utils.py
vendored
Normal file
353
third_party/vllm/tests/entrypoints/pooling/score/test_utils.py
vendored
Normal file
@@ -0,0 +1,353 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
|
||||
from vllm.entrypoints.pooling.score.utils import (
|
||||
get_score_prompt,
|
||||
)
|
||||
from vllm.inputs import TokensPrompt
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
|
||||
# A cross-encoder model for testing
|
||||
CROSS_ENCODER_MODEL_ID = "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
|
||||
|
||||
def assert_prompt_tokenization_consistent(
|
||||
tokenizer, full_prompt, engine_prompt, add_special_tokens=True
|
||||
):
|
||||
"""Verify that engine_prompt token_ids match tokenizing full_prompt."""
|
||||
expected_ids = tokenizer(full_prompt, add_special_tokens=add_special_tokens)[
|
||||
"input_ids"
|
||||
]
|
||||
actual_ids = engine_prompt["prompt_token_ids"]
|
||||
assert actual_ids == expected_ids, (
|
||||
f"Token IDs don't match.\nExpected: {expected_ids}\nActual: {actual_ids}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cross_encoder_model_config():
|
||||
return ModelConfig(
|
||||
CROSS_ENCODER_MODEL_ID,
|
||||
runner="pooling",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cross_encoder_tokenizer(cross_encoder_model_config):
|
||||
return get_tokenizer(
|
||||
CROSS_ENCODER_MODEL_ID,
|
||||
trust_remote_code=cross_encoder_model_config.trust_remote_code,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm_reranker_model_config():
|
||||
"""Model config for LLM-as-reranker style (no pad token)."""
|
||||
config = ModelConfig(
|
||||
CROSS_ENCODER_MODEL_ID,
|
||||
runner="pooling",
|
||||
)
|
||||
# use_sep_token is a property that reads from hf_config,
|
||||
# so we set it there to override the default (True)
|
||||
config.hf_config.use_sep_token = False
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tokenization_kwargs():
|
||||
"""Common tokenization kwargs used across tests."""
|
||||
return {"add_special_tokens": True, "return_tensors": None}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model_with_score_template():
|
||||
"""Mock model class that supports score template and tracks post_process calls."""
|
||||
|
||||
class MockModelWithScoreTemplate:
|
||||
supports_score_template = True
|
||||
post_process_called: list[TokensPrompt] = []
|
||||
|
||||
@staticmethod
|
||||
def get_score_template(p1: str, p2: str) -> str:
|
||||
return f"[QUERY]{p1}[SEP][DOC]{p2}"
|
||||
|
||||
@staticmethod
|
||||
def post_process_tokens(prompt: TokensPrompt) -> None:
|
||||
MockModelWithScoreTemplate.post_process_called.append(prompt)
|
||||
|
||||
return MockModelWithScoreTemplate
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model_no_score_template():
|
||||
"""Mock model class that does not support score template."""
|
||||
|
||||
class MockModelNoScoreTemplate:
|
||||
supports_score_template = False
|
||||
|
||||
return MockModelNoScoreTemplate
|
||||
|
||||
|
||||
class TestGetScorePrompt:
|
||||
"""Tests for the get_score_prompt function."""
|
||||
|
||||
def test_tokenization_kwargs_passed_through(
|
||||
self,
|
||||
llm_reranker_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
):
|
||||
"""Test that tokenization kwargs are properly passed through."""
|
||||
data_1 = "Query text"
|
||||
data_2 = "Document text"
|
||||
|
||||
# Test with truncation - custom kwargs for this test
|
||||
custom_tokenization_kwargs = {
|
||||
"add_special_tokens": True,
|
||||
"return_tensors": None,
|
||||
"truncation": True,
|
||||
"max_length": 20,
|
||||
}
|
||||
|
||||
full_prompt, engine_prompt = get_score_prompt(
|
||||
llm_reranker_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
custom_tokenization_kwargs,
|
||||
data_1,
|
||||
data_2,
|
||||
)
|
||||
|
||||
assert isinstance(full_prompt, str)
|
||||
assert "prompt_token_ids" in engine_prompt
|
||||
# With max_length=20 and truncation, should not exceed this
|
||||
assert len(engine_prompt["prompt_token_ids"]) <= 20
|
||||
# Since truncation was applied, token_ids should be a prefix of full encoding
|
||||
full_ids = cross_encoder_tokenizer(full_prompt, add_special_tokens=True)[
|
||||
"input_ids"
|
||||
]
|
||||
actual_ids = engine_prompt["prompt_token_ids"]
|
||||
assert full_ids[: len(actual_ids)] == actual_ids, (
|
||||
f"Token IDs are not a prefix of full encoding.\n"
|
||||
f"Full IDs: {full_ids}\n"
|
||||
f"Actual IDs: {actual_ids}"
|
||||
)
|
||||
|
||||
def test_model_supports_score_template(
|
||||
self,
|
||||
cross_encoder_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
mock_model_with_score_template,
|
||||
):
|
||||
"""Test when model supports score template (no score_template arg)."""
|
||||
with patch(
|
||||
"vllm.model_executor.model_loader.get_model_cls",
|
||||
return_value=mock_model_with_score_template,
|
||||
):
|
||||
full_prompt, engine_prompt = get_score_prompt(
|
||||
cross_encoder_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
"query text",
|
||||
"document text",
|
||||
)
|
||||
|
||||
assert full_prompt == "[QUERY]query text[SEP][DOC]document text"
|
||||
assert "prompt_token_ids" in engine_prompt
|
||||
assert len(engine_prompt["prompt_token_ids"]) > 0
|
||||
assert_prompt_tokenization_consistent(
|
||||
cross_encoder_tokenizer, full_prompt, engine_prompt
|
||||
)
|
||||
|
||||
def test_model_supports_score_template_but_custom_template_provided(
|
||||
self,
|
||||
cross_encoder_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
mock_model_with_score_template,
|
||||
):
|
||||
"""Test when model supports score template but custom template is provided."""
|
||||
template = (
|
||||
'TEMPLATE_USED {{ messages[0]["content"] }} {{ messages[1]["content"] }}'
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"vllm.model_executor.model_loader.get_model_cls",
|
||||
return_value=mock_model_with_score_template,
|
||||
),
|
||||
):
|
||||
full_prompt, engine_prompt = get_score_prompt(
|
||||
cross_encoder_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
"query",
|
||||
"doc",
|
||||
score_template=template, # Providing a template
|
||||
)
|
||||
|
||||
assert "prompt_token_ids" in engine_prompt
|
||||
assert full_prompt == "TEMPLATE_USED query doc"
|
||||
|
||||
assert_prompt_tokenization_consistent(
|
||||
cross_encoder_tokenizer, full_prompt, engine_prompt
|
||||
)
|
||||
|
||||
def test_not_using_default_template(
|
||||
self,
|
||||
llm_reranker_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
mock_model_no_score_template,
|
||||
):
|
||||
# FIXME: For now, we only apply a template when one is explicitly provided.
|
||||
# We cannot rely on the tokenizer's chat template because many models
|
||||
# inherit junk templates from their base LLM, which breaks both the models
|
||||
# and the tests that use them.
|
||||
with (
|
||||
patch(
|
||||
"vllm.model_executor.model_loader.get_model_cls",
|
||||
return_value=mock_model_no_score_template,
|
||||
),
|
||||
patch(
|
||||
"vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
|
||||
return_value="test querytest doc",
|
||||
),
|
||||
):
|
||||
full_prompt, engine_prompt = get_score_prompt(
|
||||
llm_reranker_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
"test query",
|
||||
"test doc",
|
||||
)
|
||||
|
||||
assert full_prompt == "test querytest doc"
|
||||
assert "prompt_token_ids" in engine_prompt
|
||||
assert_prompt_tokenization_consistent(
|
||||
cross_encoder_tokenizer, full_prompt, engine_prompt
|
||||
)
|
||||
|
||||
def test_fallback_with_sep_token(
|
||||
self,
|
||||
cross_encoder_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
mock_model_no_score_template,
|
||||
):
|
||||
"""Test fallback path when ChatTemplateResolutionError
|
||||
and use_sep_token=True."""
|
||||
with (
|
||||
patch(
|
||||
"vllm.model_executor.model_loader.get_model_cls",
|
||||
return_value=mock_model_no_score_template,
|
||||
),
|
||||
patch(
|
||||
"vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
|
||||
side_effect=ChatTemplateResolutionError("No template"),
|
||||
),
|
||||
):
|
||||
full_prompt, engine_prompt = get_score_prompt(
|
||||
cross_encoder_model_config, # use_sep_token=True
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
"query",
|
||||
"document",
|
||||
)
|
||||
|
||||
assert "prompt_token_ids" in engine_prompt
|
||||
# Should have token_type_ids from text_pair encoding
|
||||
assert "token_type_ids" in engine_prompt
|
||||
assert "query" in full_prompt
|
||||
assert "document" in full_prompt
|
||||
assert full_prompt != "querydocument"
|
||||
assert (
|
||||
engine_prompt["prompt_token_ids"]
|
||||
== cross_encoder_tokenizer(
|
||||
"query", text_pair="document", add_special_tokens=True
|
||||
)["input_ids"]
|
||||
)
|
||||
|
||||
# FIXME(?): add_special_tokens=False is needed because in this case
|
||||
# full_prompt is obtained by decoding the tokenized prompt, which includes
|
||||
# special tokens and we would get duplicated special tokens otherwise.
|
||||
# This is inconsistent with other cases.
|
||||
assert_prompt_tokenization_consistent(
|
||||
cross_encoder_tokenizer,
|
||||
full_prompt,
|
||||
engine_prompt,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
|
||||
def test_fallback_without_sep_token(
|
||||
self,
|
||||
llm_reranker_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
mock_model_no_score_template,
|
||||
):
|
||||
"""Test fallback path when ChatTemplateResolutionError
|
||||
and use_sep_token=False."""
|
||||
with (
|
||||
patch(
|
||||
"vllm.model_executor.model_loader.get_model_cls",
|
||||
return_value=mock_model_no_score_template,
|
||||
),
|
||||
patch(
|
||||
"vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
|
||||
side_effect=ChatTemplateResolutionError("No template"),
|
||||
),
|
||||
):
|
||||
full_prompt, engine_prompt = get_score_prompt(
|
||||
llm_reranker_model_config, # use_sep_token=False
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
"query",
|
||||
"document",
|
||||
)
|
||||
|
||||
assert full_prompt == "querydocument"
|
||||
assert "prompt_token_ids" in engine_prompt
|
||||
assert_prompt_tokenization_consistent(
|
||||
cross_encoder_tokenizer, full_prompt, engine_prompt
|
||||
)
|
||||
|
||||
def test_post_process_tokens_called(
|
||||
self,
|
||||
cross_encoder_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
mock_model_with_score_template,
|
||||
):
|
||||
"""Test that post_process_tokens is called on the engine prompt."""
|
||||
# Reset the call tracker
|
||||
mock_model_with_score_template.post_process_called.clear()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.model_executor.model_loader.get_model_cls",
|
||||
return_value=mock_model_with_score_template,
|
||||
),
|
||||
patch(
|
||||
"vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
|
||||
side_effect=ChatTemplateResolutionError("No template"),
|
||||
),
|
||||
):
|
||||
full_prompt, engine_prompt = get_score_prompt(
|
||||
cross_encoder_model_config,
|
||||
cross_encoder_tokenizer,
|
||||
tokenization_kwargs,
|
||||
"query",
|
||||
"doc",
|
||||
)
|
||||
|
||||
# post_process_tokens should have been called once
|
||||
assert len(mock_model_with_score_template.post_process_called) == 1
|
||||
assert mock_model_with_score_template.post_process_called[0] is engine_prompt
|
||||
assert_prompt_tokenization_consistent(
|
||||
cross_encoder_tokenizer, full_prompt, engine_prompt
|
||||
)
|
||||
Reference in New Issue
Block a user