refactor: accept review

Sigrid Jin (Sionic AI) · Sigrid Jin (Sionic AI) · commit a635648569cc · 2025-07-18T00:26:40.000+09:00
Signed-off-by: Sigrid Jin (Sionic AI) &lt;sigrid@sionic.ai&gt;
diff --git a/benchmarks/jina_embeddings_v4_validation.py b/benchmarks/jina_embeddings_v4_validation.py
@@ -155,9 +155,9 @@ def compute_vllm_embeddings(
 
     for output in outputs:
         # Extract based on token type
-        if 151652 in output.prompt_token_ids:  # VISION_START_TOKEN_ID
-            img_start = output.prompt_token_ids.index(151652)
-            img_end = output.prompt_token_ids.index(151653)
+        if VISION_START_TOKEN_ID in output.prompt_token_ids:
+            img_start = output.prompt_token_ids.index(VISION_START_TOKEN_ID)
+            img_end = output.prompt_token_ids.index(VISION_END_TOKEN_ID)
             embedding_data = output.outputs.data[img_start : img_end + 1]
         else:
             embedding_data = output.outputs.data
diff --git a/tests/models/pooling/test_jina_embeddings_v4.py b/tests/models/pooling/test_jina_embeddings_v4.py
diff --git a/vllm/model_executor/models/jina_embeddings_v4.py b/vllm/model_executor/models/jina_embeddings_v4.py
@@ -35,8 +35,6 @@
 
 PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]
 
-# Triton kernel for optimized vision token extraction
-
 
 @MULTIMODAL_REGISTRY.register_processor(Qwen2VLMultiModalProcessor,
                                         info=Qwen2VLProcessingInfo,
