| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Example script demonstrating long text embedding with chunked processing in vLLM. |
| 5 | |
| 6 | This example shows how to use vLLM's chunked processing feature to handle text |
| 7 | inputs that exceed the model's maximum token length. The feature automatically |
| 8 | splits long text into chunks and handles different pooling types optimally. |
| 9 | |
| 10 | Prerequisites: |
| 11 | 1. Start vLLM server with chunked processing enabled: |
| 12 | |
| 13 | # MEAN pooling (processes all chunks, recommended for complete coverage) |
| 14 | vllm serve intfloat/multilingual-e5-large \ |
| 15 | --pooler-config \ |
| 16 | '{"pooling_type": "MEAN", "use_activation": true, ' \ |
| 17 | '"enable_chunked_processing": true, "max_embed_len": 3072000}' \ |
| 18 | --served-model-name multilingual-e5-large \ |
| 19 | --trust-remote-code \ |
| 20 | --port 31090 \ |
| 21 | --api-key your-api-key |
| 22 | |
| 23 | # OR CLS pooling (native CLS within chunks, MEAN aggregation across chunks) |
| 24 | vllm serve BAAI/bge-large-en-v1.5 \ |
| 25 | --pooler-config \ |
| 26 | '{"pooling_type": "CLS", "use_activation": true, ' \ |
| 27 | '"enable_chunked_processing": true, "max_embed_len": 1048576}' \ |
| 28 | --served-model-name bge-large-en-v1.5 \ |
| 29 | --trust-remote-code \ |
| 30 | --port 31090 \ |
| 31 | --api-key your-api-key |
| 32 | |
| 33 | 2. Install required dependencies: |
| 34 | pip install openai requests |
| 35 | """ |
| 36 | |
| 37 | import time |
| 38 | |
| 39 | import numpy as np |
| 40 | from openai import OpenAI |
| 41 | |
| 42 | # Configuration |
| 43 | API_KEY = "your-api-key" # Replace with your actual API key |
| 44 | BASE_URL = "http://localhost:31090/v1" |
| 45 | MODEL_NAME = "multilingual-e5-large" |
| 46 | |
| 47 | |
| 48 | def generate_long_text(base_text: str, repeat_count: int) -> str: |
| 49 | """Generate long text by repeating base text.""" |
| 50 | return base_text * repeat_count |
| 51 | |
| 52 | |
| 53 | def test_embedding_with_different_lengths(): |
| 54 | """Test embedding generation with different text lengths.""" |
| 55 | client = OpenAI(api_key=API_KEY, base_url=BASE_URL) |
| 56 | |
| 57 | # Test cases with different text lengths |
| 58 | test_cases = [ |
| 59 | { |
| 60 | "name": "Short Text", |
| 61 | "text": "Hello, this is a short text for embedding.", |
| 62 | "expected_chunks": 1, |
| 63 | }, |
| 64 | { |
| 65 | "name": "Medium Text", |
| 66 | "text": generate_long_text( |
| 67 | "This is a medium-length text that should fit within the " |
| 68 | "model's context window. " * 20, |
| 69 | 2, |
| 70 | ), |
| 71 | "expected_chunks": 1, |
| 72 | }, |
| 73 | { |
| 74 | "name": "Long Text (2 chunks)", |
| 75 | "text": generate_long_text( |
| 76 | "This is a very long text that will exceed the model's " |
| 77 | "maximum context length and trigger chunked processing. " * 50, |
| 78 | 5, |
| 79 | ), |
| 80 | "expected_chunks": 2, |
| 81 | }, |
| 82 | { |
| 83 | "name": "Very Long Text (3+ chunks)", |
| 84 | "text": generate_long_text( |
| 85 | "This text is extremely long and will definitely " |
| 86 | "require multiple chunks for processing. " * 100, |
| 87 | 10, |
| 88 | ), |
| 89 | "expected_chunks": 3, |
| 90 | }, |
| 91 | ] |
| 92 | |
| 93 | print("🧪 Testing vLLM Long Text Embedding with Chunked Processing") |
| 94 | print("=" * 70) |
| 95 | |
| 96 | for i, test_case in enumerate(test_cases, 1): |
| 97 | print(f"\n📝 Test {i}: {test_case['name']}") |
| 98 | print(f"Text length: {len(test_case['text'])} characters") |
| 99 | |
| 100 | try: |
| 101 | start_time = time.time() |
| 102 | |
| 103 | response = client.embeddings.create( |
| 104 | input=test_case["text"], model=MODEL_NAME, encoding_format="float" |
| 105 | ) |
| 106 | |
| 107 | end_time = time.time() |
| 108 | processing_time = end_time - start_time |
| 109 | |
| 110 | # Extract embedding data |
| 111 | embedding = response.data[0].embedding |
| 112 | embedding_dim = len(embedding) |
| 113 | |
| 114 | print("✅ Success!") |
| 115 | print(f" - Embedding dimension: {embedding_dim}") |
| 116 | print(f" - Processing time: {processing_time:.2f}s") |
| 117 | print(f" - Expected chunks: ~{test_case['expected_chunks']}") |
| 118 | print(f" - First 5 values: {embedding[:5]}") |
| 119 | |
| 120 | except Exception as e: |
| 121 | print(f"❌ Failed: {str(e)}") |
| 122 | |
| 123 | |
| 124 | def test_batch_embedding(): |
| 125 | """Test batch embedding with mixed-length inputs.""" |
| 126 | client = OpenAI(api_key=API_KEY, base_url=BASE_URL) |
| 127 | |
| 128 | print("\n🔄 Testing Batch Embedding with Mixed Lengths") |
| 129 | print("=" * 50) |
| 130 | |
| 131 | # Mix of short and long texts |
| 132 | batch_inputs = [ |
| 133 | "Short text 1", |
| 134 | generate_long_text("Medium length text that fits in one chunk. " * 20, 1), |
| 135 | "Another short text", |
| 136 | generate_long_text("Long text requiring chunked processing. " * 100, 5), |
| 137 | ] |
| 138 | |
| 139 | try: |
| 140 | start_time = time.time() |
| 141 | |
| 142 | response = client.embeddings.create( |
| 143 | input=batch_inputs, model=MODEL_NAME, encoding_format="float" |
| 144 | ) |
| 145 | |
| 146 | end_time = time.time() |
| 147 | processing_time = end_time - start_time |
| 148 | |
| 149 | print("✅ Batch processing successful!") |
| 150 | print(f" - Number of inputs: {len(batch_inputs)}") |
| 151 | print(f" - Number of embeddings: {len(response.data)}") |
| 152 | print(f" - Total processing time: {processing_time:.2f}s") |
| 153 | print( |
| 154 | f" - Average time per input: {processing_time / len(batch_inputs):.2f}s" |
| 155 | ) |
| 156 | |
| 157 | for i, data in enumerate(response.data): |
| 158 | input_length = len(batch_inputs[i]) |
| 159 | embedding_dim = len(data.embedding) |
| 160 | print( |
| 161 | f" - Input {i + 1}: {input_length} chars → {embedding_dim}D embedding" |
| 162 | ) |
| 163 | |
| 164 | except Exception as e: |
| 165 | print(f"❌ Batch processing failed: {str(e)}") |
| 166 | |
| 167 | |
| 168 | def test_multiple_long_texts_batch(): |
| 169 | """Test batch processing with multiple long texts to verify chunk ID uniqueness.""" |
| 170 | client = OpenAI(api_key=API_KEY, base_url=BASE_URL) |
| 171 | |
| 172 | print("\n🔧 Testing Multiple Long Texts in Batch (Chunk ID Fix Verification)") |
| 173 | print("=" * 70) |
| 174 | |
| 175 | # Create multiple distinct long texts that will all require chunking |
| 176 | # Note: All pooling types now use MEAN aggregation across chunks: |
| 177 | # - Native pooling (MEAN/CLS/LAST) is used within each chunk |
| 178 | # - MEAN aggregation combines results across all chunks |
| 179 | # - Full semantic coverage for all pooling types |
| 180 | long_texts = [ |
| 181 | generate_long_text( |
| 182 | "First long document about artificial intelligence and machine learning. " |
| 183 | * 80, |
| 184 | 6, |
| 185 | ), |
| 186 | generate_long_text( |
| 187 | "Second long document about natural language processing and transformers. " |
| 188 | * 80, |
| 189 | 6, |
| 190 | ), |
| 191 | generate_long_text( |
| 192 | "Third long document about computer vision and neural networks. " * 80, 6 |
| 193 | ), |
| 194 | ] |
| 195 | |
| 196 | # Add some short texts to mix things up |
| 197 | batch_inputs = [ |
| 198 | "Short text before long texts", |
| 199 | long_texts[0], |
| 200 | "Short text between long texts", |
| 201 | long_texts[1], |
| 202 | long_texts[2], |
| 203 | "Short text after long texts", |
| 204 | ] |
| 205 | |
| 206 | print("📊 Batch composition:") |
| 207 | for i, text in enumerate(batch_inputs): |
| 208 | length = len(text) |
| 209 | text_type = "Long (will be chunked)" if length > 5000 else "Short" |
| 210 | print(f" - Input {i + 1}: {length} chars ({text_type})") |
| 211 | |
| 212 | try: |
| 213 | start_time = time.time() |
| 214 | |
| 215 | response = client.embeddings.create( |
| 216 | input=batch_inputs, model=MODEL_NAME, encoding_format="float" |
| 217 | ) |
| 218 | |
| 219 | end_time = time.time() |
| 220 | processing_time = end_time - start_time |
| 221 | |
| 222 | print("\n✅ Multiple long texts batch processing successful!") |
| 223 | print(f" - Number of inputs: {len(batch_inputs)}") |
| 224 | print(f" - Number of embeddings returned: {len(response.data)}") |
| 225 | print(f" - Total processing time: {processing_time:.2f}s") |
| 226 | |
| 227 | # Verify each embedding is different (no incorrect aggregation) |
| 228 | embeddings = [data.embedding for data in response.data] |
| 229 | |
| 230 | if len(embeddings) >= 3: |
| 231 | import numpy as np |
| 232 | |
| 233 | # Compare embeddings of the long texts (indices 1, 3, 4) |
| 234 | long_embeddings = [ |
| 235 | np.array(embeddings[1]), # First long text |
| 236 | np.array(embeddings[3]), # Second long text |
| 237 | np.array(embeddings[4]), # Third long text |
| 238 | ] |
| 239 | |
| 240 | print("\n🔍 Verifying embedding uniqueness:") |
| 241 | for i in range(len(long_embeddings)): |
| 242 | for j in range(i + 1, len(long_embeddings)): |
| 243 | cosine_sim = np.dot(long_embeddings[i], long_embeddings[j]) / ( |
| 244 | np.linalg.norm(long_embeddings[i]) |
| 245 | * np.linalg.norm(long_embeddings[j]) |
| 246 | ) |
| 247 | print( |
| 248 | f" - Similarity between long text {i + 1} and {j + 1}: " |
| 249 | f"{cosine_sim:.4f}" |
| 250 | ) |
| 251 | |
| 252 | if ( |
| 253 | cosine_sim < 0.9 |
| 254 | ): # Different content should have lower similarity |
| 255 | print(" ✅ Good: Embeddings are appropriately different") |
| 256 | else: |
| 257 | print( |
| 258 | " ⚠️ High similarity - may indicate chunk " |
| 259 | "aggregation issue" |
| 260 | ) |
| 261 | |
| 262 | print("\n📋 Per-input results:") |
| 263 | for i, data in enumerate(response.data): |
| 264 | input_length = len(batch_inputs[i]) |
| 265 | embedding_dim = len(data.embedding) |
| 266 | embedding_norm = np.linalg.norm(data.embedding) |
| 267 | print( |
| 268 | f" - Input {i + 1}: {input_length} chars → {embedding_dim}D " |
| 269 | f"embedding (norm: {embedding_norm:.4f})" |
| 270 | ) |
| 271 | |
| 272 | print( |
| 273 | "\n✅ This test verifies the fix for chunk ID collisions in " |
| 274 | "batch processing" |
| 275 | ) |
| 276 | print(" - Before fix: Multiple long texts would have conflicting chunk IDs") |
| 277 | print(" - After fix: Each prompt's chunks have unique IDs with prompt index") |
| 278 | |
| 279 | except Exception as e: |
| 280 | print(f"❌ Multiple long texts batch test failed: {str(e)}") |
| 281 | print(" This might indicate the chunk ID collision bug is present!") |
| 282 | |
| 283 | |
| 284 | def test_embedding_consistency(): |
| 285 | """Test that chunked processing produces consistent results.""" |
| 286 | client = OpenAI(api_key=API_KEY, base_url=BASE_URL) |
| 287 | |
| 288 | print("\n🔍 Testing Embedding Consistency") |
| 289 | print("=" * 40) |
| 290 | |
| 291 | # Use the same long text multiple times |
| 292 | long_text = generate_long_text( |
| 293 | "Consistency test text for chunked processing validation. " * 50, 3 |
| 294 | ) |
| 295 | |
| 296 | embeddings = [] |
| 297 | |
| 298 | try: |
| 299 | for i in range(3): |
| 300 | response = client.embeddings.create( |
| 301 | input=long_text, model=MODEL_NAME, encoding_format="float" |
| 302 | ) |
| 303 | embeddings.append(response.data[0].embedding) |
| 304 | print(f" - Generated embedding {i + 1}") |
| 305 | |
| 306 | # Check consistency (embeddings should be identical) |
| 307 | if len(embeddings) >= 2: |
| 308 | # Calculate similarity between first two embeddings |
| 309 | |
| 310 | emb1 = np.array(embeddings[0]) |
| 311 | emb2 = np.array(embeddings[1]) |
| 312 | |
| 313 | # Cosine similarity |
| 314 | cosine_sim = np.dot(emb1, emb2) / ( |
| 315 | np.linalg.norm(emb1) * np.linalg.norm(emb2) |
| 316 | ) |
| 317 | |
| 318 | print("✅ Consistency test completed!") |
| 319 | print(f" - Cosine similarity between runs: {cosine_sim:.6f}") |
| 320 | print(" - Expected: ~1.0 (identical embeddings)") |
| 321 | |
| 322 | if cosine_sim > 0.999: |
| 323 | print(" - ✅ High consistency achieved!") |
| 324 | else: |
| 325 | print(" - ⚠️ Consistency may vary due to numerical precision") |
| 326 | |
| 327 | except Exception as e: |
| 328 | print(f"❌ Consistency test failed: {str(e)}") |
| 329 | |
| 330 | |
| 331 | def main(): |
| 332 | """Main function to run all tests.""" |
| 333 | print("🚀 vLLM Long Text Embedding Client") |
| 334 | print(f"📡 Connecting to: {BASE_URL}") |
| 335 | print(f"🤖 Model: {MODEL_NAME}") |
| 336 | masked_key = "*" * (len(API_KEY) - 4) + API_KEY[-4:] if len(API_KEY) > 4 else "****" |
| 337 | print(f"🔑 API Key: {masked_key}") |
| 338 | |
| 339 | # Run all test cases |
| 340 | test_embedding_with_different_lengths() |
| 341 | test_batch_embedding() |
| 342 | test_multiple_long_texts_batch() |
| 343 | test_embedding_consistency() |
| 344 | |
| 345 | print("\n" + "=" * 70) |
| 346 | print("🎉 All tests completed!") |
| 347 | print("\n💡 Key Features Demonstrated:") |
| 348 | print(" - ✅ Automatic chunked processing for long text") |
| 349 | print(" - ✅ Seamless handling of mixed-length batches") |
| 350 | print(" - ✅ Multiple long texts in single batch (chunk ID fix)") |
| 351 | print(" - ✅ Unified chunked processing:") |
| 352 | print(" • Native pooling used within each chunk") |
| 353 | print(" • MEAN aggregation across all chunks") |
| 354 | print(" • Complete semantic coverage for all pooling types") |
| 355 | print(" - ✅ Consistent embedding generation") |
| 356 | print(" - ✅ Backward compatibility with short text") |
| 357 | print("\n📚 For more information, see:") |
| 358 | print( |
| 359 | " - Documentation: https://docs.vllm.ai/en/latest/models/pooling_models.html" |
| 360 | ) |
| 361 | print(" - Chunked Processing Guide: openai_embedding_long_text.md") |
| 362 | |
| 363 | |
| 364 | if __name__ == "__main__": |
| 365 | main() |