72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
from tests.utils import RemoteOpenAIServer
|
|
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,
|
|
"--pooler-config.task",
|
|
"token_classify",
|
|
]
|
|
|
|
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
|
yield remote_server
|
|
|
|
|
|
@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,
|
|
},
|
|
)
|
|
|
|
if task == "plugin":
|
|
err_msg = "No IOProcessor plugin installed."
|
|
else:
|
|
err_msg = f"Unsupported task: {task!r}"
|
|
assert response.json()["error"]["message"].startswith(err_msg)
|