feat: use enum for the vector type

modeling_jina_embeddings_v4.py

@@ -30,8 +30,12 @@ class PromptType(str, Enum):
     passage = "passage"
 
 
+class VectorType(str, Enum):
+    single = "single"
+    multi = "multi"
+
+
 PREFIX_DICT = {"query": "Query", "passage": "Passage"}
-VECTOR_TYPES = ["single", "multi"]
 
 
 class JinaEmbeddingsV4Processor(Qwen2_5_VLProcessor):
@@ -320,7 +324,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         task_label: Union[str, List[str]],
         processor_fn: Callable,
         desc: str,
-        vector_type: str = "single",
+        vector_type: Union[str, VectorType] = VectorType.single,
         return_numpy: bool = False,
         batch_size: int = 32,
         truncate_dim: Optional[int] = None,
@@ -340,7 +344,8 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
                     device_type=torch.device(self.device).type, dtype=torch.bfloat16
                 ):
                     embeddings = self(**batch, task_label=task_label)
-                    if vector_type == "single":
+                    vector_type_str = vector_type.value if isinstance(vector_type, VectorType) else vector_type
+                    if vector_type_str == VectorType.single.value:
                         embeddings = embeddings.single_vec_emb
                         if truncate_dim is not None:
                             embeddings = embeddings[:, :truncate_dim]
@@ -357,7 +362,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
 
     def _validate_encoding_params(
         self,
-        vector_type: Optional[str] = None,
+        vector_type: Optional[Union[str, VectorType]] = None,
         truncate_dim: Optional[int] = None,
         prompt_name: Optional[str] = None,
     ) -> Dict[str, Any]:
@@ -374,13 +379,17 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
                     else PREFIX_DICT["query"]
                 )
 
-        vector_type = vector_type or "single"
-        if vector_type not in VECTOR_TYPES:
-            raise ValueError(
-                f"Invalid vector_type: {vector_type}. Must be one of {VECTOR_TYPES}."
-            )
+        vector_type = vector_type or VectorType.single
+        if isinstance(vector_type, VectorType):
+            encode_kwargs["vector_type"] = vector_type.value
         else:
-            encode_kwargs["vector_type"] = vector_type
+            try:
+                vector_type_enum = VectorType(vector_type)
+                encode_kwargs["vector_type"] = vector_type_enum.value
+            except ValueError:
+                raise ValueError(
+                    f"Invalid vector_type: {vector_type}. Must be one of {[v.value for v in VectorType]}."
+                )
 
         truncate_dim = truncate_dim or self.config.truncate_dim
         if truncate_dim is not None and truncate_dim not in self.config.matryoshka_dims:
@@ -413,7 +422,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         task: Optional[str] = None,
         max_length: int = 8192,
         batch_size: int = 8,
-        vector_type: Optional[str] = None,
+        vector_type: Optional[Union[str, VectorType]] = None,
         return_numpy: bool = False,
         truncate_dim: Optional[int] = None,
         prompt_name: Optional[str] = None,
@@ -425,7 +434,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             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' or 'multi')
+            vector_type: Type of embedding vector to generate (VectorType.single or VectorType.multi)
             return_numpy: Whether to return numpy arrays instead of torch tensors
             truncate_dim: Dimension to truncate embeddings to (128, 256, 512, or 1024)
             prompt_name: Type of text being encoded ('query' or 'passage')
@@ -477,7 +486,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         images: List[Union[str, Image.Image]],
         task: Optional[str] = None,
         batch_size: int = 8,
-        vector_type: Optional[str] = None,
+        vector_type: Optional[Union[str, VectorType]] = None,
         return_numpy: bool = False,
         truncate_dim: Optional[int] = None,
         max_pixels: Optional[int] = None,
@@ -488,7 +497,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         Args:
             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' or 'multi')
+            vector_type: Type of embedding vector to generate (VectorType.single or VectorType.multi)
             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