main
py 99 lines 3.19 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import Any
4
5 import openai
6 import pytest
7 import pytest_asyncio
8
9 from tests.utils import RemoteOpenAIServer
10
11 MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
12 max_model_len = 128
13
14 input = """Immerse yourself in the enchanting chronicle of calculus, a
15 mathematical domain that has radically transformed our comprehension of
16 change and motion. Despite its roots in ancient civilizations, the
17 formal birth of calculus predominantly occurred in the 17th century,
18 primarily under the influential guidance of Sir Isaac Newton and Gottfried
19 Wilhelm Leibniz. The earliest traces of calculus concepts are found in
20 ancient Greek mathematics,most notably in the works of Eudoxus and
21 Archimedes, around 300 BCE. They utilized the 'method of exhaustion'—a
22 technique for computing areas and volumes through the use of finite sums.
23 This methodology laid crucial foundational work for integral calculus.
24 In the 17th century, both Newton and Leibniz independently pioneered
25 calculus, each contributing unique perspectives that would shape this new
26 field."""
27
28
29 @pytest.fixture(scope="module")
30 def server():
31 args = [
32 "--runner",
33 "pooling",
34 "--dtype",
35 "bfloat16",
36 "--enforce-eager",
37 "--max-model-len",
38 str(max_model_len),
39 ]
40
41 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
42 yield remote_server
43
44
45 @pytest_asyncio.fixture
46 async def client(server):
47 async with server.get_async_client() as async_client:
48 yield async_client
49
50
51 @pytest.mark.asyncio
52 async def test_smaller_truncation_size(client: openai.AsyncOpenAI):
53 truncation_size = 10
54 kwargs: dict[str, Any] = {
55 "model": MODEL_NAME,
56 "input": input,
57 "truncate_prompt_tokens": truncation_size,
58 }
59
60 response = await client.post(path="embeddings", cast_to=object, body={**kwargs})
61
62 assert response["usage"]["prompt_tokens"] == truncation_size
63
64
65 @pytest.mark.asyncio
66 async def test_bigger_truncation_size(client: openai.AsyncOpenAI):
67 truncation_size = max_model_len + 1
68 kwargs: dict[str, Any] = {
69 "model": MODEL_NAME,
70 "input": input,
71 "truncate_prompt_tokens": truncation_size,
72 }
73
74 with pytest.raises(openai.BadRequestError) as err:
75 await client.post(path="embeddings", cast_to=object, body={**kwargs})
76
77 assert err.value.status_code == 400
78 error_details = err.value.response.json()["error"]
79 assert error_details["type"] == "BadRequestError"
80 expected_message = (
81 "truncate_prompt_tokens value is "
82 "greater than max_model_len."
83 " Please request a smaller truncation size."
84 )
85 assert error_details["message"] == expected_message
86
87
88 @pytest.mark.asyncio
89 async def test_max_truncation_size(client: openai.AsyncOpenAI):
90 truncation_size = -1
91 kwargs: dict[str, Any] = {
92 "model": MODEL_NAME,
93 "input": input,
94 "truncate_prompt_tokens": truncation_size,
95 }
96
97 response = await client.post(path="embeddings", cast_to=object, body={**kwargs})
98
99 assert response["usage"]["prompt_tokens"] == max_model_len