Integrate with Sentence Transformers v5.4

config_sentence_transformers.json

@@ -4,10 +4,7 @@
       "transformers": "4.46.2",
       "pytorch": "2.2.2"
     },
-    "prompts":{
-      "image": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is shown in this image?<|im_end|>\n<|endoftext|>",
-      "query": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Query: %s<|im_end|>\n<|endoftext|>"
-    },
+    "prompts": {},
     "default_prompt_name": null,
     "similarity_fn_name": "cosine"
   }

custom_st.py

@@ -42,22 +42,23 @@ class Transformer(nn.Module):
         self.max_pixels = max_pixels
         self.min_pixels = min_pixels
         self.max_seq_length = max_seq_length
-        
+
         # Handle args
         model_kwargs = model_args or {}
-        model_kwargs.update(kwargs)
-
         processor_kwargs = processor_args or {}
+
+        if cache_dir is not None:
+            model_kwargs['cache_dir'] = cache_dir
+            processor_kwargs['cache_dir'] = cache_dir
+
         processor_kwargs.update({
             'min_pixels': min_pixels,
             'max_pixels': max_pixels,
-            'cache_dir': cache_dir
         })
         
         # Initialize model
         self.model = Qwen2VLForConditionalGeneration.from_pretrained(
             model_name_or_path,
-            cache_dir=cache_dir,
             **model_kwargs
         ).eval()
 
@@ -271,7 +272,7 @@ class Transformer(nn.Module):
         )
         return features
 
-    def tokenize(self, texts: List[Union[str, Image.Image, bytes]], padding: str = 'longest') -> Dict[str, torch.Tensor]:
+    def tokenize(self, texts: List[Union[str, Image.Image, bytes]], padding: str = 'longest', **kwargs) -> Dict[str, torch.Tensor]:
         processed_texts, processed_images = self._process_input(texts)
         
         return self.processor(
@@ -311,4 +312,8 @@ class Transformer(nn.Module):
         else:
             config = {'model_name_or_path': input_path}
 
-        return Transformer(**config)
+        return Transformer(**config)
+    
+    @property
+    def modalities(self) -> List[str]:
+        return ['text', 'image']