support-max-pixels (#6)

- feat: configure max pixels (509c5332bacf540daad3c8fcead31c9abf2e2452)

README.md

@@ -24,8 +24,8 @@ images = [Image.open(path) for path in image_paths]
 # Example 1: Text matching task with single vector embeddings
 model.set_task(task='text-matching')
 
-# Generate embeddings with dimension truncation (256)
-img_embeddings = model.encode_images(images=images, truncate_dim=256)
+# Generate embeddings with dimension truncation (256), decrease max_pixels
+img_embeddings = model.encode_images(images=images, truncate_dim=256, max_pixels=602112)
 text_embeddings = model.encode_texts(texts=texts, truncate_dim=256, max_length=512)
 
 # Example 2: Retrieval task with multi-vector embeddings

modeling_jina_embeddings_v4.py

@@ -374,6 +374,21 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         truncate_dim: Optional[int] = None,
         text_type: Optional[str] = None,
     ) -> List[torch.Tensor]:
+        """
+        Encodes a list of texts into embeddings.
+
+        Args:
+            texts: List of text strings to encode
+            max_length: Maximum token length for text processing
+            batch_size: Number of texts 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
+            truncate_dim: Dimension to truncate embeddings to (128, 256, 512, or 1024)
+            text_type: Type of text being encoded ('query' or 'passage')
+
+        Returns:
+            List of text embeddings as tensors or numpy arrays
+        """
         text_type = text_type or "query"
         encode_kwargs = self._validate_encoding_params(
             vector_type, truncate_dim, text_type
@@ -404,7 +419,26 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         vector_type: Optional[str] = None,
         return_numpy: bool = False,
         truncate_dim: Optional[int] = None,
+        max_pixels: Optional[int] = None,
     ) -> List[torch.Tensor]:
+        """
+        Encodes a list of images into embeddings.
+
+        Args:
+            images: List of PIL images 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
+            truncate_dim: Dimension to truncate embeddings to (128, 256, 512, or 1024)
+            max_pixels: Maximum number of pixels to process per image
+
+        Returns:
+            List of image embeddings as tensors or numpy arrays
+        """
+        if max_pixels:
+            default_max_pixels = self.processor.image_processor.max_pixels
+            self.processor.image_processor.max_pixels = max_pixels  # change during encoding
+
         encode_kwargs = self._validate_encoding_params(vector_type, truncate_dim)
 
         is_single = len(images) == 1
@@ -417,6 +451,9 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             **encode_kwargs,
         )
 
+        if max_pixels:
+            self.processor.image_processor.max_pixels = default_max_pixels
+
         return embeddings[0] if is_single else embeddings
 
     @classmethod
@@ -426,6 +463,9 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         *args,
         **kwargs,
     ):
+        """
+        Loads a pretrained model and configures it with the appropriate task adapter (`retrieval` by default).
+        """
         if "torch_dtype" not in kwargs:
             kwargs["torch_dtype"] = "auto"