test_classification.py 7.74 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
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from vllm.entrypoints.openai.protocol import ClassificationResponse

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
166
167
168
169


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

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

    default = await get_outputs(activation=None)
    w_activation = await get_outputs(activation=True)
    wo_activation = await get_outputs(activation=False)

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


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_pooling(server: RemoteOpenAIServer, model_name: str):
    # pooling api uses ALL pooling, which does not support chunked prefill.
    response = requests.post(
        server.url_for("pooling"),
198
        json={"model": model_name, "input": "test", "encoding_format": "float"},
199
200
    )
    assert response.json()["error"]["type"] == "BadRequestError"
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
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,
            "text_1": "ping",
            "text_2": "pong",
        },
    )
    assert response.json()["error"]["type"] == "BadRequestError"


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
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()["error"]["type"] == "BadRequestError"