initial sentence-transformers support

config_sentence_transformers.json

@@ -0,0 +1,13 @@
+{
+    "__version__": {
+      "sentence_transformers": "3.3.0",
+      "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|>"
+    },
+    "default_prompt_name": null,
+    "similarity_fn_name": "cosine"
+  }

custom_st.py

@@ -0,0 +1,143 @@
+import base64
+import json
+import os
+import math
+from io import BytesIO
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import requests
+import torch
+from PIL import Image
+from torch import nn
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
+
+class Transformer(nn.Module):
+    save_in_root: bool = True
+    
+    def __init__(
+        self,
+        model_name_or_path: str = 'llamaindex/vdr-2b-multi-v1',
+        processor_name_or_path: Optional[str] = None,
+        max_pixels: int = 768 * 28 * 28,
+        min_pixels: int = 1 * 28 * 28,
+        dimension: int = 2048,
+        cache_dir: Optional[str] = None,
+        device: str = 'cuda:0',
+        **kwargs,
+    ) -> None:
+        super(Transformer, self).__init__()
+        
+        self.device = device
+        self.dimension = dimension
+        self.max_pixels = max_pixels
+        self.min_pixels = min_pixels
+        
+        # Initialize model
+        self.model = Qwen2VLForConditionalGeneration.from_pretrained(
+            model_name_or_path,
+            attn_implementation="flash_attention_2",
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+            cache_dir=cache_dir,
+            **kwargs
+        ).eval()
+
+        # Initialize processor
+        self.processor = AutoProcessor.from_pretrained(
+            processor_name_or_path or model_name_or_path,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+            cache_dir=cache_dir
+        )
+
+        self.model.padding_side = "left"
+        self.processor.tokenizer.padding_side = "left"
+
+        self.document_prompt = "<|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|>"
+        self.query_prompt = "<|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|>"
+
+    def _smart_resize(self, height: int, width: int) -> tuple[int, int]:
+        h_bar = max(28, self._round_by_factor(height, 28))
+        w_bar = max(28, self._round_by_factor(width, 28))
+        if h_bar * w_bar > self.max_pixels:
+            beta = math.sqrt((height * width) / self.max_pixels)
+            h_bar = self._floor_by_factor(height / beta, 28)
+            w_bar = self._floor_by_factor(width / beta, 28)
+        elif h_bar * w_bar < self.min_pixels:
+            beta = math.sqrt(self.min_pixels / (height * width))
+            h_bar = self._ceil_by_factor(height * beta, 28)
+            w_bar = self._ceil_by_factor(width * beta, 28)
+        return w_bar, h_bar
+
+    @staticmethod
+    def _round_by_factor(number: float, factor: int) -> int:
+        return round(number / factor) * factor
+
+    @staticmethod
+    def _ceil_by_factor(number: float, factor: int) -> int:
+        return math.ceil(number / factor) * factor
+
+    @staticmethod
+    def _floor_by_factor(number: float, factor: int) -> int:
+        return math.floor(number / factor) * factor
+
+    def _resize_image(self, image: Image.Image) -> Image.Image:
+        new_size = self._smart_resize(image.height, image.width)
+        return image.resize(new_size)
+
+    @staticmethod
+    def _decode_data_image(data_image_str: str) -> Image.Image:
+        header, data = data_image_str.split(',', 1)
+        image_data = base64.b64decode(data)
+        return Image.open(BytesIO(image_data))
+
+    def _process_input(self, texts: List[Union[str, Image.Image]]) -> tuple[List[str], List[Image.Image]]:
+        processed_texts = []
+        processed_images = []
+        dummy_image = Image.new('RGB', (56, 56))
+
+        for sample in texts:
+            if isinstance(sample, str):
+                processed_texts.append(self.query_prompt % sample)
+                processed_images.append(dummy_image)
+            elif isinstance(sample, Image.Image):
+                processed_texts.append(self.document_prompt)
+                processed_images.append(self._resize_image(sample))
+
+        return processed_texts, processed_images
+
+    def forward(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        cache_position = torch.arange(0, features['input_ids'].shape[0])
+        inputs = self.model.prepare_inputs_for_generation(
+            **features, cache_position=cache_position, use_cache=False
+        )
+
+        with torch.no_grad():
+            output = self.model(
+                **inputs,
+                return_dict=True,
+                output_hidden_states=True
+            )
+
+        embeddings = output.hidden_states[-1][:, -1]
+        features['sentence_embedding'] = torch.nn.functional.normalize(
+            embeddings[:, :self.dimension], p=2, dim=-1
+        )
+        return features
+
+    def tokenize(self, texts: List[Union[str, Image.Image]], padding: str = 'longest') -> Dict[str, torch.Tensor]:
+        processed_texts, processed_images = self._process_input(texts)
+        
+        inputs = self.processor(
+            text=processed_texts,
+            images=processed_images,
+            videos=None,
+            padding=padding,
+            return_tensors='pt'
+        )
+        
+        return {k: v.to(self.device) for k, v in inputs.items()}
+
+    def save(self, output_path: str, safe_serialization: bool = True) -> None:
+        self.model.save_pretrained(output_path, safe_serialization=safe_serialization)
+        self.processor.save_pretrained(output_path)

modules.json

@@ -0,0 +1,19 @@
+[
+    {
+        "idx": 0,
+        "name": "transformer",
+        "path": "",
+        "type": "custom_st.Transformer",
+        "model_name_or_path": "llamaindex/vdr-2b-multi-v1",
+        "dimension": 2048,
+        "max_pixels": 602112,
+        "min_pixels": 784,
+        "device": "cuda:0"
+    },
+    {
+        "idx": 1,
+        "name": "normalizer",
+        "path": "1_Normalize",
+        "type": "sentence_transformers.models.Normalize"
+    }
+]