refactor-image-processing (#16)

- refactor: support urls, fast processor, flash attn check (9624180b21e596eb410896719e1c1708aaed343c)
- refactor: image loading in st wrapper (9ef2e43d97b27bc27da6b71bc68d6160d317da20)

custom_st.py

@@ -1,32 +1,34 @@
+from io import BytesIO
+from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, Union
 
+import requests
 import torch
 from PIL import Image
 from torch import nn
-from transformers import AutoConfig, AutoProcessor, AutoModel
+from transformers import AutoConfig, AutoModel, AutoProcessor
 
 
 class Transformer(nn.Module):
 
     save_in_root: bool = True
-    
+
     def __init__(
         self,
-        model_name_or_path: str = 'jinaai/jina-embeddings-v4',
+        model_name_or_path: str = "jinaai/jina-embeddings-v4",
         max_seq_length: Optional[int] = None,
         config_args: Optional[Dict[str, Any]] = None,
         model_args: Optional[Dict[str, Any]] = None,
         tokenizer_args: Optional[Dict[str, Any]] = None,
         cache_dir: Optional[str] = None,
-        backend: Literal['torch', 'onnx', 'openvino'] = 'torch',
+        backend: Literal["torch", "onnx", "openvino"] = "torch",
         **kwargs,
     ) -> None:
         super(Transformer, self).__init__()
-        if backend != 'torch':
+        if backend != "torch":
             raise ValueError(
-                f'Backend \'{backend}\' is not supported, please use \'torch\' instead'
+                f"Backend '{backend}' is not supported, please use 'torch' instead"
             )
-
         config_kwargs = config_args or {}
         model_kwargs = model_args or {}
         tokenizer_kwargs = tokenizer_args or {}
@@ -34,9 +36,11 @@ class Transformer(nn.Module):
         self.config = AutoConfig.from_pretrained(
             model_name_or_path, cache_dir=cache_dir, **config_kwargs
         )
-        self.default_task = model_args.pop('default_task', None)
+        self.default_task = model_args.pop("default_task", None)
         if self.default_task and self.default_task not in self.config.task_names:
-            raise ValueError(f"Invalid task: {self.default_task}. Must be one of {self.config.task_names}.")
+            raise ValueError(
+                f"Invalid task: {self.default_task}. Must be one of {self.config.task_names}."
+            )
 
         self.model = AutoModel.from_pretrained(
             model_name_or_path, config=self.config, cache_dir=cache_dir, **model_kwargs
@@ -45,6 +49,7 @@ class Transformer(nn.Module):
         self.processor = AutoProcessor.from_pretrained(
             model_name_or_path,
             cache_dir=cache_dir,
+            use_fast=True,
             **tokenizer_kwargs,
         )
         self.max_seq_length = max_seq_length or 8192
@@ -55,33 +60,52 @@ class Transformer(nn.Module):
         encoding = {}
         text_indices = []
         image_indices = []
-        
         for i, text in enumerate(texts):
             if isinstance(text, str):
-                text_indices.append(i)
+                # Remove Query: or Passage: prefixes when checking for URLs or file paths
+                clean_text = text
+                if text.startswith("Query: "):
+                    clean_text = text[len("Query: ") :]
+                elif text.startswith("Passage: "):
+                    clean_text = text[len("Passage: ") :]
+
+                if clean_text.startswith("http"):
+                    response = requests.get(clean_text)
+                    texts[i] = Image.open(BytesIO(response.content)).convert("RGB")
+                    image_indices.append(i)
+                elif Path(clean_text).is_file():
+                    try:
+                        texts[i] = Image.open(clean_text).convert("RGB")
+                        image_indices.append(i)
+                    except Exception as e:
+                        text_indices.append(i)
+                else:
+                    text_indices.append(i)
             elif isinstance(text, Image.Image):
                 image_indices.append(i)
             else:
-                raise ValueError(f'Invalid input type: {type(text)}')
-        
+                raise ValueError(f"Invalid input type: {type(text)}")
         if text_indices:
             _texts = [texts[i] for i in text_indices]
-            text_features = self.processor.process_texts(_texts, max_length=self.max_seq_length)
+            text_features = self.processor.process_texts(
+                _texts, max_length=self.max_seq_length
+            )
             for key, value in text_features.items():
-                encoding[f'text_{key}'] = value
-            encoding['text_indices'] = text_indices
-        
+                encoding[f"text_{key}"] = value
+            encoding["text_indices"] = text_indices
+
         if image_indices:
             _images = [texts[i] for i in image_indices]
             img_features = self.processor.process_images(_images)
             for key, value in img_features.items():
-                encoding[f'image_{key}'] = value
-            encoding['image_indices'] = image_indices
-            
+                encoding[f"image_{key}"] = value
+            encoding["image_indices"] = image_indices
+
         return encoding
-    
 
-    def forward(self, features: Dict[str, torch.Tensor], task: Optional[str] = None) -> Dict[str, torch.Tensor]:
+    def forward(
+        self, features: Dict[str, torch.Tensor], task: Optional[str] = None
+    ) -> Dict[str, torch.Tensor]:
         self.model.eval()
 
         if task is None:
@@ -94,41 +118,55 @@ class Transformer(nn.Module):
             task = self.default_task
         else:
             if task not in self.config.task_names:
-                raise ValueError(f"Invalid task: {task}. Must be one of {self.config.task_names}.")
+                raise ValueError(
+                    f"Invalid task: {task}. Must be one of {self.config.task_names}."
+                )
 
         device = self.model.device.type
         all_embeddings = []
-        
+
         with torch.no_grad():
-            if any(k.startswith('text_') for k in features.keys()):
-                text_batch = {k[len('text_'):]: v.to(device) for k, v in features.items() if k.startswith('text_') and k != 'text_indices'}
-                text_indices = features.get('text_indices', [])
-                
+            if any(k.startswith("text_") for k in features.keys()):
+                text_batch = {
+                    k[len("text_") :]: v.to(device)
+                    for k, v in features.items()
+                    if k.startswith("text_") and k != "text_indices"
+                }
+                text_indices = features.get("text_indices", [])
+
                 with torch.autocast(device_type=device):
-                    text_embeddings = self.model(**text_batch, task_label=task).single_vec_emb
+                    text_embeddings = self.model(
+                        **text_batch, task_label=task
+                    ).single_vec_emb
                     if self.config.truncate_dim:
-                        text_embeddings = text_embeddings[:, :self.config.truncate_dim]
-                
+                        text_embeddings = text_embeddings[:, : self.config.truncate_dim]
+
                 for i, embedding in enumerate(text_embeddings):
                     all_embeddings.append((text_indices[i], embedding))
-            
-            if any(k.startswith('image_') for k in features.keys()):
-                image_batch = {k[len('image_'):]: v.to(device) for k, v in features.items() if k.startswith('image_') and k != 'image_indices'}
-                image_indices = features.get('image_indices', [])
-                
+
+            if any(k.startswith("image_") for k in features.keys()):
+                image_batch = {
+                    k[len("image_") :]: v.to(device)
+                    for k, v in features.items()
+                    if k.startswith("image_") and k != "image_indices"
+                }
+                image_indices = features.get("image_indices", [])
+
                 with torch.autocast(device_type=device):
-                    img_embeddings = self.model(**image_batch, task_label=task).single_vec_emb
+                    img_embeddings = self.model(
+                        **image_batch, task_label=task
+                    ).single_vec_emb
                     if self.config.truncate_dim:
-                        img_embeddings = img_embeddings[:, :self.config.truncate_dim]
-                
+                        img_embeddings = img_embeddings[:, : self.config.truncate_dim]
+
                 for i, embedding in enumerate(img_embeddings):
                     all_embeddings.append((image_indices[i], embedding))
 
         if not all_embeddings:
-            raise RuntimeError('No embeddings were generated')
+            raise RuntimeError("No embeddings were generated")
 
         all_embeddings.sort(key=lambda x: x[0])  # sort by original index
         combined_embeddings = torch.stack([emb for _, emb in all_embeddings])
-        features['sentence_embedding'] = combined_embeddings
-        
+        features["sentence_embedding"] = combined_embeddings
+
         return features

modeling_jina_embeddings_v4.py

@@ -5,20 +5,24 @@ import os
 from dataclasses import dataclass
 from enum import Enum
 from functools import partial
+from io import BytesIO
 from typing import Any, Callable, ClassVar, Dict, List, Optional, Union, cast
 
 import numpy as np
+import requests
 import torch
 from huggingface_hub import snapshot_download
-from peft import PeftModel, LoraConfig
+from peft import LoraConfig, PeftModel
 from PIL import Image
 from torch import nn
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 from transformers import BatchFeature
-from .qwen2_5_vl import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
+from transformers.utils import is_flash_attn_2_available
+
 from .configuration_jina_embeddings_v4 import JinaEmbeddingsV4Config
 from .custom_lora_module import MultiAdapterLinear
+from .qwen2_5_vl import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
 
 
 class PromptType(str, Enum):
@@ -140,7 +144,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         self._init_projection_layers(config)
         self.post_init()
         self.processor = JinaEmbeddingsV4Processor.from_pretrained(
-            self.name_or_path, trust_remote_code=True
+            self.name_or_path, trust_remote_code=True, use_fast=True
         )
         self.single_vector_projector_dim = config.single_vector_projector_dim
         self.multi_vector_projector_dim = config.multi_vector_projector_dim
@@ -160,7 +164,9 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             task (str): The task name. Must be one of ['retrieval', 'text-matching', 'code']
         """
         if task not in self.config.task_names:
-            raise ValueError(f"Invalid task: {task}. Must be one of {self.config.task_names}.")
+            raise ValueError(
+                f"Invalid task: {task}. Must be one of {self.config.task_names}."
+            )
         self._task = task
 
     def get_last_hidden_states(
@@ -342,7 +348,9 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         for batch in tqdm(dataloader, desc=desc):
             with torch.no_grad():
                 batch = {k: v.to(self.device) for k, v in batch.items()}
-                with torch.autocast(device_type=torch.device(self.device).type, dtype=torch.bfloat16):
+                with torch.autocast(
+                    device_type=torch.device(self.device).type, dtype=torch.bfloat16
+                ):
                     embeddings = self(**batch, task_label=task_label)
                     if vector_type == "single_vector":
                         embeddings = embeddings.single_vec_emb
@@ -395,7 +403,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             encode_kwargs["truncate_dim"] = truncate_dim
 
         return encode_kwargs
-    
+
     def _validate_task(self, task: Optional[str] = None) -> str:
         if task is None:
             if self.task is None:
@@ -406,7 +414,9 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             task = self.task
         else:
             if task not in self.config.task_names:
-                raise ValueError(f"Invalid task: {task}. Must be one of {self.config.task_names}.")
+                raise ValueError(
+                    f"Invalid task: {task}. Must be one of {self.config.task_names}."
+                )
         return task
 
     def encode_texts(
@@ -460,9 +470,23 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
 
         return embeddings
 
+    def _load_images_if_needed(
+        self, images: List[Union[str, Image.Image]]
+    ) -> List[Image.Image]:
+        loaded_images = []
+        for image in images:
+            if isinstance(image, str):
+                if image.startswith("http"):
+                    response = requests.get(image)
+                    image = Image.open(BytesIO(response.content)).convert("RGB")
+                else:
+                    image = Image.open(image).convert("RGB")
+            loaded_images.append(image)
+        return loaded_images
+
     def encode_images(
         self,
-        images: List[Image.Image],
+        images: List[Union[str, Image.Image]],
         task: Optional[str] = None,
         batch_size: int = 8,
         vector_type: Optional[str] = None,
@@ -474,7 +498,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         Encodes a list of images into embeddings.
 
         Args:
-            images: List of PIL images to encode
+            images: List of PIL images, URLs, or local file paths to encode
             batch_size: Number of images to process at once
             vector_type: Type of embedding vector to generate ('single_vector' or 'multi_vector')
             return_numpy: Whether to return numpy arrays instead of torch tensors
@@ -489,9 +513,9 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             self.processor.image_processor.max_pixels = (
                 max_pixels  # change during encoding
             )
-
         encode_kwargs = self._validate_encoding_params(vector_type, truncate_dim)
         task = self._validate_task(task)
+        images = self._load_images_if_needed(images)
         embeddings = self._process_batches(
             data=images,
             processor_fn=self.processor.process_images,
@@ -519,8 +543,10 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         """
         if "torch_dtype" not in kwargs:
             kwargs["torch_dtype"] = "auto"
-                    
+
         kwargs["key_mapping"] = super()._checkpoint_conversion_mapping
+        if not is_flash_attn_2_available():
+            kwargs["attn_implementation"] = "sdpa"
 
         base_model = super().from_pretrained(
             pretrained_model_name_or_path, *args, **kwargs
@@ -547,19 +573,19 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             model_id=adapter_dir,
             config=lora_config,
         )
-        
+
         @property
         def task(self):
             return self.model.task
-        
+
         @task.setter
         def task(self, value):
             self.model.task = value
-        
+
         peft_model.task = property(task.fget, task.fset)
         peft_model.__class__.task = property(
             lambda self: self.model.task,
-            lambda self, value: setattr(self.model, 'task', value)
+            lambda self, value: setattr(self.model, "task", value),
         )
 
         return peft_model

tokenizer_config.json

@@ -202,7 +202,7 @@
   "extra_special_tokens": {},
   "model_max_length": 131072,
   "pad_token": "<|endoftext|>",
-  "processor_class": "ColQwen25DuoProcessor",
+  "processor_class": "JinaEmbeddingsV4Processor",
   "split_special_tokens": false,
   "tokenizer_class": "Qwen2Tokenizer",
   "unk_token": null