A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 21:00 UTC
Fix entrypoint tests for embedding models (#14052)
Harry Mellor committed
Feb 28, 2025 at 16:56 UTC
4be4b26cb728f0746cb0f9d7900b8a7f2988a647
1 file changed
+11
-11
tests/entrypoints/openai/test_embedding.py
+11
-11
@@ -27,7 +27,7 @@ def server():
27
"bfloat16",
28
"--enforce-eager",
29
"--max-model-len",
30
- "8192",
30
+ "512",
31
"--chat-template",
32
DUMMY_CHAT_TEMPLATE,
33
]
@@ -60,10 +60,10 @@ async def test_single_embedding(client: openai.AsyncOpenAI, model_name: str):
60
61
assert embeddings.id is not None
62
assert len(embeddings.data) == 1
63
- assert len(embeddings.data[0].embedding) == 4096
63
+ assert len(embeddings.data[0].embedding) == 384
64
assert embeddings.usage.completion_tokens == 0
65
- assert embeddings.usage.prompt_tokens == 9
66
- assert embeddings.usage.total_tokens == 9
65
+ assert embeddings.usage.prompt_tokens == 11
66
+ assert embeddings.usage.total_tokens == 11
67
68
# test using token IDs
69
input_tokens = [1, 1, 1, 1, 1]
@@ -77,7 +77,7 @@ async def test_single_embedding(client: openai.AsyncOpenAI, model_name: str):
77
78
assert embeddings.id is not None
79
assert len(embeddings.data) == 1
80
- assert len(embeddings.data[0].embedding) == 4096
80
+ assert len(embeddings.data[0].embedding) == 384
81
assert embeddings.usage.completion_tokens == 0
82
assert embeddings.usage.prompt_tokens == 5
83
assert embeddings.usage.total_tokens == 5
@@ -101,10 +101,10 @@ async def test_batch_embedding(client: openai.AsyncOpenAI, model_name: str):
101
102
assert embeddings.id is not None
103
assert len(embeddings.data) == 3
104
- assert len(embeddings.data[0].embedding) == 4096
104
+ assert len(embeddings.data[0].embedding) == 384
105
assert embeddings.usage.completion_tokens == 0
106
- assert embeddings.usage.prompt_tokens == 32
107
- assert embeddings.usage.total_tokens == 32
106
+ assert embeddings.usage.prompt_tokens == 33
107
+ assert embeddings.usage.total_tokens == 33
108
109
# test List[List[int]]
110
input_tokens = [[4, 5, 7, 9, 20], [15, 29, 499], [24, 24, 24, 24, 24],
@@ -119,7 +119,7 @@ async def test_batch_embedding(client: openai.AsyncOpenAI, model_name: str):
119
120
assert embeddings.id is not None
121
assert len(embeddings.data) == 4
122
- assert len(embeddings.data[0].embedding) == 4096
122
+ assert len(embeddings.data[0].embedding) == 384
123
assert embeddings.usage.completion_tokens == 0
124
assert embeddings.usage.prompt_tokens == 17
125
assert embeddings.usage.total_tokens == 17
@@ -234,7 +234,7 @@ async def test_single_embedding_truncation(client: openai.AsyncOpenAI,
234
235
assert embeddings.id is not None
236
assert len(embeddings.data) == 1
237
- assert len(embeddings.data[0].embedding) == 4096
237
+ assert len(embeddings.data[0].embedding) == 384
238
assert embeddings.usage.completion_tokens == 0
239
assert embeddings.usage.prompt_tokens == 10
240
assert embeddings.usage.total_tokens == 10
@@ -252,7 +252,7 @@ async def test_single_embedding_truncation(client: openai.AsyncOpenAI,
252
253
assert embeddings.id is not None
254
assert len(embeddings.data) == 1
255
- assert len(embeddings.data[0].embedding) == 4096
255
+ assert len(embeddings.data[0].embedding) == 384
256
assert embeddings.usage.completion_tokens == 0
257
assert embeddings.usage.prompt_tokens == 10
258
assert embeddings.usage.total_tokens == 10