test_classification.py 9.29 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5

import pytest
import requests
6
7
import torch
import torch.nn.functional as F
8

9
from tests.utils import RemoteOpenAIServer
10
from vllm.entrypoints.openai.protocol import ClassificationResponse, PoolingResponse
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
DTYPE = "float32"  # Use float32 to avoid NaN issue


@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])
31
def test_single_input_classification(server: RemoteOpenAIServer, model_name: str):
32
33
34
35
    input_text = "This product was excellent and exceeded my expectations"

    classification_response = requests.post(
        server.url_for("classify"),
36
        json={"model": model_name, "input": input_text},
37
38
39
    )

    classification_response.raise_for_status()
40
    output = ClassificationResponse.model_validate(classification_response.json())
41
42
43
44
45
46
47
48
49

    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])
50
def test_multiple_inputs_classification(server: RemoteOpenAIServer, model_name: str):
51
52
53
54
55
56
57
58
59
60
61
    input_texts = [
        "The product arrived on time and works perfectly",
        "I'm very satisfied with my purchase, would buy again",
        "The customer service was helpful and resolved my issue quickly",
        "This product broke after one week, terrible quality",
        "I'm very disappointed with this purchase, complete waste of money",
        "The customer service was rude and unhelpful",
    ]

    classification_response = requests.post(
        server.url_for("classify"),
62
        json={"model": model_name, "input": input_texts},
63
    )
64
    output = ClassificationResponse.model_validate(classification_response.json())
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

    assert len(output.data) == len(input_texts)
    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_truncate_prompt_tokens(server: RemoteOpenAIServer, model_name: str):
    long_text = "hello " * 600

    classification_response = requests.post(
        server.url_for("classify"),
81
        json={"model": model_name, "input": long_text, "truncate_prompt_tokens": 5},
82
83
84
    )

    classification_response.raise_for_status()
85
    output = ClassificationResponse.model_validate(classification_response.json())
86
87
88
89
90
91
92
93
94

    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


@pytest.mark.parametrize("model_name", [MODEL_NAME])
95
96
97
def test_invalid_truncate_prompt_tokens_error(
    server: RemoteOpenAIServer, model_name: str
):
98
99
    classification_response = requests.post(
        server.url_for("classify"),
100
        json={"model": model_name, "input": "test", "truncate_prompt_tokens": 513},
101
102
103
104
    )

    error = classification_response.json()
    assert classification_response.status_code == 400
105
    assert "truncate_prompt_tokens" in error["error"]["message"]
106
107
108
109
110
111


@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"),
112
        json={"model": model_name, "input": ""},
113
114
115
116
    )

    error = classification_response.json()
    assert classification_response.status_code == 400
117
    assert "error" in error
118
119
120


@pytest.mark.parametrize("model_name", [MODEL_NAME])
121
def test_batch_classification_empty_list(server: RemoteOpenAIServer, model_name: str):
122
123
    classification_response = requests.post(
        server.url_for("classify"),
124
        json={"model": model_name, "input": []},
125
126
    )
    classification_response.raise_for_status()
127
    output = ClassificationResponse.model_validate(classification_response.json())
128
129
130
131

    assert output.object == "list"
    assert isinstance(output.data, list)
    assert len(output.data) == 0
132
133
134
135
136
137


@pytest.mark.asyncio
async def test_invocations(server: RemoteOpenAIServer):
    request_args = {
        "model": MODEL_NAME,
138
        "input": "This product was excellent and exceeded my expectations",
139
140
    }

141
142
143
    classification_response = requests.post(
        server.url_for("classify"), json=request_args
    )
144
145
    classification_response.raise_for_status()

146
147
148
    invocation_response = requests.post(
        server.url_for("invocations"), json=request_args
    )
149
150
151
152
153
154
    invocation_response.raise_for_status()

    classification_output = classification_response.json()
    invocation_output = invocation_response.json()

    assert classification_output.keys() == invocation_output.keys()
155
    for classification_data, invocation_data in zip(
156
157
        classification_output["data"], invocation_output["data"]
    ):
158
159
        assert classification_data.keys() == invocation_data.keys()
        assert classification_data["probs"] == pytest.approx(
160
161
            invocation_data["probs"], rel=0.01
        )
162
163
164
165


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
166
async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
167
168
    input_text = ["This product was excellent and exceeded my expectations"]

169
    async def get_outputs(use_activation):
170
171
        response = requests.post(
            server.url_for("classify"),
172
173
174
175
176
            json={
                "model": model_name,
                "input": input_text,
                "use_activation": use_activation,
            },
177
        )
178
        outputs = response.json()
179
        return torch.tensor([x["probs"] for x in outputs["data"]])
180

181
182
183
    default = await get_outputs(use_activation=None)
    w_activation = await get_outputs(use_activation=True)
    wo_activation = await get_outputs(use_activation=False)
184

185
186
187
188
189
190
191
192
193
    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)."
    )
194
195
196
197


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
198
async def test_score(server: RemoteOpenAIServer, model_name: str):
199
200
201
202
203
204
205
206
207
208
209
210
211
212
    # score api is only enabled for num_labels == 1.
    response = requests.post(
        server.url_for("score"),
        json={
            "model": model_name,
            "text_1": "ping",
            "text_2": "pong",
        },
    )
    assert response.json()["error"]["type"] == "BadRequestError"


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
213
async def test_rerank(server: RemoteOpenAIServer, model_name: str):
214
215
216
217
218
219
220
221
222
223
    # 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()["error"]["type"] == "BadRequestError"
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
    input_text = "This product was excellent and exceeded my expectations"
    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):
    # token_classify uses ALL pooling, which does not support chunked prefill.
    task = "token_classify"
    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"Task {task} is not supported"
    )


@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": "test",
            "encoding_format": "float",
            "task": task,
        },
    )
    assert response.json()["error"]["type"] == "BadRequestError"
    assert response.json()["error"]["message"].startswith(
        f"Task {task} is not supported"
    )