Finalize Sentence Transformers integration

1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+    "word_embedding_dimension": 4096,
+    "pooling_mode_cls_token": false,
+    "pooling_mode_mean_tokens": false,
+    "pooling_mode_max_tokens": false,
+    "pooling_mode_mean_sqrt_len_tokens": false,
+    "pooling_mode_weightedmean_tokens": false,
+    "pooling_mode_lasttoken": true,
+    "include_prompt": true
+  }

README.md

@@ -2,6 +2,7 @@
 tags:
 - mmeb
 - transformers
+- sentence-transformers
 language:
 - en
 - ar
@@ -34,15 +35,10 @@ Our model achieves SOTA performance on MMEB benchmark.
 
 ## Usage
 
-Below is an example we adapted from [VLM2Vec](https://huggingface.co/TIGER-Lab/VLM2Vec-Full).
+### Transformers
 
-First clone github
-```bash
-git clone https://github.com/haon-chen/mmE5.git
-pip install -r requirements.txt
-```
+Below is an example we adapted from [VLM2Vec](https://huggingface.co/TIGER-Lab/VLM2Vec-Full).
 
-Then you can enter the directory to run the following command.
 ```python
 import torch
 import requests
@@ -107,6 +103,52 @@ print(string, '=', compute_similarity(qry_output, tgt_output))
 ## <|image|><|begin_of_text|> Represent the given image. = tensor([[0.3887]], device='cuda:0', dtype=torch.bfloat16)
 ```
 
+### Sentence Transformers
+
+You can also use Sentence Transformers, where the majority of the pre- and post-processing has been abstracted.
+
+```python
+from sentence_transformers import SentenceTransformer
+import requests
+
+# Load the model
+model = SentenceTransformer("intfloat/mmE5-mllama-11b-instruct", trust_remote_code=True)
+
+# Download an example image of a cat and a dog
+dog_cat_image_bytes = requests.get('https://github.com/haon-chen/mmE5/blob/main/figures/example.jpg?raw=true', stream=True).raw.read()
+with open("cat_dog_example.jpg", "wb") as f:
+    f.write(dog_cat_image_bytes)
+
+# Image + Text -> Text
+image_embeddings = model.encode([{
+    "image": "cat_dog_example.jpg",
+    "text": "Represent the given image with the following question: What is in the image",
+}])
+text_embeddings = model.encode([
+    {"text": "A cat and a dog"},
+    {"text": "A cat and a tiger"},
+])
+
+similarity = model.similarity(image_embeddings, text_embeddings)
+print(similarity)
+# tensor([[0.3967, 0.3090]])
+# ✅ The first text is most similar to the image
+
+# Text -> Image
+image_embeddings = model.encode([
+    {"image": dog_cat_image_bytes, "text": "Represent the given image."},
+])
+text_embeddings = model.encode([
+    {"text": "Find me an everyday image that matches the given caption: A cat and a dog."},
+    {"text": "Find me an everyday image that matches the given caption: A cat and a tiger."},
+])
+
+similarity = model.similarity(image_embeddings, text_embeddings)
+print(similarity)
+# tensor([[0.4250, 0.3896]])
+# ✅ The first text is most similar to the image
+```
+
 ## Citation
 ```
 @article{chen2025mmE5,

custom_st.py

@@ -17,6 +17,7 @@ class MultiModalTransformer(BaseTransformer):
         super().__init__(model_name_or_path, **kwargs)
         if tokenizer_args is None:
             tokenizer_args = {}
+        tokenizer_args.pop("trust_remote_code", None)
 
         # Initialize processor
         self.processor = AutoProcessor.from_pretrained(
@@ -32,6 +33,7 @@ class MultiModalTransformer(BaseTransformer):
             is_peft_model: bool,
             **model_args,
     ) -> None:
+        model_args.pop("trust_remote_code", None)
         self.auto_model = MllamaForConditionalGeneration.from_pretrained(
             model_name_or_path, torch_dtype=torch.bfloat16, cache_dir=cache_dir, **model_args
         )
@@ -47,49 +49,38 @@ class MultiModalTransformer(BaseTransformer):
             **kwargs
         )
 
-        # Apply last pooling and normalization
-        last_hidden_state = outputs.hidden_states[-1]
-        attention_mask = features["attention_mask"]
-        sentence_embedding = self._last_pooling(last_hidden_state, attention_mask)
-
-        features.update({"sentence_embedding": sentence_embedding})
+        features.update({"token_embeddings": outputs.hidden_states[-1]})
         return features
 
-    def _last_pooling(self, last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-        """Apply last token pooling and L2 normalization"""
-        sequence_lengths = attention_mask.sum(dim=1) - 1
-        batch_size = last_hidden_state.shape[0]
-        reps = last_hidden_state[torch.arange(batch_size, device=last_hidden_state.device), sequence_lengths]
-        return torch.nn.functional.normalize(reps, p=2, dim=-1)
-
     def tokenize(self, texts: List[List[Dict]] | List[str]) -> Dict[str, torch.Tensor]:
         def process_text_item(item):
             if isinstance(item, str):
-                return item, []
+                return item, None
 
-            text, images = "", []
-            for sub_item in item:
-                if sub_item["type"] == "text":
-                    text += sub_item["content"]
-                elif sub_item["type"] in ["image_bytes", "image_path"]:
-                    text += "<|image|>"
-                    if sub_item["type"] == "image_bytes":
-                        img = Image.open(BytesIO(sub_item["content"])).convert("RGB")
-                    else:
-                        img = Image.open(sub_item["content"]).convert("RGB")
-                    images.append(img)
-                else:
-                    raise ValueError(f"Unknown data type {sub_item['type']}")
-            return text, images
+            text, img = "", None
+            if "image" in item:
+                text += "<|image|>"
+                img = item["image"]
+                if isinstance(img, bytes):
+                    img = Image.open(BytesIO(img)).convert("RGB")
+                elif isinstance(img, str):
+                    img = Image.open(img).convert("RGB")
+                elif not isinstance(img, Image):
+                    raise ValueError(f"Unknown image type {type(img)}")
+            if "text" in item:
+                if text:
+                    text += "<|begin_of_text|> "
+                text += item["text"].lstrip()
+            
+            return text, img
 
         all_texts, all_images = [], []
         for item in texts:
             text, images = process_text_item(item)
             all_texts.append(text)
-            all_images.extend(images)
+            all_images.append(images)
 
-        # Process inputs through the processor
-        if all_images:
+        if all_images != [None] * len(all_images):
             inputs = self.processor(
                 text=all_texts,
                 images=all_images,

modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "custom_st.MultiModalTransformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.models.Normalize"
+  }
+]