Add Sentence Transformers integration (custom module)

config_sentence_transformers.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "SentenceTransformer",
+  "__version__": {
+    "sentence_transformers": "5.7.0",
+    "transformers": "5.14.1",
+    "pytorch": "2.13.0"
+  },
+  "prompts": {
+    "query": "Retrieve images or text relevant to the user's query.",
+    "document": "Represent the user's input."
+  },
+  "default_prompt_name": "query",
+  "similarity_fn_name": "cosine"
+}

custom_st.py

@@ -0,0 +1,407 @@
+"""Sentence Transformers module for fusion-embedding-2.
+
+A thin adapter that exposes the released fusion-embedding model through the
+Sentence Transformers multimodal ``encode`` API (text, image, audio, video).
+All embedding math runs through the ``fusion_embedding`` package's own
+``UnifiedEmbedder`` (the model's native loading path), so vectors produced here
+are identical to ``fusion_embedding.UnifiedEmbedder.from_pretrained(...)``:
+
+* text: chat-template instruction, EOS pooling, text-side whitening (fp32);
+* image: the frozen base's native vision path (no whitening, no adapters);
+* video: the released video preprocessing over the frozen base's video path;
+* audio: soxr resampling to 16 kHz, Whisper-style mel, trained resampler, and
+  the frozen decoder with ONLY the audio adapter gate open.
+
+Every vector is L2-normalized at the full interop dimension (2048). Shorter
+Matryoshka rungs: pass ``truncate_dim=<rung>`` together with
+``normalize_embeddings=True`` to ``encode`` (truncate-then-renormalize equals
+the native MRL readout).
+
+Requirements (beyond sentence-transformers>=5.5.1):
+
+    pip install "fusion-embedding[sense]>=0.3.0" torchvision
+
+The ``sense`` extra pulls the audio decode/resample stack (soundfile, librosa);
+transformers itself ships with sentence-transformers. torchvision is required to
+load the model at all, not only for video, because the base processor builds a
+video processor during construction; this applies to the native loader too.
+Embedding an image needs Pillow, which torchvision carries. Embedding a video by
+file path additionally requires torchcodec, which in turn needs FFmpeg.
+
+Supported inputs per item (one modality per item; the model has no fused
+multi-modality input):
+
+* ``str`` — text, or a local image/audio/video file path (auto-detected);
+* ``PIL.Image.Image`` or an HxWxC uint8 array — image;
+* ``{"audio": {"array": waveform, "sampling_rate": sr}}`` (or the inner dict
+  directly) — audio; a bare 1-D array is rejected because the sampling rate
+  would be unknown;
+* a ``[T, C, H, W]`` uint8 frame tensor/array — video.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Optional
+
+import torch
+
+try:
+    from sentence_transformers.base.modality import infer_modality
+    from sentence_transformers.base.modules.input_module import InputModule
+except ImportError as exc:  # pragma: no cover - version guard
+    raise ImportError(
+        "The fusion-embedding-2 Sentence Transformers integration requires "
+        "sentence-transformers>=5.5.1 (multimodal encode). "
+        "Upgrade with: pip install -U sentence-transformers"
+    ) from exc
+
+try:
+    from fusion_embedding.config import INSTRUCTION_REGISTRY
+    from fusion_embedding.model import last_token_pool
+    from fusion_embedding.unified import UnifiedEmbedder, _chat
+except ImportError as exc:  # pragma: no cover - dependency guard
+    raise ImportError(
+        "The fusion-embedding-2 Sentence Transformers integration needs the "
+        "fusion-embedding package for the model implementation. Install it "
+        "with: pip install 'fusion-embedding[sense]>=0.3.0'"
+    ) from exc
+
+_MIN_ST_VERSION = (5, 5, 1)
+
+
+def _require_min_st_version() -> None:
+    """Fail with the real reason on Sentence Transformers older than 5.5.1.
+
+    Single-key modality dicts such as {"audio": {"array": ..., "sampling_rate": ...}}
+    are classified as a tuple by infer_modality before 5.5.1, so the encode call is
+    rejected by Sentence Transformers itself with a message claiming the modality is
+    unsupported. The import guard above cannot catch that: sentence_transformers.base
+    imports cleanly on 5.4.x.
+    """
+    import sentence_transformers
+
+    raw = getattr(sentence_transformers, "__version__", "0")
+    parts = []
+    for chunk in raw.split(".")[:3]:
+        digits = "".join(c for c in chunk if c.isdigit())
+        parts.append(int(digits) if digits else 0)
+    while len(parts) < 3:
+        parts.append(0)
+    if tuple(parts) < _MIN_ST_VERSION:
+        raise ImportError(
+            "The fusion-embedding-2 Sentence Transformers integration requires "
+            f"sentence-transformers>=5.5.1, found {raw}. Earlier versions reject "
+            "single-key modality dicts such as "
+            '{"audio": {"array": ..., "sampling_rate": ...}} before this module is '
+            "reached. Upgrade with: pip install -U 'sentence-transformers>=5.5.1'"
+        )
+
+
+CKPT_FILENAME = "fusion-embedding-2-2b-preview.pt"
+
+
+class FusionEmbedding2Module(InputModule):
+    """Single Sentence Transformers module wrapping the full fusion-embedding-2
+    encoder (all modalities plus the canonical readout, so no separate Pooling
+    or Normalize module is needed: ``forward`` emits ``sentence_embedding``
+    directly)."""
+
+    config_file_name = "sentence_bert_config.json"
+    config_keys = ["ckpt_filename", "max_seq_length"]
+    save_in_root = True
+
+    def __init__(
+        self,
+        model_name_or_path: Optional[str] = None,
+        ckpt_filename: str = CKPT_FILENAME,
+        max_seq_length: int = 512,
+        revision: Optional[str] = None,
+        token: "bool | str | None" = None,
+        cache_folder: Optional[str] = None,
+        local_files_only: bool = False,
+        model_kwargs: Optional[dict] = None,
+        embedder: Optional[UnifiedEmbedder] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        _require_min_st_version()
+        self.ckpt_filename = ckpt_filename
+        self.max_seq_length = max_seq_length
+
+        if embedder is None:
+            if model_name_or_path is None:
+                raise ValueError("model_name_or_path is required (or pass embedder=)")
+            model_kwargs = dict(model_kwargs or {})
+            dtype = model_kwargs.pop("torch_dtype", model_kwargs.pop("dtype", torch.bfloat16))
+            if isinstance(dtype, str):
+                dtype = getattr(torch, dtype)
+            device = model_kwargs.pop(
+                "device", "cuda" if torch.cuda.is_available() else "cpu"
+            )
+            ckpt_path = self.load_file_path(
+                model_name_or_path,
+                filename=ckpt_filename,
+                token=token,
+                cache_folder=cache_folder,
+                revision=revision,
+                local_files_only=local_files_only,
+            )
+            if ckpt_path is None:
+                raise FileNotFoundError(
+                    f"checkpoint {ckpt_filename!r} not found in {model_name_or_path!r}"
+                )
+            embedder = UnifiedEmbedder.from_pretrained(ckpt_path, device=device, dtype=dtype)
+
+        self._emb = embedder
+        # Wire the video seam the UnifiedEmbedder anticipates: the released video
+        # preprocessing (fusion_embedding.multimodal) over the frozen base.
+        self._emb._video_pooler = self._video_pooled
+
+        # Register the underlying torch modules so Sentence Transformers device
+        # management (`model.to(device)`) moves the whole stack.
+        self.fusion_model = embedder.model
+        if embedder.full is not None:
+            self.base = embedder.full
+        if embedder.tok is not None:
+            self.tokenizer = embedder.tok
+
+    # ------------------------------------------------------------------ loading
+    @classmethod
+    def load(
+        cls,
+        model_name_or_path: str,
+        subfolder: str = "",
+        token: "bool | str | None" = None,
+        cache_folder: Optional[str] = None,
+        revision: Optional[str] = None,
+        local_files_only: bool = False,
+        trust_remote_code: bool = False,
+        model_kwargs: Optional[dict] = None,
+        processor_kwargs: Optional[dict] = None,
+        config_kwargs: Optional[dict] = None,
+        backend: str = "torch",
+        **kwargs,
+    ) -> "FusionEmbedding2Module":
+        if backend != "torch":
+            raise ValueError(
+                f"fusion-embedding-2 only supports the torch backend, got {backend!r}"
+            )
+        config = cls.load_config(
+            model_name_or_path,
+            subfolder=subfolder,
+            token=token,
+            cache_folder=cache_folder,
+            revision=revision,
+            local_files_only=local_files_only,
+        )
+        config.pop("model_name_or_path", None)
+        if config_kwargs:
+            config.update(config_kwargs)
+        return cls(
+            model_name_or_path,
+            revision=revision,
+            token=token,
+            cache_folder=cache_folder,
+            local_files_only=local_files_only,
+            model_kwargs=model_kwargs,
+            **config,
+        )
+
+    # -------------------------------------------------------------- ST contract
+    @property
+    def modalities(self) -> list:
+        return ["text", "image", "audio", "video"]
+
+    def get_embedding_dimension(self) -> int:
+        return int(self._emb.contract.dim)
+
+    def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
+        # Configuration only: the 2B weights live in the model repository's
+        # checkpoint file and are not duplicated by Sentence Transformers saves.
+        self.save_config(output_path)
+
+    # ------------------------------------------------------------- input parsing
+    def preprocess(self, inputs: list, prompt: Optional[str] = None, **kwargs) -> dict:
+        items = []
+        for item in inputs:
+            modality = infer_modality(item, supported_modalities=self.modalities)
+            if isinstance(modality, tuple):
+                raise ValueError(
+                    "fusion-embedding-2 embeds one modality per input item; "
+                    f"got a combined input with {modality}. Encode each modality "
+                    "separately (the shared space makes the vectors comparable)."
+                )
+            if isinstance(item, dict) and set(item.keys()) == {modality}:
+                item = item[modality]
+            items.append((modality, self._parse(modality, item)))
+        return {"fusion_inputs": items, "fusion_prompt": prompt}
+
+    def _parse(self, modality: str, item: Any) -> Any:
+        if modality == "text":
+            return item
+        if modality == "image":
+            return self._parse_image(item)
+        if modality == "audio":
+            return self._parse_audio(item)
+        if modality == "video":
+            return self._parse_video(item)
+        raise ValueError(f"unsupported modality {modality!r}")
+
+    @staticmethod
+    def _parse_image(item):
+        import numpy as np
+
+        # Guarded so transformers' trust_remote_code import check does not make
+        # Pillow a load-time requirement: dynamic_module_utils.get_imports skips
+        # ast.Try blocks, and an unguarded import here is otherwise treated as
+        # mandatory at construction even for text-only use.
+        try:
+            from PIL import Image
+        except ImportError as exc:  # pragma: no cover - optional dependency
+            raise ImportError(
+                "embedding an image requires Pillow (pip install pillow)"
+            ) from exc
+
+        if isinstance(item, Image.Image):
+            return item
+        if isinstance(item, str):
+            if item.startswith(("http://", "https://", "data:")):
+                raise ValueError(
+                    "image URLs / data URIs are not supported; download the file "
+                    "and pass a local path or a PIL image"
+                )
+            return item  # local path; decoded by the native path (PIL)
+        if isinstance(item, torch.Tensor):
+            item = item.cpu().numpy()
+        if isinstance(item, np.ndarray):
+            if item.ndim == 3 and item.shape[0] in (1, 3, 4) and item.shape[-1] not in (1, 3, 4):
+                item = np.transpose(item, (1, 2, 0))  # CHW -> HWC
+            if item.ndim != 3 or item.shape[-1] not in (1, 3, 4):
+                raise ValueError(f"expected an HxWxC image array, got shape {item.shape}")
+            if item.dtype != np.uint8:
+                item = np.clip(item, 0, 255).astype(np.uint8)
+            return Image.fromarray(item.squeeze(-1) if item.shape[-1] == 1 else item)
+        raise ValueError(f"unsupported image input type {type(item).__name__}")
+
+    @staticmethod
+    def _parse_audio(item):
+        """Return (payload, sampling_rate_or_None); paths carry their own rate."""
+        import numpy as np
+
+        if isinstance(item, str):
+            if item.startswith(("http://", "https://")):
+                raise ValueError(
+                    "audio URLs are not supported; download the file and pass a "
+                    "local path or {'array': ..., 'sampling_rate': ...}"
+                )
+            return (item, None)
+        if isinstance(item, dict):
+            if "array" not in item or "sampling_rate" not in item:
+                raise ValueError(
+                    "audio dicts must have the form "
+                    "{'array': waveform, 'sampling_rate': sr}"
+                )
+            array, sr = item["array"], int(item["sampling_rate"])
+        else:
+            try:  # torchcodec AudioDecoder (optional dependency)
+                from torchcodec.decoders import AudioDecoder
+            except ImportError:
+                AudioDecoder = None
+            if AudioDecoder is not None and isinstance(item, AudioDecoder):
+                samples = item.get_all_samples()
+                return (samples.data.mean(dim=0).cpu().numpy(), int(samples.sample_rate))
+            raise ValueError(
+                "a bare audio array has no sampling rate; pass "
+                "{'audio': {'array': waveform, 'sampling_rate': sr}} instead"
+            )
+        if isinstance(array, torch.Tensor):
+            array = array.cpu().numpy()
+        array = np.asarray(array)
+        if array.ndim == 2 and array.shape[0] < array.shape[1]:
+            array = array.T  # (channels, samples) -> (samples, channels)
+        if array.ndim > 2:
+            raise ValueError(f"expected a 1-D or 2-D waveform, got shape {array.shape}")
+        return (array.astype(np.float32, copy=False), sr)
+
+    @staticmethod
+    def _parse_video(item):
+        import numpy as np
+
+        if isinstance(item, str):
+            if item.startswith(("http://", "https://")):
+                raise ValueError(
+                    "video URLs are not supported; download the file and pass a "
+                    "local path or a [T, C, H, W] frame tensor"
+                )
+            return item  # local path; decoded natively (torchcodec, 1 fps, <=64 frames)
+        if isinstance(item, dict):
+            # {"array": frames, "video_metadata": ...}: the released frame-tensor
+            # path derives its own metadata, so user metadata is not consumed.
+            item = item["array"]
+        if isinstance(item, np.ndarray):
+            item = torch.from_numpy(np.ascontiguousarray(item))
+        if isinstance(item, torch.Tensor):
+            if item.ndim == 5 and item.shape[0] == 1:
+                item = item.squeeze(0)
+            if item.ndim == 4 and item.shape[-1] in (1, 3) and item.shape[1] not in (1, 3):
+                item = item.permute(0, 3, 1, 2)  # THWC -> TCHW
+            if item.ndim != 4:
+                raise ValueError(f"expected a [T, C, H, W] frame tensor, got shape {list(item.shape)}")
+            return item
+        raise ValueError(f"unsupported video input type {type(item).__name__}")
+
+    # ------------------------------------------------------------------ forward
+    def forward(self, features: dict, **kwargs) -> dict:
+        self._sync_device()
+        prompt = features.get("fusion_prompt")
+        vectors = []
+        for modality, payload in features["fusion_inputs"]:
+            if modality == "text":
+                vectors.append(self._emb.embed_text(payload, instruction=prompt or None))
+            elif modality == "image":
+                vectors.append(self._emb.embed_image(payload))
+            elif modality == "audio":
+                array, sr = payload
+                vectors.append(self._emb.embed_audio(array, sr=sr))
+            elif modality == "video":
+                vectors.append(self._emb.embed_video(payload))
+            else:  # pragma: no cover - guarded in preprocess
+                raise ValueError(f"unsupported modality {modality!r}")
+        features["sentence_embedding"] = torch.stack(vectors)
+        return features
+
+    def _sync_device(self) -> None:
+        """Follow Sentence Transformers device moves (`model.to(...)`)."""
+        param = next(self.parameters(), None)
+        if param is not None:
+            self._emb.device = param.device
+
+    # ----------------------------------------------------------- native video path
+    @torch.no_grad()
+    def _video_pooled(self, video, fps, max_frames) -> torch.Tensor:
+        """The released fusion-embedding video path: reference-exact frame
+        preprocessing (fusion_embedding.multimodal) -> frozen base's video
+        forward -> EOS pooling. Runs with every adapter gate closed."""
+        from fusion_embedding.config import VIDEO_USER_CONTENT
+        from fusion_embedding.multimodal import _v_prepare, _v_resize_video
+
+        emb = self._emb
+        gate = getattr(emb.model, "_adapter_gate", None)
+        if gate is not None and gate.active:
+            raise RuntimeError("adapter gate is open during a video embed")
+        if emb.full is None or emb.proc is None:
+            raise RuntimeError("video embedding needs the real processor + base")
+        frames, metadata = _v_prepare(video, fps, max_frames)
+        frames = _v_resize_video(frames)
+        text = _chat(INSTRUCTION_REGISTRY["doc"], VIDEO_USER_CONTENT)
+        inputs = emb.proc(
+            text=[text],
+            videos=[frames],
+            video_metadata=[metadata],
+            do_resize=False,
+            do_sample_frames=False,
+            return_tensors="pt",
+        ).to(emb.device)
+        hidden = emb.full(**inputs).last_hidden_state
+        return last_token_pool(hidden, inputs["attention_mask"])

modules.json

@@ -0,0 +1,8 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "custom_st.FusionEmbedding2Module"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "ckpt_filename": "fusion-embedding-2-2b-preview.pt",
+  "max_seq_length": 512
+}