| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test cases for long text embedding with automatic chunking mechanism. |
| 4 | |
| 5 | This test suite validates vLLM's automatic chunking functionality for handling |
| 6 | text inputs that exceed the model's maximum token length, specifically targeting |
| 7 | the intfloat/multilingual-e5-small model (max token length: 512). |
| 8 | """ |
| 9 | |
| 10 | import random |
| 11 | |
| 12 | import openai |
| 13 | import pytest |
| 14 | import pytest_asyncio |
| 15 | |
| 16 | from tests.utils import RemoteOpenAIServer |
| 17 | from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse |
| 18 | |
| 19 | |
| 20 | def _generate_random_text(word_count: int) -> str: |
| 21 | """Generate deterministic text with approximately the specified word count.""" |
| 22 | rng = random.Random(word_count) |
| 23 | |
| 24 | # Common English words with focus on verbs and nouns for realistic text |
| 25 | common_words = [ |
| 26 | # Essential articles and pronouns (minimal) |
| 27 | "the", |
| 28 | "and", |
| 29 | "you", |
| 30 | "they", |
| 31 | "this", |
| 32 | "that", |
| 33 | "these", |
| 34 | "those", |
| 35 | # Action verbs |
| 36 | "create", |
| 37 | "build", |
| 38 | "develop", |
| 39 | "design", |
| 40 | "implement", |
| 41 | "execute", |
| 42 | "analyze", |
| 43 | "process", |
| 44 | "generate", |
| 45 | "calculate", |
| 46 | "evaluate", |
| 47 | "optimize", |
| 48 | "transform", |
| 49 | "integrate", |
| 50 | "configure", |
| 51 | "deploy", |
| 52 | "monitor", |
| 53 | "manage", |
| 54 | "discover", |
| 55 | "explore", |
| 56 | "investigate", |
| 57 | "research", |
| 58 | "study", |
| 59 | "examine", |
| 60 | "improve", |
| 61 | "enhance", |
| 62 | "upgrade", |
| 63 | "modify", |
| 64 | "update", |
| 65 | "maintain", |
| 66 | "solve", |
| 67 | "resolve", |
| 68 | "handle", |
| 69 | "address", |
| 70 | "tackle", |
| 71 | "overcome", |
| 72 | "communicate", |
| 73 | "collaborate", |
| 74 | "coordinate", |
| 75 | "organize", |
| 76 | "plan", |
| 77 | "achieve", |
| 78 | "accomplish", |
| 79 | "complete", |
| 80 | "finish", |
| 81 | "deliver", |
| 82 | "provide", |
| 83 | # Technology and science nouns |
| 84 | "system", |
| 85 | "application", |
| 86 | "software", |
| 87 | "hardware", |
| 88 | "network", |
| 89 | "database", |
| 90 | "algorithm", |
| 91 | "model", |
| 92 | "framework", |
| 93 | "platform", |
| 94 | "interface", |
| 95 | "protocol", |
| 96 | "architecture", |
| 97 | "infrastructure", |
| 98 | "component", |
| 99 | "module", |
| 100 | "service", |
| 101 | "technology", |
| 102 | "innovation", |
| 103 | "solution", |
| 104 | "methodology", |
| 105 | "approach", |
| 106 | "artificial", |
| 107 | "intelligence", |
| 108 | "machine", |
| 109 | "learning", |
| 110 | "neural", |
| 111 | "network", |
| 112 | "computer", |
| 113 | "processor", |
| 114 | "memory", |
| 115 | "storage", |
| 116 | "computation", |
| 117 | "data", |
| 118 | "information", |
| 119 | "knowledge", |
| 120 | "insight", |
| 121 | "pattern", |
| 122 | "trend", |
| 123 | "analysis", |
| 124 | "research", |
| 125 | "development", |
| 126 | "engineering", |
| 127 | "science", |
| 128 | "mathematics", |
| 129 | "statistics", |
| 130 | "probability", |
| 131 | "optimization", |
| 132 | "performance", |
| 133 | "efficiency", |
| 134 | # General nouns |
| 135 | "project", |
| 136 | "team", |
| 137 | "organization", |
| 138 | "company", |
| 139 | "business", |
| 140 | "industry", |
| 141 | "market", |
| 142 | "customer", |
| 143 | "user", |
| 144 | "client", |
| 145 | "product", |
| 146 | "feature", |
| 147 | "function", |
| 148 | "requirement", |
| 149 | "specification", |
| 150 | "documentation", |
| 151 | "report", |
| 152 | "result", |
| 153 | "outcome", |
| 154 | "impact", |
| 155 | "benefit", |
| 156 | "advantage", |
| 157 | "challenge", |
| 158 | "problem", |
| 159 | "opportunity", |
| 160 | "strategy", |
| 161 | "goal", |
| 162 | "objective", |
| 163 | "target", |
| 164 | "milestone", |
| 165 | "process", |
| 166 | "procedure", |
| 167 | "workflow", |
| 168 | "pipeline", |
| 169 | "operation", |
| 170 | "task", |
| 171 | "activity", |
| 172 | "event", |
| 173 | "session", |
| 174 | "meeting", |
| 175 | "discussion", |
| 176 | "decision", |
| 177 | ] |
| 178 | |
| 179 | words = [] |
| 180 | for _ in range(word_count): |
| 181 | words.append(rng.choice(common_words)) |
| 182 | |
| 183 | # Add some punctuation for more realistic text |
| 184 | text = " ".join(words) |
| 185 | # Add periods every 10-20 words |
| 186 | words_list = text.split() |
| 187 | result = [] |
| 188 | for i, word in enumerate(words_list): |
| 189 | result.append(word) |
| 190 | if (i + 1) % rng.randint(10, 20) == 0 and i < len(words_list) - 1: |
| 191 | result[-1] += "." |
| 192 | |
| 193 | return " ".join(result) |
| 194 | |
| 195 | |
| 196 | MODEL_NAME = "intfloat/multilingual-e5-small" |
| 197 | DTYPE = "bfloat16" |
| 198 | |
| 199 | # Test text: Generate text with approximately 1500 words to exceed 1024 tokens |
| 200 | LONG_TEXT_1500_WORDS = _generate_random_text(1500) |
| 201 | |
| 202 | # Test text: Generate text with approximately 2500 words to exceed 2048 tokens |
| 203 | LONG_TEXT_2500_WORDS = _generate_random_text(2500) |
| 204 | |
| 205 | |
| 206 | @pytest.fixture(scope="module") |
| 207 | def server_with_chunked_processing(): |
| 208 | """Start server with automatic chunking processing enabled.""" |
| 209 | args = [ |
| 210 | "--runner", |
| 211 | "pooling", |
| 212 | "--dtype", |
| 213 | DTYPE, |
| 214 | "--enforce-eager", |
| 215 | "--max-model-len", |
| 216 | "512", # Set smaller max_model_len to trigger chunking mechanism |
| 217 | "--pooler-config", |
| 218 | ( |
| 219 | '{"pooling_type": "MEAN", "use_activation": true, ' |
| 220 | '"enable_chunked_processing": true, "max_embed_len": 10000}' |
| 221 | ), |
| 222 | "--gpu-memory-utilization", |
| 223 | "0.8", |
| 224 | ] |
| 225 | |
| 226 | with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: |
| 227 | yield remote_server |
| 228 | |
| 229 | |
| 230 | @pytest_asyncio.fixture |
| 231 | async def client_with_chunked_processing(server_with_chunked_processing): |
| 232 | """Create async client with chunking processing support.""" |
| 233 | async with server_with_chunked_processing.get_async_client() as async_client: |
| 234 | yield async_client |
| 235 | |
| 236 | |
| 237 | @pytest.mark.asyncio |
| 238 | @pytest.mark.parametrize("model_name", [MODEL_NAME]) |
| 239 | async def test_long_text_embedding_1500_chars( |
| 240 | client_with_chunked_processing: openai.AsyncOpenAI, model_name: str |
| 241 | ): |
| 242 | """Test embedding processing for ~1500 character long text |
| 243 | (~1028 tokens, exceeding 512 token limit).""" |
| 244 | # Verify text length |
| 245 | # Verify text has sufficient word count (approximately 1500 words) |
| 246 | word_count = len(LONG_TEXT_1500_WORDS.split()) |
| 247 | assert word_count >= 1400, f"Test text word count insufficient: {word_count} words" |
| 248 | |
| 249 | # Send embedding request |
| 250 | embedding_response = await client_with_chunked_processing.embeddings.create( |
| 251 | model=model_name, |
| 252 | input=[LONG_TEXT_1500_WORDS], |
| 253 | encoding_format="float", |
| 254 | ) |
| 255 | |
| 256 | # Verify response structure |
| 257 | embeddings = EmbeddingResponse.model_validate( |
| 258 | embedding_response.model_dump(mode="json") |
| 259 | ) |
| 260 | |
| 261 | assert embeddings.id is not None |
| 262 | assert len(embeddings.data) == 1 |
| 263 | assert ( |
| 264 | len(embeddings.data[0].embedding) == 384 |
| 265 | ) # multilingual-e5-small embedding dimension |
| 266 | assert embeddings.usage.completion_tokens == 0 |
| 267 | # Due to chunked processing, token count should |
| 268 | # reflect actual processed tokens |
| 269 | # With ~1500 words, we expect roughly |
| 270 | # 1024+ tokens (exceeding 512 token limit) |
| 271 | # Should exceed single chunk limit of 512 |
| 272 | assert embeddings.usage.prompt_tokens > 800 |
| 273 | assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens |
| 274 | |
| 275 | # Verify embedding vector validity |
| 276 | embedding_vector = embeddings.data[0].embedding |
| 277 | assert isinstance(embedding_vector, list) |
| 278 | assert all(isinstance(x, float) for x in embedding_vector), ( |
| 279 | "Embedding vector should contain floats" |
| 280 | ) |
| 281 | assert not all(x == 0 for x in embedding_vector), ( |
| 282 | "Embedding vector should not be all zeros" |
| 283 | ) |
| 284 | |
| 285 | |
| 286 | @pytest.mark.asyncio |
| 287 | @pytest.mark.parametrize("model_name", [MODEL_NAME]) |
| 288 | async def test_long_text_embedding_2500_chars( |
| 289 | client_with_chunked_processing: openai.AsyncOpenAI, model_name: str |
| 290 | ): |
| 291 | """Test embedding processing for ~2500 character long text |
| 292 | (~2048 tokens, requiring multiple chunks).""" |
| 293 | # Verify text length |
| 294 | # Verify text has sufficient word count (approximately 2500 words) |
| 295 | word_count = len(LONG_TEXT_2500_WORDS.split()) |
| 296 | assert word_count >= 2300, f"Test text word count insufficient: {word_count} words" |
| 297 | |
| 298 | # Send embedding request |
| 299 | embedding_response = await client_with_chunked_processing.embeddings.create( |
| 300 | model=model_name, |
| 301 | input=[LONG_TEXT_2500_WORDS], |
| 302 | encoding_format="float", |
| 303 | ) |
| 304 | |
| 305 | # Verify response structure |
| 306 | embeddings = EmbeddingResponse.model_validate( |
| 307 | embedding_response.model_dump(mode="json") |
| 308 | ) |
| 309 | |
| 310 | assert embeddings.id is not None |
| 311 | assert len(embeddings.data) == 1 |
| 312 | assert ( |
| 313 | len(embeddings.data[0].embedding) == 384 |
| 314 | ) # multilingual-e5-small embedding dimension |
| 315 | assert embeddings.usage.completion_tokens == 0 |
| 316 | # Due to chunked processing, token count should |
| 317 | # reflect actual processed tokens |
| 318 | # With ~2500 words, we expect |
| 319 | # roughly 2048+ tokens (requiring multiple chunks) |
| 320 | # Should require multiple chunks for processing |
| 321 | assert embeddings.usage.prompt_tokens > 1500 |
| 322 | assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens |
| 323 | |
| 324 | # Verify embedding vector validity |
| 325 | embedding_vector = embeddings.data[0].embedding |
| 326 | assert isinstance(embedding_vector, list) |
| 327 | assert all(isinstance(x, float) for x in embedding_vector), ( |
| 328 | "Embedding vector should contain floats" |
| 329 | ) |
| 330 | assert not all(x == 0 for x in embedding_vector), ( |
| 331 | "Embedding vector should not be all zeros" |
| 332 | ) |
| 333 | |
| 334 | |
| 335 | @pytest.mark.asyncio |
| 336 | @pytest.mark.parametrize("model_name", [MODEL_NAME]) |
| 337 | async def test_batch_long_text_embedding( |
| 338 | client_with_chunked_processing: openai.AsyncOpenAI, model_name: str |
| 339 | ): |
| 340 | """Test batch long text embedding processing.""" |
| 341 | input_texts = [ |
| 342 | LONG_TEXT_1500_WORDS, |
| 343 | LONG_TEXT_2500_WORDS, |
| 344 | "This is a short text test.", # Short text for comparison |
| 345 | ] |
| 346 | |
| 347 | # Send batch embedding request |
| 348 | embedding_response = await client_with_chunked_processing.embeddings.create( |
| 349 | model=model_name, |
| 350 | input=input_texts, |
| 351 | encoding_format="float", |
| 352 | ) |
| 353 | |
| 354 | # Verify response structure |
| 355 | embeddings = EmbeddingResponse.model_validate( |
| 356 | embedding_response.model_dump(mode="json") |
| 357 | ) |
| 358 | |
| 359 | assert embeddings.id is not None |
| 360 | assert len(embeddings.data) == 3 # Three input texts |
| 361 | |
| 362 | # Verify each embedding dimension |
| 363 | for i, embedding_data in enumerate(embeddings.data): |
| 364 | assert len(embedding_data.embedding) == 384 |
| 365 | assert embedding_data.index == i |
| 366 | |
| 367 | # Verify embedding vector validity |
| 368 | embedding_vector = embedding_data.embedding |
| 369 | assert all(isinstance(x, float) for x in embedding_vector) |
| 370 | assert not all(x == 0 for x in embedding_vector) |
| 371 | |
| 372 | # Verify token usage |
| 373 | assert embeddings.usage.completion_tokens == 0 |
| 374 | # Total token count should be very substantial |
| 375 | assert embeddings.usage.prompt_tokens > 1000 |
| 376 | assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens |
| 377 | |
| 378 | |
| 379 | @pytest.mark.asyncio |
| 380 | @pytest.mark.parametrize("model_name", [MODEL_NAME]) |
| 381 | async def test_chunked_vs_normal_consistency( |
| 382 | client_with_chunked_processing: openai.AsyncOpenAI, model_name: str |
| 383 | ): |
| 384 | """Test consistency between chunked and |
| 385 | normal processing (using short text).""" |
| 386 | # Use a short text within the 512 token limit |
| 387 | short_text = ( |
| 388 | "Artificial intelligence technology is changing our world, " |
| 389 | "bringing unprecedented opportunities and challenges." |
| 390 | ) |
| 391 | |
| 392 | # Send embedding request |
| 393 | embedding_response = await client_with_chunked_processing.embeddings.create( |
| 394 | model=model_name, |
| 395 | input=[short_text], |
| 396 | encoding_format="float", |
| 397 | ) |
| 398 | |
| 399 | # Verify response structure |
| 400 | embeddings = EmbeddingResponse.model_validate( |
| 401 | embedding_response.model_dump(mode="json") |
| 402 | ) |
| 403 | |
| 404 | assert embeddings.id is not None |
| 405 | assert len(embeddings.data) == 1 |
| 406 | assert len(embeddings.data[0].embedding) == 384 |
| 407 | assert embeddings.usage.completion_tokens == 0 |
| 408 | # Short text should not require chunked processing |
| 409 | assert embeddings.usage.prompt_tokens < 512 |
| 410 | assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens |
| 411 | |
| 412 | # 验证embedding向量的有效性 |
| 413 | embedding_vector = embeddings.data[0].embedding |
| 414 | assert isinstance(embedding_vector, list) |
| 415 | assert all(isinstance(x, float) for x in embedding_vector) |
| 416 | assert not all(x == 0 for x in embedding_vector) |
| 417 | |
| 418 | |
| 419 | @pytest.mark.asyncio |
| 420 | @pytest.mark.parametrize("model_name", [MODEL_NAME]) |
| 421 | async def test_chunked_processing_response_format( |
| 422 | client_with_chunked_processing: openai.AsyncOpenAI, model_name: str |
| 423 | ): |
| 424 | """Test response format and structure during chunked processing.""" |
| 425 | # Test with long text to trigger chunking |
| 426 | embedding_response = await client_with_chunked_processing.embeddings.create( |
| 427 | model=model_name, |
| 428 | input=[LONG_TEXT_1500_WORDS], |
| 429 | encoding_format="float", |
| 430 | ) |
| 431 | |
| 432 | # Verify response structure |
| 433 | embeddings = EmbeddingResponse.model_validate( |
| 434 | embedding_response.model_dump(mode="json") |
| 435 | ) |
| 436 | |
| 437 | assert embeddings.id is not None |
| 438 | assert len(embeddings.data) == 1 |
| 439 | assert embeddings.data[0].object == "embedding" |
| 440 | assert embeddings.data[0].index == 0 |
| 441 | |
| 442 | # Verify embedding vector properties |
| 443 | embedding_vector = embeddings.data[0].embedding |
| 444 | assert isinstance(embedding_vector, list) |
| 445 | import math |
| 446 | |
| 447 | vector_norm = math.sqrt(sum(x * x for x in embedding_vector)) |
| 448 | # Check that the vector is normalized |
| 449 | # (default behavior for most embedding models) |
| 450 | assert 0.8 < vector_norm < 1.2, ( |
| 451 | f"Vector norm should be reasonable, actual: {vector_norm}" |
| 452 | ) |