quadembed 0.1.0__tar.gz

quadembed-0.1.0/.gitignore

@@ -0,0 +1,24 @@
+.venv/
+__pycache__/
+*.pyc
+.claude/
+scratch/
+
+# large binary checkpoints and cached datasets -- not meant for git,
+# the trained model is published separately on Hugging Face:
+# https://huggingface.co/Mithil-21/quadembed-nano
+checkpoints/
+checkpoints_round*_backup/
+data_cache/
+*.pt
+
+# staging folder for the Hugging Face upload, mirrors published content
+hf_repo/
+
+# python build artifacts
+package/dist/
+*.egg-info/
+
+# personal blog draft and its working assets -- not meant for the public repo
+medium-blog-draft.md
+blog_assets/

quadembed-0.1.0/LICENSE

@@ -0,0 +1,45 @@
+Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)
+
+Copyright (c) 2026 Mithil Maske
+
+This package and its associated model weights are licensed under the Creative
+Commons Attribution-NonCommercial 4.0 International License.
+
+You are free to:
+
+  Share  - copy and redistribute the material in any medium or format
+  Adapt  - remix, transform, and build upon the material
+
+Under the following terms:
+
+  Attribution  - You must give appropriate credit, provide a link to the
+                 license, and indicate if changes were made.
+
+  NonCommercial - You may not use the material for commercial purposes.
+
+Full license text: https://creativecommons.org/licenses/by-nc/4.0/legalcode
+
+
+WHY NON-COMMERCIAL
+------------------
+
+QuadEmbed is built on top of three frozen pretrained encoders. One of them,
+jinaai/jina-embeddings-v5-text-nano, is itself licensed CC BY-NC 4.0, and it
+defines the embedding space that the trained projectors map into. That
+non-commercial restriction therefore carries through to this work.
+
+The other two frozen encoders are permissively licensed:
+  - google/siglip2-base-patch16-naflex  (Apache-2.0)
+  - openai/whisper-large-v3             (Apache-2.0)
+
+If you need commercial use, you would need a separate license for the text
+encoder from Jina AI, or to retrain the projectors against a permissively
+licensed text encoder instead.
+
+
+ATTRIBUTION
+-----------
+
+The architecture reproduced here (GELATO) is described in Jina AI's paper on
+jina-embeddings-v5-omni, arXiv:2605.08384. This project is an independent
+reproduction and is not affiliated with or endorsed by Jina AI.

quadembed-0.1.0/PKG-INFO

@@ -0,0 +1,157 @@
+Metadata-Version: 2.5
+Name: quadembed
+Version: 0.1.0
+Summary: Multimodal embeddings for text, image, audio, and video in one shared vector space, trained on a single consumer GPU
+Project-URL: Homepage, https://github.com/mithilai/QuadEmbed
+Project-URL: Repository, https://github.com/mithilai/QuadEmbed
+Project-URL: Model Weights, https://huggingface.co/Mithil-21/quadembed-nano
+Project-URL: Write-up, https://medium.com/@mithilmaske/i-built-a-multimodal-embedding-model-from-scratch-on-an-rtx-4060-text-image-audio-and-video-ab1fef04f1cd
+Author: Mithil Maske
+License: CC-BY-NC-4.0
+License-File: LICENSE
+Keywords: audio,clip,contrastive-learning,embeddings,multimodal,retrieval,semantic-search,text-to-image,video
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Multimedia :: Sound/Audio :: Analysis
+Classifier: Topic :: Multimedia :: Video
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: huggingface-hub>=0.24
+Requires-Dist: numpy>=1.24
+Requires-Dist: pillow>=9.0
+Requires-Dist: torch>=2.1
+Requires-Dist: transformers>=4.45
+Provides-Extra: all
+Requires-Dist: librosa>=0.10; extra == 'all'
+Requires-Dist: opencv-python>=4.8; extra == 'all'
+Requires-Dist: soundfile>=0.12; extra == 'all'
+Provides-Extra: audio
+Requires-Dist: librosa>=0.10; extra == 'audio'
+Requires-Dist: soundfile>=0.12; extra == 'audio'
+Provides-Extra: video
+Requires-Dist: opencv-python>=4.8; extra == 'video'
+Description-Content-Type: text/markdown
+
+# QuadEmbed
+
+Multimodal embeddings for **text, image, audio, and video** in one shared
+768-dimensional vector space. Compare any modality against any other with a
+dot product.
+
+Trained end-to-end on a single RTX 4060 laptop GPU (8GB VRAM) by freezing
+three pretrained encoders and training only two small projection heads,
+reproducing the GELATO architecture behind Jina AI's jina-embeddings-v5-omni.
+
+> ⚠️ **Non-commercial license.** QuadEmbed builds on
+> `jina-embeddings-v5-text-nano`, which is CC-BY-NC-4.0. That restriction
+> carries through to this package and its weights. Research and educational
+> use is fine; commercial use is not, without a separate license from Jina AI.
+
+## Install
+
+```bash
+pip install quadembed              # text + image + audio
+pip install quadembed[video]       # adds video support (OpenCV)
+pip install quadembed[all]         # everything, plus audio file loading helpers
+```
+
+## Quickstart
+
+```python
+from quadembed import QuadEmbed
+from PIL import Image
+
+model = QuadEmbed.from_pretrained()   # downloads weights from the Hub on first run
+
+texts = ["a dog running on the beach", "a bowl of ramen noodles"]
+text_embeds = model.embed_text(texts)
+image_embeds = model.embed_image([Image.open("photo.jpg").convert("RGB")])
+
+# embeddings are L2-normalized, so this is cosine similarity
+similarity = text_embeds @ image_embeds.T
+for text, score in zip(texts, similarity[:, 0].tolist()):
+    print(f"{score:.3f}  {text}")
+```
+
+### Audio
+
+Pass mono float32 arrays at 16 kHz (what the frozen Whisper encoder expects):
+
+```python
+import soundfile as sf
+
+audio, sr = sf.read("clip.wav")
+audio_embeds = model.embed_audio([audio.astype("float32")], sampling_rate=sr)
+similarity = model.embed_text(["a dog barking"]) @ audio_embeds.T
+```
+
+### Video
+
+```python
+model = QuadEmbed.from_pretrained(modalities=("text", "video"))
+video_embeds = model.embed_video_file("clip.mp4", num_frames=4)
+similarity = model.embed_text(["a person skateboarding"]) @ video_embeds.T
+```
+
+### Loading only what you need
+
+Each encoder costs memory and download time. Load a subset:
+
+```python
+model = QuadEmbed.from_pretrained(modalities=("text", "vision"))   # skip audio
+model = QuadEmbed.from_pretrained(device="cpu")                    # force CPU
+```
+
+## How it works
+
+Three frozen encoders, two trained projectors:
+
+| Role | Model (frozen) | Trainable on top |
+|---|---|---|
+| Text (anchor) | `jinaai/jina-embeddings-v5-text-nano` (239M) | nothing, it defines the space |
+| Vision | `google/siglip2-base-patch16-naflex` | vision projector (2.36M params) |
+| Audio | `openai/whisper-large-v3` encoder | audio projector (0.98M params) |
+
+Only ~3.3M parameters were ever trained, against nearly a billion frozen ones.
+Video needs no encoder or projector of its own: frames are sampled, run
+through the vision path, and mean-pooled over time.
+
+Training used bidirectional in-batch InfoNCE (temperature 0.02) plus
+Matryoshka representation learning over prefix dims {32, 64, 128, 256, 768},
+so truncated embeddings remain usable if you need a smaller index.
+
+## Measured performance
+
+Cross-modal retrieval recall@k on held-out splits, text-query direction:
+
+| Modality | R@1 | R@5 | R@10 | n |
+|---|---|---|---|---|
+| Image | 13.7% | 68.6% | 81.1% | 1024 |
+| Audio | 67% | 97% | 100% | 33 |
+| Video | 40% | 86% | 94% | 50 |
+
+Random-chance R@1 on the 1024-candidate image eval is ~0.1%.
+
+These are honest small-scale numbers. Vision was trained on ~172k
+image-caption pairs, orders of magnitude less than production embedding
+models, and image R@1 is the weakest metric as a result. R@5 and R@10 are
+considerably stronger, so this is more useful for candidate retrieval and
+reranking than for exact top-1 matching.
+
+## Links
+
+- **Source and training code:** [github.com/mithilai/QuadEmbed](https://github.com/mithilai/QuadEmbed)
+- **Model weights:** [huggingface.co/Mithil-21/quadembed-nano](https://huggingface.co/Mithil-21/quadembed-nano)
+- **Full write-up:** [I Built a Multimodal Embedding Model From Scratch on an RTX 4060](https://medium.com/@mithilmaske/i-built-a-multimodal-embedding-model-from-scratch-on-an-rtx-4060-text-image-audio-and-video-ab1fef04f1cd)
+
+## Credit
+
+QuadEmbed reproduces the architecture described in Jina AI's GELATO paper
+([arXiv:2605.08384](https://arxiv.org/abs/2605.08384)). All credit for the
+original architecture and training recipe belongs there. This is an
+independent reproduction, not affiliated with or endorsed by Jina AI.

quadembed-0.1.0/README.md

@@ -0,0 +1,118 @@
+# QuadEmbed
+
+Multimodal embeddings for **text, image, audio, and video** in one shared
+768-dimensional vector space. Compare any modality against any other with a
+dot product.
+
+Trained end-to-end on a single RTX 4060 laptop GPU (8GB VRAM) by freezing
+three pretrained encoders and training only two small projection heads,
+reproducing the GELATO architecture behind Jina AI's jina-embeddings-v5-omni.
+
+> ⚠️ **Non-commercial license.** QuadEmbed builds on
+> `jina-embeddings-v5-text-nano`, which is CC-BY-NC-4.0. That restriction
+> carries through to this package and its weights. Research and educational
+> use is fine; commercial use is not, without a separate license from Jina AI.
+
+## Install
+
+```bash
+pip install quadembed              # text + image + audio
+pip install quadembed[video]       # adds video support (OpenCV)
+pip install quadembed[all]         # everything, plus audio file loading helpers
+```
+
+## Quickstart
+
+```python
+from quadembed import QuadEmbed
+from PIL import Image
+
+model = QuadEmbed.from_pretrained()   # downloads weights from the Hub on first run
+
+texts = ["a dog running on the beach", "a bowl of ramen noodles"]
+text_embeds = model.embed_text(texts)
+image_embeds = model.embed_image([Image.open("photo.jpg").convert("RGB")])
+
+# embeddings are L2-normalized, so this is cosine similarity
+similarity = text_embeds @ image_embeds.T
+for text, score in zip(texts, similarity[:, 0].tolist()):
+    print(f"{score:.3f}  {text}")
+```
+
+### Audio
+
+Pass mono float32 arrays at 16 kHz (what the frozen Whisper encoder expects):
+
+```python
+import soundfile as sf
+
+audio, sr = sf.read("clip.wav")
+audio_embeds = model.embed_audio([audio.astype("float32")], sampling_rate=sr)
+similarity = model.embed_text(["a dog barking"]) @ audio_embeds.T
+```
+
+### Video
+
+```python
+model = QuadEmbed.from_pretrained(modalities=("text", "video"))
+video_embeds = model.embed_video_file("clip.mp4", num_frames=4)
+similarity = model.embed_text(["a person skateboarding"]) @ video_embeds.T
+```
+
+### Loading only what you need
+
+Each encoder costs memory and download time. Load a subset:
+
+```python
+model = QuadEmbed.from_pretrained(modalities=("text", "vision"))   # skip audio
+model = QuadEmbed.from_pretrained(device="cpu")                    # force CPU
+```
+
+## How it works
+
+Three frozen encoders, two trained projectors:
+
+| Role | Model (frozen) | Trainable on top |
+|---|---|---|
+| Text (anchor) | `jinaai/jina-embeddings-v5-text-nano` (239M) | nothing, it defines the space |
+| Vision | `google/siglip2-base-patch16-naflex` | vision projector (2.36M params) |
+| Audio | `openai/whisper-large-v3` encoder | audio projector (0.98M params) |
+
+Only ~3.3M parameters were ever trained, against nearly a billion frozen ones.
+Video needs no encoder or projector of its own: frames are sampled, run
+through the vision path, and mean-pooled over time.
+
+Training used bidirectional in-batch InfoNCE (temperature 0.02) plus
+Matryoshka representation learning over prefix dims {32, 64, 128, 256, 768},
+so truncated embeddings remain usable if you need a smaller index.
+
+## Measured performance
+
+Cross-modal retrieval recall@k on held-out splits, text-query direction:
+
+| Modality | R@1 | R@5 | R@10 | n |
+|---|---|---|---|---|
+| Image | 13.7% | 68.6% | 81.1% | 1024 |
+| Audio | 67% | 97% | 100% | 33 |
+| Video | 40% | 86% | 94% | 50 |
+
+Random-chance R@1 on the 1024-candidate image eval is ~0.1%.
+
+These are honest small-scale numbers. Vision was trained on ~172k
+image-caption pairs, orders of magnitude less than production embedding
+models, and image R@1 is the weakest metric as a result. R@5 and R@10 are
+considerably stronger, so this is more useful for candidate retrieval and
+reranking than for exact top-1 matching.
+
+## Links
+
+- **Source and training code:** [github.com/mithilai/QuadEmbed](https://github.com/mithilai/QuadEmbed)
+- **Model weights:** [huggingface.co/Mithil-21/quadembed-nano](https://huggingface.co/Mithil-21/quadembed-nano)
+- **Full write-up:** [I Built a Multimodal Embedding Model From Scratch on an RTX 4060](https://medium.com/@mithilmaske/i-built-a-multimodal-embedding-model-from-scratch-on-an-rtx-4060-text-image-audio-and-video-ab1fef04f1cd)
+
+## Credit
+
+QuadEmbed reproduces the architecture described in Jina AI's GELATO paper
+([arXiv:2605.08384](https://arxiv.org/abs/2605.08384)). All credit for the
+original architecture and training recipe belongs there. This is an
+independent reproduction, not affiliated with or endorsed by Jina AI.

quadembed-0.1.0/pyproject.toml

@@ -0,0 +1,59 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "quadembed"
+version = "0.1.0"
+description = "Multimodal embeddings for text, image, audio, and video in one shared vector space, trained on a single consumer GPU"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "CC-BY-NC-4.0" }
+authors = [{ name = "Mithil Maske" }]
+keywords = [
+    "embeddings",
+    "multimodal",
+    "retrieval",
+    "contrastive-learning",
+    "clip",
+    "text-to-image",
+    "audio",
+    "video",
+    "semantic-search",
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Multimedia :: Sound/Audio :: Analysis",
+    "Topic :: Multimedia :: Video",
+]
+dependencies = [
+    "torch>=2.1",
+    "transformers>=4.45",
+    "huggingface-hub>=0.24",
+    "pillow>=9.0",
+    "numpy>=1.24",
+]
+
+[project.optional-dependencies]
+video = ["opencv-python>=4.8"]
+audio = ["soundfile>=0.12", "librosa>=0.10"]
+all = ["opencv-python>=4.8", "soundfile>=0.12", "librosa>=0.10"]
+
+[project.urls]
+Homepage = "https://github.com/mithilai/QuadEmbed"
+Repository = "https://github.com/mithilai/QuadEmbed"
+"Model Weights" = "https://huggingface.co/Mithil-21/quadembed-nano"
+"Write-up" = "https://medium.com/@mithilmaske/i-built-a-multimodal-embedding-model-from-scratch-on-an-rtx-4060-text-image-audio-and-video-ab1fef04f1cd"
+
+[tool.hatch.build.targets.wheel]
+packages = ["quadembed"]
+
+[tool.hatch.build.targets.sdist]
+include = ["quadembed", "README.md", "LICENSE"]

quadembed-0.1.0/quadembed/__init__.py

@@ -0,0 +1,168 @@
+"""QuadEmbed: a local multimodal embedding model covering text, image, audio,
+and video in one shared 768-dim space.
+
+Reproduces the GELATO architecture (jina-embeddings-v5-omni): three frozen
+pretrained encoders plus two small trained projectors. Only the projectors
+were trained; every encoder is used as published.
+
+    from quadembed import QuadEmbed
+
+    model = QuadEmbed.from_pretrained()          # pulls weights from the Hub
+    text = model.embed_text(["a dog on the beach"])
+    image = model.embed_image([Image.open("photo.jpg")])
+    print(text @ image.T)                        # cosine similarity
+"""
+import os
+
+import torch
+import torch.nn.functional as F
+
+from .encoders import AudioEncoder, TextEncoder, VisionEncoder
+from .losses import bidirectional_infonce, matryoshka_infonce
+from .projectors import AudioProjector, VisionProjector, embed_video_from_frames
+from .video import embed_video, sample_frames
+
+__version__ = "0.1.0"
+
+DEFAULT_REPO_ID = "Mithil-21/quadembed-nano"
+
+__all__ = [
+    "QuadEmbed",
+    "TextEncoder",
+    "VisionEncoder",
+    "AudioEncoder",
+    "VisionProjector",
+    "AudioProjector",
+    "matryoshka_infonce",
+    "bidirectional_infonce",
+    "embed_video",
+    "embed_video_from_frames",
+    "sample_frames",
+    "__version__",
+]
+
+_CHECKPOINTS = {
+    "vision": "checkpoints/vision_projector.pt",
+    "audio": "checkpoints/audio_projector.pt",
+    "video": "checkpoints/video_projector.pt",
+}
+
+
+def _default_device() -> str:
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class QuadEmbed:
+    """Three frozen encoders plus the trained projectors, with one `embed_*`
+    method per modality.
+
+    Every method returns an L2-normalized ``[batch, 768]`` tensor in the same
+    shared space, so any two outputs from any two modalities can be compared
+    with a plain matrix product (which is then cosine similarity).
+
+    Loading all three encoders takes a few GB of VRAM and a minute or two on
+    first run while the frozen backbones download from the Hub. Pass
+    ``modalities=`` to load only what you need.
+    """
+
+    def __init__(self, device: str | None = None, modalities=("text", "vision", "audio")):
+        self.device = device or _default_device()
+        self.modalities = tuple(modalities)
+
+        self.text_encoder = TextEncoder(self.device) if "text" in self.modalities else None
+
+        needs_vision = "vision" in self.modalities or "video" in self.modalities
+        self.vision_encoder = VisionEncoder(self.device) if needs_vision else None
+        self.audio_encoder = AudioEncoder(self.device) if "audio" in self.modalities else None
+
+        self.vision_projector = (
+            VisionProjector(patch_dim=768, out_dim=768).to(self.device).eval()
+            if "vision" in self.modalities else None
+        )
+        self.audio_projector = (
+            AudioProjector(in_dim=1280, out_dim=768).to(self.device).eval()
+            if "audio" in self.modalities else None
+        )
+        self.video_projector = (
+            VisionProjector(patch_dim=768, out_dim=768).to(self.device).eval()
+            if "video" in self.modalities else None
+        )
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id_or_path: str = DEFAULT_REPO_ID,
+        device: str | None = None,
+        modalities=("text", "vision", "audio"),
+        revision: str | None = None,
+    ) -> "QuadEmbed":
+        """Load trained projector weights.
+
+        ``repo_id_or_path`` is either a Hugging Face repo id (the default
+        pulls the published checkpoints) or a local directory containing a
+        ``checkpoints/`` folder in the same layout.
+        """
+        model = cls(device=device, modalities=modalities)
+
+        targets = [
+            ("vision", model.vision_projector),
+            ("audio", model.audio_projector),
+            ("video", model.video_projector),
+        ]
+        for name, projector in targets:
+            if projector is None:
+                continue
+            path = model._resolve(repo_id_or_path, _CHECKPOINTS[name], revision)
+            projector.load_state_dict(torch.load(path, map_location=model.device))
+        return model
+
+    @staticmethod
+    def _resolve(repo_id_or_path: str, filename: str, revision: str | None) -> str:
+        local = os.path.join(repo_id_or_path, filename)
+        if os.path.isfile(local):
+            return local
+        if os.path.isdir(repo_id_or_path):
+            raise FileNotFoundError(f"{local} not found in local directory {repo_id_or_path}")
+
+        from huggingface_hub import hf_hub_download
+
+        return hf_hub_download(repo_id=repo_id_or_path, filename=filename, revision=revision)
+
+    def _check(self, attr: str, modality: str):
+        value = getattr(self, attr)
+        if value is None:
+            raise RuntimeError(
+                f"'{modality}' was not loaded. Construct with "
+                f"modalities=(..., '{modality}') to enable it."
+            )
+        return value
+
+    @torch.no_grad()
+    def embed_text(self, texts: list[str]) -> torch.Tensor:
+        encoder = self._check("text_encoder", "text")
+        return F.normalize(encoder.embed(texts).float(), dim=-1)
+
+    @torch.no_grad()
+    def embed_image(self, images: list) -> torch.Tensor:
+        encoder = self._check("vision_encoder", "vision")
+        projector = self._check("vision_projector", "vision")
+        patch_tokens, mask, spatial_shapes = encoder.patch_tokens(images)
+        return F.normalize(projector(patch_tokens.float(), mask, spatial_shapes), dim=-1)
+
+    @torch.no_grad()
+    def embed_audio(self, arrays: list, sampling_rate: int = 16000) -> torch.Tensor:
+        """``arrays``: list of mono float32 numpy arrays at ``sampling_rate``
+        (16 kHz is what the frozen Whisper encoder expects)."""
+        encoder = self._check("audio_encoder", "audio")
+        projector = self._check("audio_projector", "audio")
+        frame_tokens = encoder.frame_tokens(arrays, sampling_rate=sampling_rate)
+        return F.normalize(projector(frame_tokens.float()), dim=-1)
+
+    @torch.no_grad()
+    def embed_video_file(self, video_path: str, num_frames: int = 4) -> torch.Tensor:
+        """Samples ``num_frames`` frames and mean-pools their embeddings.
+        Needs the video extra: ``pip install quadembed[video]``."""
+        encoder = self._check("vision_encoder", "video")
+        projector = self._check("video_projector", "video")
+        out = embed_video(encoder, projector, video_path, self.device, num_frames=num_frames)
+        return F.normalize(out.unsqueeze(0), dim=-1)

quadembed-0.1.0/quadembed/encoders.py

@@ -0,0 +1,103 @@
+"""Frozen encoder loaders.
+
+Substitutions vs. the GELATO paper's exact checkpoints, since the paper's
+adapted encoders (Qwen3.5's vision tower, Qwen2.5-Omni's audio tower) aren't
+distributed as standalone checkpoints, and the paper itself says they were
+*adapted from* these two source models -- so we use the source models
+directly, which the paper's own dimension numbers (1280 audio, patch-merge
+producing 3072 for nano) line up with:
+  text  : jinaai/jina-embeddings-v5-text-nano (exact match, publicly released)
+  vision: google/siglip2-base-patch16-naflex   (source SigLIP2 GELATO adapted)
+  audio : openai/whisper-large-v3 encoder      (source Whisper GELATO adapted)
+"""
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoModel,
+    AutoTokenizer,
+    AutoImageProcessor,
+    WhisperModel,
+    WhisperFeatureExtractor,
+)
+
+try:
+    from transformers import Siglip2VisionModel as _VisionModelClass
+except ImportError:
+    from transformers import AutoModel as _VisionModelClass
+
+TEXT_MODEL_ID = "jinaai/jina-embeddings-v5-text-nano"
+VISION_MODEL_ID = "google/siglip2-base-patch16-naflex"
+AUDIO_MODEL_ID = "openai/whisper-large-v3"
+
+
+def _freeze(module: nn.Module) -> nn.Module:
+    module.eval()
+    for p in module.parameters():
+        p.requires_grad_(False)
+    return module
+
+
+def mean_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    mask = attention_mask.unsqueeze(-1).to(last_hidden_state.dtype)
+    summed = (last_hidden_state * mask).sum(dim=1)
+    counts = mask.sum(dim=1).clamp(min=1e-9)
+    return summed / counts
+
+
+class TextEncoder:
+    def __init__(self, device: str, dtype=torch.bfloat16, max_length: int = 512):
+        self.tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID, trust_remote_code=True)
+        self.model = _freeze(
+            AutoModel.from_pretrained(TEXT_MODEL_ID, trust_remote_code=True, dtype=dtype).to(device)
+        )
+        self.device = device
+        self.max_length = max_length
+
+    @torch.no_grad()
+    def embed(self, texts: list[str]) -> torch.Tensor:
+        inputs = self.tokenizer(
+            texts, padding=True, truncation=True, max_length=self.max_length, return_tensors="pt"
+        ).to(self.device)
+        out = self.model(**inputs)
+        return mean_pool(out.last_hidden_state, inputs["attention_mask"])
+
+
+class VisionEncoder:
+    def __init__(self, device: str, dtype=torch.bfloat16):
+        self.processor = AutoImageProcessor.from_pretrained(VISION_MODEL_ID)
+        self.model = _freeze(_VisionModelClass.from_pretrained(VISION_MODEL_ID, dtype=dtype).to(device))
+        self.device = device
+
+    @torch.no_grad()
+    def patch_tokens(self, images: list) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        """Returns (patch_tokens [B,N,D], attention_mask [B,N] or None,
+        spatial_shapes [B,2] or None -- per-image (H_patches, W_patches)
+        before NaFlex pads every image in the batch to the same N, needed
+        for a real spatial 2x2 merge instead of grouping arbitrary tokens."""
+        inputs = self.processor(images=images, return_tensors="pt").to(self.device)
+        vision_module = getattr(self.model, "vision_model", self.model)
+        out = vision_module(**inputs)
+        mask = inputs.get("pixel_attention_mask")
+        spatial_shapes = inputs.get("spatial_shapes")
+        return out.last_hidden_state, mask, spatial_shapes
+
+
+class AudioEncoder:
+    def __init__(self, device: str, dtype=torch.bfloat16):
+        self.feature_extractor = WhisperFeatureExtractor.from_pretrained(AUDIO_MODEL_ID)
+        full = WhisperModel.from_pretrained(AUDIO_MODEL_ID, dtype=dtype)
+        self.model = full.encoder
+        del full.decoder
+        del full
+        self.model = _freeze(self.model.to(device))
+        self.device = device
+        self.dtype = dtype
+
+    @torch.no_grad()
+    def frame_tokens(self, audio_arrays: list, sampling_rate: int = 16000) -> torch.Tensor:
+        inputs = self.feature_extractor(
+            audio_arrays, sampling_rate=sampling_rate, return_tensors="pt"
+        ).to(self.device)
+        features = inputs["input_features"].to(self.dtype)
+        out = self.model(features)
+        return out.last_hidden_state

quadembed-0.1.0/quadembed/losses.py

@@ -0,0 +1,30 @@
+"""Bidirectional in-batch InfoNCE + Matryoshka loss, matching GELATO's Stage
+1/2 objective (temperature 0.02, prefix dims for the nano joint space)."""
+import torch
+import torch.nn.functional as F
+
+NANO_MATRYOSHKA_DIMS = (32, 64, 128, 256, 768)
+
+
+def bidirectional_infonce(a: torch.Tensor, b: torch.Tensor, temperature: float = 0.02) -> torch.Tensor:
+    """a, b: [B, D] paired embeddings (e.g. text, image), same batch order = positives."""
+    a = F.normalize(a, dim=-1)
+    b = F.normalize(b, dim=-1)
+    logits = a @ b.T / temperature  # [B, B]
+    labels = torch.arange(a.shape[0], device=a.device)
+    loss_a2b = F.cross_entropy(logits, labels)
+    loss_b2a = F.cross_entropy(logits.T, labels)
+    return (loss_a2b + loss_b2a) / 2
+
+
+def matryoshka_infonce(a: torch.Tensor, b: torch.Tensor, dims=NANO_MATRYOSHKA_DIMS, temperature: float = 0.02) -> torch.Tensor:
+    """Sum of bidirectional InfoNCE over truncated embedding prefixes, so
+    every prefix length is independently a usable (Matryoshka-truncatable)
+    embedding, not just the full 768-dim vector."""
+    losses = []
+    full_dim = a.shape[-1]
+    for d in dims:
+        if d > full_dim:
+            continue
+        losses.append(bidirectional_infonce(a[..., :d], b[..., :d], temperature))
+    return torch.stack(losses).mean()

quadembed-0.1.0/quadembed/projectors.py

@@ -0,0 +1,111 @@
+"""Trainable projectors that map frozen vision/audio encoder outputs into the
+frozen text encoder's embedding space -- the only parameters GELATO trains.
+
+Everything upstream (SigLIP2, Whisper encoder, jina-v5-text-nano) stays frozen;
+only VisionProjector.fc_vision_2 and AudioProjector.fc_audio get gradients.
+"""
+import torch
+import torch.nn as nn
+
+
+class VisionProjector(nn.Module):
+    """LayerNorm -> real 2x2 spatial merge -> trainable linear, matching
+    GELATO's "LayerNorm, 2x2 spatial merge, fc_vision_2" description
+    (nano: 3072->768).
+
+    Uses NaFlex's per-image `spatial_shapes` (H_patches, W_patches) to group
+    each image's own patch grid into genuine 2x2 spatial blocks before
+    concatenating channel-wise -- earlier versions approximated this with
+    sequential 4-token grouping (mixing patches from unrelated rows), which
+    is architecturally wrong and is kept here only as a fallback for
+    callers that don't have spatial_shapes available.
+    """
+
+    def __init__(self, patch_dim: int, out_dim: int, merge: int = 4, hidden_dim: int | None = None):
+        super().__init__()
+        self.merge = merge
+        self.norm = nn.LayerNorm(patch_dim, elementwise_affine=False)
+        if hidden_dim:
+            # not part of GELATO's own recipe (which uses a single linear
+            # fc_vision_2) -- an experiment to test whether a linear
+            # projector's capacity, not data, is capping R@1. See README's
+            # "round 6" for the result.
+            self.fc_vision_2 = nn.Sequential(
+                nn.Linear(patch_dim * merge, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, out_dim)
+            )
+        else:
+            self.fc_vision_2 = nn.Linear(patch_dim * merge, out_dim)
+        self.modality_delim = nn.Parameter(torch.zeros(out_dim))
+
+    def _spatial_merge_one(self, tokens: torch.Tensor, h: int, w: int) -> torch.Tensor:
+        # tokens: [N, D] for one image, first h*w entries are the valid grid
+        d = tokens.shape[-1]
+        h2, w2 = h - (h % 2), w - (w % 2)  # drop a trailing odd row/col rather than fabricate one
+        grid = tokens[: h * w].view(h, w, d)[:h2, :w2]
+        blocks = grid.reshape(h2 // 2, 2, w2 // 2, 2, d).permute(0, 2, 1, 3, 4).reshape(-1, 4 * d)
+        return blocks
+
+    def forward(
+        self,
+        patch_tokens: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        spatial_shapes: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # patch_tokens: [B, N, patch_dim]
+        b, n, d = patch_tokens.shape
+        patch_tokens = self.norm(patch_tokens)
+
+        if spatial_shapes is not None:
+            pooled = []
+            for i in range(b):
+                h, w = int(spatial_shapes[i, 0]), int(spatial_shapes[i, 1])
+                blocks = self._spatial_merge_one(patch_tokens[i], h, w)
+                projected = self.fc_vision_2(blocks)  # [num_blocks, out_dim]
+                pooled.append(projected.mean(dim=0))
+            pooled = torch.stack(pooled, dim=0)
+        else:
+            pad = (-n) % self.merge
+            pt = patch_tokens
+            if pad:
+                pt = nn.functional.pad(pt, (0, 0, 0, pad))
+                if attention_mask is not None:
+                    attention_mask = nn.functional.pad(attention_mask, (0, pad))
+            merged = pt.reshape(b, pt.shape[1] // self.merge, d * self.merge)
+            projected = self.fc_vision_2(merged)  # [B, N/merge, out_dim]
+            if attention_mask is not None:
+                merged_mask = attention_mask.reshape(b, -1, self.merge).amax(dim=-1).to(projected.dtype)
+                weights = merged_mask.unsqueeze(-1)
+                pooled = (projected * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1e-6)
+            else:
+                pooled = projected.mean(dim=1)
+
+        return pooled + self.modality_delim
+
+
+class AudioProjector(nn.Module):
+    """Trainable linear (1280->768 for nano), matching GELATO's fc_audio."""
+
+    def __init__(self, in_dim: int, out_dim: int):
+        super().__init__()
+        self.fc_audio = nn.Linear(in_dim, out_dim)
+        self.modality_delim = nn.Parameter(torch.zeros(out_dim))
+
+    def forward(self, frame_tokens: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        # frame_tokens: [B, T, in_dim]
+        projected = self.fc_audio(frame_tokens)
+        if attention_mask is not None:
+            weights = attention_mask.unsqueeze(-1).to(projected.dtype)
+            pooled = (projected * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1e-6)
+        else:
+            pooled = projected.mean(dim=1)
+        return pooled + self.modality_delim
+
+
+def embed_video_from_frames(vision_projector: VisionProjector, per_frame_patch_tokens: list[torch.Tensor]) -> torch.Tensor:
+    """Video = mean of per-frame vision-projector embeddings.
+
+    No dedicated video encoder, matching GELATO: sample frames, run each
+    through the (shared, already-trained) vision projector, pool over time.
+    """
+    frame_embeds = torch.stack([vision_projector(tokens) for tokens in per_frame_patch_tokens], dim=1)  # [B, F, out_dim]
+    return frame_embeds.mean(dim=1)

quadembed-0.1.0/quadembed/video.py

@@ -0,0 +1,47 @@
+"""Frame sampling + video embedding. No video encoder exists here -- a video
+is just N sampled frames run through the already-trained VisionProjector and
+mean-pooled, exactly as GELATO does."""
+import numpy as np
+import torch
+from PIL import Image
+
+
+def _require_cv2():
+    try:
+        import cv2
+    except ImportError as e:  # pragma: no cover
+        raise ImportError(
+            "Video support needs OpenCV. Install it with:\n"
+            "    pip install quadembed[video]"
+        ) from e
+    return cv2
+
+
+def sample_frames(video_path: str, num_frames: int = 4) -> list:
+    cv2 = _require_cv2()
+    cap = cv2.VideoCapture(video_path)
+    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    if total <= 0:
+        cap.release()
+        raise ValueError(f"no frames read from {video_path}")
+    indices = np.linspace(0, total - 1, num=min(num_frames, total), dtype=int)
+    frames = []
+    for idx in indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
+        ok, frame_bgr = cap.read()
+        if not ok:
+            continue
+        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+        frames.append(Image.fromarray(frame_rgb))
+    cap.release()
+    if not frames:
+        raise ValueError(f"no frames decoded from {video_path}")
+    return frames
+
+
+@torch.no_grad()
+def embed_video(vision_encoder, vision_projector, video_path: str, device: str, num_frames: int = 4) -> torch.Tensor:
+    frames = sample_frames(video_path, num_frames=num_frames)
+    patch_tokens, mask, spatial_shapes = vision_encoder.patch_tokens(frames)  # [num_frames, N, 768]
+    frame_embeds = vision_projector(patch_tokens.float(), mask, spatial_shapes)  # [num_frames, 768]
+    return frame_embeds.mean(dim=0)