test_embedding.py 13.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
6
import base64

import numpy as np
7
8
import openai
import pytest
9
import pytest_asyncio
10
import requests
11
12
import torch
import torch.nn.functional as F
13

14
from vllm.entrypoints.openai.protocol import EmbeddingResponse
15
from vllm.transformers_utils.tokenizer import get_tokenizer
16

17
18
from ...models.language.pooling.embed_utils import (
    run_embedding_correctness_test)
19
from ...models.utils import check_embeddings_close
20
from ...utils import RemoteOpenAIServer
21

22
MODEL_NAME = "intfloat/multilingual-e5-small"
23
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}"""  # noqa: E501
24
DTYPE = "bfloat16"
25
26
27


@pytest.fixture(scope="module")
28
def server():
29
    args = [
30
31
        "--runner",
        "pooling",
32
33
        # use half precision for speed and memory savings in CI environment
        "--dtype",
34
        DTYPE,
35
36
        "--enforce-eager",
        "--max-model-len",
37
        "512",
38
39
        "--chat-template",
        DUMMY_CHAT_TEMPLATE,
40
41
    ]

42
    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
43
        yield remote_server
44
45


46
@pytest_asyncio.fixture
47
48
async def client(server):
    async with server.get_async_client() as async_client:
49
        yield async_client
50
51


52
53
54
55
56
57
58
@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


59
@pytest.mark.asyncio
60
@pytest.mark.parametrize("model_name", [MODEL_NAME])
61
62
async def test_single_embedding(hf_model, client: openai.AsyncOpenAI,
                                model_name: str):
63
64
65
66
67
    input_texts = [
        "The chef prepared a delicious meal.",
    ]

    # test single embedding
68
    embedding_response = await client.embeddings.create(
69
70
71
72
        model=model_name,
        input=input_texts,
        encoding_format="float",
    )
73
74
75
    embeddings = EmbeddingResponse.model_validate(
        embedding_response.model_dump(mode="json"))

76
77
    assert embeddings.id is not None
    assert len(embeddings.data) == 1
78
    assert len(embeddings.data[0].embedding) == 384
79
    assert embeddings.usage.completion_tokens == 0
80
81
    assert embeddings.usage.prompt_tokens == 11
    assert embeddings.usage.total_tokens == 11
82

83
    vllm_outputs = [d.embedding for d in embeddings.data]
84
    run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
85

86
87
    # test using token IDs
    input_tokens = [1, 1, 1, 1, 1]
88
    embedding_response = await client.embeddings.create(
89
90
91
92
        model=model_name,
        input=input_tokens,
        encoding_format="float",
    )
93
94
95
    embeddings = EmbeddingResponse.model_validate(
        embedding_response.model_dump(mode="json"))

96
97
    assert embeddings.id is not None
    assert len(embeddings.data) == 1
98
    assert len(embeddings.data[0].embedding) == 384
99
100
101
102
103
104
    assert embeddings.usage.completion_tokens == 0
    assert embeddings.usage.prompt_tokens == 5
    assert embeddings.usage.total_tokens == 5


@pytest.mark.asyncio
105
@pytest.mark.parametrize("model_name", [MODEL_NAME])
106
107
async def test_batch_embedding(hf_model, client: openai.AsyncOpenAI,
                               model_name: str):
108
    # test list[str]
109
110
111
112
    input_texts = [
        "The cat sat on the mat.", "A feline was resting on a rug.",
        "Stars twinkle brightly in the night sky."
    ]
113
    embedding_response = await client.embeddings.create(
114
115
116
117
        model=model_name,
        input=input_texts,
        encoding_format="float",
    )
118
119
120
    embeddings = EmbeddingResponse.model_validate(
        embedding_response.model_dump(mode="json"))

121
122
    assert embeddings.id is not None
    assert len(embeddings.data) == 3
123
    assert len(embeddings.data[0].embedding) == 384
124
    assert embeddings.usage.completion_tokens == 0
125
126
    assert embeddings.usage.prompt_tokens == 33
    assert embeddings.usage.total_tokens == 33
127

128
    vllm_outputs = [d.embedding for d in embeddings.data]
129
    run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
130

131
    # test list[list[int]]
132
133
    input_tokens = [[4, 5, 7, 9, 20], [15, 29, 499], [24, 24, 24, 24, 24],
                    [25, 32, 64, 77]]
134
    embedding_response = await client.embeddings.create(
135
136
137
138
        model=model_name,
        input=input_tokens,
        encoding_format="float",
    )
139
140
141
    embeddings = EmbeddingResponse.model_validate(
        embedding_response.model_dump(mode="json"))

142
143
    assert embeddings.id is not None
    assert len(embeddings.data) == 4
144
    assert len(embeddings.data[0].embedding) == 384
145
146
147
    assert embeddings.usage.completion_tokens == 0
    assert embeddings.usage.prompt_tokens == 17
    assert embeddings.usage.total_tokens == 17
148
149
150


@pytest.mark.asyncio
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_conversation_embedding(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.",
    }]

166
167
168
169
170
171
172
173
    chat_response = requests.post(
        server.url_for("v1/embeddings"),
        json={
            "model": model_name,
            "messages": messages,
            "encoding_format": "float",
        },
    )
174
    chat_response.raise_for_status()
175
    chat_embeddings = EmbeddingResponse.model_validate(chat_response.json())
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191

    tokenizer = get_tokenizer(tokenizer_name=model_name, tokenizer_mode="fast")
    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},
    )
192
193
    completion_embeddings = EmbeddingResponse.model_validate(
        completion_response.model_dump(mode="json"))
194

195
196
197
198
199
200
    assert chat_embeddings.id is not None
    assert completion_embeddings.id is not None
    assert chat_embeddings.created <= completion_embeddings.created
    assert chat_embeddings.model_dump(
        exclude={"id", "created"}) == (completion_embeddings.model_dump(
            exclude={"id", "created"}))
201
202
203
204


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
205
async def test_batch_base64_embedding(hf_model, client: openai.AsyncOpenAI,
206
207
208
209
210
211
                                      model_name: str):
    input_texts = [
        "Hello my name is",
        "The best thing about vLLM is that it supports many different models"
    ]

212
213
214
    responses_float = await client.embeddings.create(input=input_texts,
                                                     model=model_name,
                                                     encoding_format="float")
215
    float_data = [d.embedding for d in responses_float.data]
216
    run_embedding_correctness_test(hf_model, input_texts, float_data)
217

218
219
220
    responses_base64 = await client.embeddings.create(input=input_texts,
                                                      model=model_name,
                                                      encoding_format="base64")
221
    base64_data = []
222
    for data in responses_base64.data:
223
        base64_data.append(
224
            np.frombuffer(base64.b64decode(data.embedding),
225
                          dtype="float32").tolist())
226

227
    run_embedding_correctness_test(hf_model, input_texts, base64_data)
228
229

    # Default response is float32 decoded from base64 by OpenAI Client
230
231
    responses_default = await client.embeddings.create(input=input_texts,
                                                       model=model_name)
232
    default_data = [d.embedding for d in responses_default.data]
233
    run_embedding_correctness_test(hf_model, input_texts, default_data)
234
235
236


@pytest.mark.asyncio
237
238
239
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_single_embedding_truncation(client: openai.AsyncOpenAI,
                                           model_name: str):
240
241
242
243
244
    input_texts = [
        "Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
    ]

    # test single embedding
245
    embedding_response = await client.embeddings.create(
246
247
248
        model=model_name,
        input=input_texts,
        extra_body={"truncate_prompt_tokens": 10})
249
250
251
    embeddings = EmbeddingResponse.model_validate(
        embedding_response.model_dump(mode="json"))

252
253
    assert embeddings.id is not None
    assert len(embeddings.data) == 1
254
    assert len(embeddings.data[0].embedding) == 384
255
256
257
258
259
260
261
262
    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
    ]
263
    embedding_response = await client.embeddings.create(
264
265
266
        model=model_name,
        input=input_tokens,
        extra_body={"truncate_prompt_tokens": 10})
267
268
    embeddings = EmbeddingResponse.model_validate(
        embedding_response.model_dump(mode="json"))
269
270
271

    assert embeddings.id is not None
    assert len(embeddings.data) == 1
272
    assert len(embeddings.data[0].embedding) == 384
273
274
275
276
277
278
    assert embeddings.usage.completion_tokens == 0
    assert embeddings.usage.prompt_tokens == 10
    assert embeddings.usage.total_tokens == 10


@pytest.mark.asyncio
279
280
281
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_single_embedding_truncation_invalid(client: openai.AsyncOpenAI,
                                                   model_name: str):
282
283
284
285
286
    input_texts = [
        "Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
    ]

    with pytest.raises(openai.BadRequestError):
287
        response = await client.embeddings.create(
288
289
290
            model=model_name,
            input=input_texts,
            extra_body={"truncate_prompt_tokens": 8193})
291
        assert "error" in response.object
292
        assert "truncate_prompt_tokens value is greater than max_model_len. "\
293
               "Please, select a smaller truncation size." in response.message
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318


@pytest.mark.asyncio
async def test_invocations(server: RemoteOpenAIServer,
                           client: openai.AsyncOpenAI):
    input_texts = [
        "The chef prepared a delicious meal.",
    ]

    request_args = {
        "model": MODEL_NAME,
        "input": input_texts,
        "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()
319
320
321
322
323
324
325
    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")
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358


@pytest.mark.asyncio
async def test_invocations_conversation(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()
359
360
361
362
363
364
365
    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")
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397


@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_normalize(server: RemoteOpenAIServer, model_name: str):
    input_text = ["The chef prepared a delicious meal."]

    async def get_outputs(normalize):
        request_args = {
            "model": MODEL_NAME,
            "input": input_text,
            "encoding_format": "float",
            "normalize": normalize
        }

        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(normalize=None)
    w_normal = await get_outputs(normalize=True)
    wo_normal = await 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)."