refactor: accept review

Signed-off-by: Sigrid Jin (Sionic AI) <[email protected]>

Author: Sigrid Jin (Sionic AI)

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

tests/models/multimodal/test_jina_embeddings_v4.py renamed to tests/models/pooling/test_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,

vllm/model_executor/models/jina_embeddings_v4.py

@@ -31,16 +31,16 @@
 # THE POSSIBILITY OF SUCH DAMAGE.
 #####
 
-import os
-import string
-import sys
-import threading
 ##
 # Python bindings for the NVML library
 ##
 from ctypes import *
 from ctypes.util import find_library
 from functools import wraps
+import sys
+import os
+import threading
+import string
 
 ## C Type mappings ##
 ## Enums