Integrate sentence transformers (#9)

- Base Integration with SentenceTransformers (2df56dc77ce0fbca90c338d38d2cffb4de4c9ea0)
- Update custom_st.py (f70ad8d07b0f89da5c767bc57d3453b8923938ac)
- Update README.md (27b4e411bf08eee10f4bd27941807530cce3099f)
- Update README.md (5b4def35a3fdbbe9c547cd2fce25896feb492d97)
- Update README.md (5ff08124217cda0a6a0a48ab93ada3aa2ac0a1c8)
- Update README.md (4f9008b5898196357c0cc9c767d767af1235cc32)


Co-authored-by: Solomatin Roman <Samoed@users.noreply.huggingface.co>

1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+    "word_embedding_dimension": 1536,
+    "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

@@ -3692,46 +3692,90 @@ The `GME` models support three types of input: **text**, **image**, and **image-
 |[`gme-Qwen2-VL-7B`](https://huggingface.co/Alibaba-NLP/gme-Qwen2-VL-7B-Instruct) | 8.29B | 32768 | 3584 | 67.48 | 69.73 | 67.44 |
 
 ## Usage 
+**Use with sentence_transformers**
+
+The `encode` function accept `str` or `dict` with key(s) in `{'text', 'image', 'prompt'}`.
+
+**Do not pass `prompt` as the argument to `encode`**, pass as the input as a `dict` with a `prompt` key.
+
+```python
+from sentence_transformers import SentenceTransformer
+
+
+t2i_prompt = 'Find an image that matches the given text.'
+texts = [
+    "The Tesla Cybertruck is a battery electric pickup truck built by Tesla, Inc. since 2023.",
+    "Alibaba office.",
+]
+images = [
+    'https://upload.wikimedia.org/wikipedia/commons/e/e9/Tesla_Cybertruck_damaged_window.jpg',
+    'https://upload.wikimedia.org/wikipedia/commons/e/e0/TaobaoCity_Alibaba_Xixi_Park.jpg',
+]
+
+
+gme_st = SentenceTransformer("Alibaba-NLP/gme-Qwen2-VL-2B-Instruct")
+
+# Single-modal embedding
+e_text = gme_st.encode(texts, convert_to_tensor=True)
+e_image = gme_st.encode([dict(image=i) for i in images], convert_to_tensor=True)
+print('Single-modal', (e_text @ e_image.T).tolist())
+## Single-modal [[0.356201171875, 0.06536865234375], [0.041717529296875, 0.37890625]]
+
+# How to set embedding instruction
+e_query = gme_st.encode([dict(text=t, prompt=t2i_prompt) for t in texts], convert_to_tensor=True)
+# If no prompt, we always use the default instruction.
+e_corpus = gme_st.encode([dict(image=i) for i in images], convert_to_tensor=True)
+print('Single-modal with instruction', (e_query @ e_corpus.T).tolist())
+## Single-modal with instruction [[0.425537109375, 0.1158447265625], [0.049835205078125, 0.413818359375]]
+
+# Fused-modal embedding
+e_fused = gme_st.encode([dict(text=t, image=i) for t, i in zip(texts, images)], convert_to_tensor=True)
+print('Fused-modal', (e_fused @ e_fused.T).tolist())
+## Fused-modal [[0.99951171875, 0.0556640625], [0.0556640625, 0.99951171875]]
+```
+
+
 **Use with custom code**
 
 ```python
 # You can find the script gme_inference.py in https://huggingface.co/Alibaba-NLP/gme-Qwen2-VL-2B-Instruct/blob/main/gme_inference.py
 from gme_inference import GmeQwen2VL
 
+t2i_prompt = 'Find an image that matches the given text.'
 texts = [
-    "What kind of car is this?",
-    "The Tesla Cybertruck is a battery electric pickup truck built by Tesla, Inc. since 2023."
+    "The Tesla Cybertruck is a battery electric pickup truck built by Tesla, Inc. since 2023.",
+    "Alibaba office.",
 ]
 images = [
-    'https://en.wikipedia.org/wiki/File:Tesla_Cybertruck_damaged_window.jpg',
-    'https://en.wikipedia.org/wiki/File:2024_Tesla_Cybertruck_Foundation_Series,_front_left_(Greenwich).jpg',
+    'https://upload.wikimedia.org/wikipedia/commons/e/e9/Tesla_Cybertruck_damaged_window.jpg',
+    'https://upload.wikimedia.org/wikipedia/commons/e/e0/TaobaoCity_Alibaba_Xixi_Park.jpg',
 ]
 
+
 gme = GmeQwen2VL("Alibaba-NLP/gme-Qwen2-VL-2B-Instruct")
 
 # Single-modal embedding
 e_text = gme.get_text_embeddings(texts=texts)
 e_image = gme.get_image_embeddings(images=images)
-print((e_text * e_image).sum(-1))
-## tensor([0.2281, 0.6001], dtype=torch.float16)
+print('Single-modal', (e_text @ e_image.T).tolist())
+## [[0.359619140625, 0.0655517578125], [0.04180908203125, 0.374755859375]]
 
 # How to set embedding instruction
-e_query = gme.get_text_embeddings(texts=texts, instruction='Find an image that matches the given text.')
+e_query = gme.get_text_embeddings(texts=texts, instruction=t2i_prompt)
 # If is_query=False, we always use the default instruction.
 e_corpus = gme.get_image_embeddings(images=images, is_query=False)
-print((e_query * e_corpus).sum(-1))
-## tensor([0.2433, 0.7051], dtype=torch.float16)
+print('Single-modal with instruction', (e_query @ e_corpus.T).tolist())
+## [[0.429931640625, 0.11505126953125], [0.049835205078125, 0.409423828125]]
 
 # Fused-modal embedding
 e_fused = gme.get_fused_embeddings(texts=texts, images=images)
-print((e_fused[0] * e_fused[1]).sum())
-## tensor(0.6108, dtype=torch.float16)
-
+print('Fused-modal', (e_fused @ e_fused.T).tolist())
+## [[1.0, 0.05511474609375], [0.05511474609375, 1.0]]
 ```
 
 ## Evaluation
 
-We validated the performance on our universal multimodal retrieval benchmark (**UMRB**) among others.
+We validated the performance on our universal multimodal retrieval benchmark (**UMRB**, see [Release UMRB](https://huggingface.co/Alibaba-NLP/gme-Qwen2-VL-7B-Instruct/discussions/2)) among others.
 
 |                    |      | Single-modal |           | Cross-modal |             |           | Fused-modal |            |            |             |  Avg.      |
 |--------------------|------|:------------:|:---------:|:-----------:|:-----------:|:---------:|:-----------:|:----------:|:----------:|:-----------:|:----------:|

config_sentence_transformers.json

@@ -0,0 +1,7 @@
+{
+  "prompts": {
+     "query": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": null
+}

custom_st.py

@@ -0,0 +1,221 @@
+from io import BytesIO
+from typing import Any, Dict, Optional, List
+import torch
+from PIL import Image
+from sentence_transformers.models import Transformer as BaseTransformer
+from transformers import AutoModelForVision2Seq, AutoProcessor
+
+
+class MultiModalTransformer(BaseTransformer):
+    def __init__(
+        self,
+        model_name_or_path: str,
+        cache_dir: Optional[str] = None,
+        tokenizer_args: Optional[Dict[str, Any]] = None,
+        min_image_tokens: int = 256,
+        max_image_tokens: int = 1280,
+        max_length: int = 1800,
+        **kwargs,
+    ):
+        super().__init__(model_name_or_path, **kwargs)
+        if tokenizer_args is None:
+            tokenizer_args = {}
+        tokenizer_args.pop("trust_remote_code", None)
+
+        # Initialize processor
+        min_pixels = min_image_tokens * 28 * 28
+        max_pixels = max_image_tokens * 28 * 28
+        self.processor = AutoProcessor.from_pretrained(
+            model_name_or_path, min_pixels=min_pixels, max_pixels=max_pixels, **kwargs
+        )
+        self.processor.tokenizer.padding_side = 'right'
+        self.sep = ' '
+        self.max_length = max_length
+        self.normalize = True
+
+    def _load_model(
+            self,
+            model_name_or_path: str,
+            config,
+            cache_dir: str,
+            backend: str,
+            is_peft_model: bool,
+            **model_args,
+    ) -> None:
+        model_args.pop("trust_remote_code", None)
+        self.auto_model = AutoModelForVision2Seq.from_pretrained(
+            model_name_or_path, torch_dtype=torch.float16, **model_args
+        )
+
+    def forward(
+        self, features: Dict[str, torch.Tensor], **kwargs
+    ) -> Dict[str, torch.Tensor]:       
+        if features.get("inputs_embeds", None) is None:
+            features["inputs_embeds"] = self.auto_model.base_model.embed_tokens(features["input_ids"])
+            if features.get("pixel_values", None) is not None:
+                features["pixel_values"] = features["pixel_values"].type(self.auto_model.visual.get_dtype())
+                image_embeds = self.auto_model.visual(
+                    features["pixel_values"], grid_thw=features["image_grid_thw"]
+                )
+                image_mask = features["input_ids"] == self.auto_model.config.image_token_id
+                features["inputs_embeds"][image_mask] = image_embeds
+                # features.pop("pixel_values")
+                # features.pop("image_grid_thw")
+        # features.pop("input_ids")
+        inputs = {k: v for k, v in features.items() if k in 'position_ids,attention_mask,inputs_embeds'}
+        outputs = self.auto_model.model(
+            **inputs,
+            return_dict=True,
+            output_hidden_states=True,
+            # **kwargs
+        )
+        # pooling_mask = features["attention_mask"] if features.get("pooling_mask", None) is None else features["pooling_mask"]
+        # left_padding = (pooling_mask[:, -1].sum() == pooling_mask.shape[0])  # TODO
+        # if left_padding:
+        #     embeddings = outputs.last_hidden_state
+        # else:
+        #     sequence_lengths = pooling_mask.sum(dim=1) - 1
+        #     embeddings = outputs.last_hidden_state[torch.arange(
+        #         outputs.last_hidden_state.shape[0], device=outputs.last_hidden_state.device
+        #     ), sequence_lengths]
+        features.update({"token_embeddings": outputs.last_hidden_state})
+        return features 
+
+    def tokenize(self, texts: List[List[Dict[str, Any]]] | List[str]) -> Dict[str, torch.Tensor]:
+        default_instruction = 'You are a helpful assistant.'
+
+        all_texts, all_images = list(), list()
+        for item in texts:
+            if isinstance(item, str):
+                txt, img, inst = item, None, default_instruction
+            elif isinstance(item, dict):
+                txt = item.get('text', None)
+                img = item.get('image', None)
+                inst = item.get('prompt', default_instruction)
+            else:
+                raise RuntimeError(f'Input format not supported! {item=}')
+
+            input_str = ''
+            if img is None:
+                all_images = None  # All examples in the same batch are consistent
+                # or will have ValueError: Could not make a flat list of images from xxxx
+            else:
+                input_str += '<|vision_start|><|image_pad|><|vision_end|>'
+                img = fetch_image(img)
+                all_images.append(img)
+            if txt is not None:
+                input_str += txt
+            msg = f'<|im_start|>system\n{inst}<|im_end|>\n<|im_start|>user\n{input_str}<|im_end|>\n<|im_start|>assistant\n<|endoftext|>'
+            all_texts.append(msg)
+
+        inputs = self.processor(
+            text=all_texts,
+            images=all_images,
+            padding="longest",
+            truncation=True,
+            max_length=self.max_seq_length,
+            return_tensors='pt'
+        )
+        return inputs
+
+
+### Copied from qwen_vl_utils.vision_process.py
+import base64
+from io import BytesIO
+import requests
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+
+    if max(h_bar, w_bar) / min(h_bar, w_bar) > MAX_RATIO:
+        logging.warning(
+            f"Absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(h_bar, w_bar) / min(h_bar, w_bar)}"
+        )
+        if h_bar > w_bar:
+            h_bar = w_bar * MAX_RATIO
+        else:
+            w_bar = h_bar * MAX_RATIO
+    return h_bar, w_bar
+
+
+def fetch_image(image: str | Image.Image, size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = image_obj.convert("RGB")
+    ## resize
+    # if "resized_height" in ele and "resized_width" in ele:
+    #     resized_height, resized_width = smart_resize(
+    #         ele["resized_height"],
+    #         ele["resized_width"],
+    #         factor=size_factor,
+    #     )
+    # else:
+    width, height = image.size
+    # min_pixels = ele.get("min_pixels", MIN_PIXELS)
+    # max_pixels = ele.get("max_pixels", MAX_PIXELS)
+    resized_height, resized_width = smart_resize(
+        height,
+        width,
+        factor=size_factor,
+        min_pixels=MIN_PIXELS,
+        max_pixels=MAX_PIXELS,
+    )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+###

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"
+  }
+]