llm-semantic-router/multi-modal-embed-small

pip install torch transformers pillow safetensors

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer, SiglipModel, SiglipProcessor, WhisperModel, WhisperFeatureExtractor
from huggingface_hub import hf_hub_download

class MultiModalEmbedder(nn.Module):
    """Standalone multimodal embedder - no external dependencies."""
    
    def __init__(self):
        super().__init__()
        # Text encoder (384d, no projection needed)
        self.text_tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        self.text_encoder = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        
        # Image encoder (768d -> 384d projection)
        self.image_processor = SiglipProcessor.from_pretrained("google/siglip-base-patch16-512")
        self.image_encoder = SiglipModel.from_pretrained("google/siglip-base-patch16-512").vision_model
        self.image_proj = nn.Linear(768, 384)
        
        # Audio encoder (384d, no projection needed)
        self.audio_processor = WhisperFeatureExtractor.from_pretrained("openai/whisper-tiny")
        self.audio_encoder = WhisperModel.from_pretrained("openai/whisper-tiny").encoder
    
    def encode_text(self, texts):
        if isinstance(texts, str):
            texts = [texts]
        inputs = self.text_tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.text_encoder(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
        return F.normalize(embeddings, p=2, dim=-1)
    
    def encode_image(self, images):
        inputs = self.image_processor(images=images, return_tensors="pt")
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.image_encoder(**inputs)
        embeddings = outputs.pooler_output
        embeddings = self.image_proj(embeddings)  # 768 -> 384
        return F.normalize(embeddings, p=2, dim=-1)
    
    def encode_audio(self, waveform):
        # waveform: numpy array or tensor at 16kHz
        if isinstance(waveform, torch.Tensor):
            waveform = waveform.squeeze().numpy()
        inputs = self.audio_processor(waveform, sampling_rate=16000, return_tensors="pt")
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.audio_encoder(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
        return F.normalize(embeddings, p=2, dim=-1)

# Load model
model = MultiModalEmbedder()

# Download and load trained weights
checkpoint_path = hf_hub_download(
    repo_id="llm-semantic-router/multi-modal-embed-small",
    filename="model.pt"
)
state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False)

# Load text encoder weights
model.text_encoder.load_state_dict({
    k.replace("text_encoder.encoder.", ""): v 
    for k, v in state_dict.items() 
    if k.startswith("text_encoder.encoder.")
})

# Load image encoder and projection weights
model.image_encoder.load_state_dict({
    k.replace("image_encoder.vision_encoder.", ""): v 
    for k, v in state_dict.items() 
    if k.startswith("image_encoder.vision_encoder.")
})
model.image_proj.load_state_dict({
    k.replace("image_encoder.projection.", ""): v 
    for k, v in state_dict.items() 
    if k.startswith("image_encoder.projection.")
})

# Load audio encoder weights
model.audio_encoder.load_state_dict({
    k.replace("audio_encoder.encoder.", ""): v 
    for k, v in state_dict.items() 
    if k.startswith("audio_encoder.encoder.")
})

model.eval()
print("Model loaded successfully!")

import torch.nn.functional as F

# Single text
text_embedding = model.encode_text("A photo of a cat")  # Shape: [1, 384]

# Batch of texts
texts = ["A fluffy orange cat", "A golden retriever dog", "A red sports car"]
text_embeddings = model.encode_text(texts)  # Shape: [3, 384]

# Compute similarity
similarities = F.cosine_similarity(text_embeddings[0:1], text_embeddings[1:], dim=-1)
print(f"Cat vs Dog: {similarities[0]:.3f}")
print(f"Cat vs Car: {similarities[1]:.3f}")

from PIL import Image
import requests
from io import BytesIO

# Load image
url = "https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg"
image = Image.open(BytesIO(requests.get(url).content)).convert('RGB')

# Get embedding
image_embedding = model.encode_image(image)  # Shape: [1, 384]

import torchaudio

# Load audio (16kHz)
waveform, sample_rate = torchaudio.load("speech.wav")
if sample_rate != 16000:
    waveform = torchaudio.functional.resample(waveform, sample_rate, 16000)

# Get embedding
audio_embedding = model.encode_audio(waveform)  # Shape: [1, 384]

# Image-to-text retrieval
image = Image.open("cat.jpg").convert('RGB')
image_emb = model.encode_image(image)

captions = [
    "A cat sleeping on a bed",
    "A dog playing in the park",
    "A car driving on the highway",
]
text_embs = model.encode_text(captions)

similarities = F.cosine_similarity(image_emb, text_embs)
best_idx = similarities.argmax().item()
print(f"Best match: {captions[best_idx]} ({similarities[best_idx]:.3f})")

# Full 384-dim embedding
full_emb = model.encode_text("Hello world")  # [1, 384]

# Truncate to smaller dimensions
emb_256 = F.normalize(full_emb[:, :256], p=2, dim=-1)  # 1.5x faster retrieval
emb_128 = F.normalize(full_emb[:, :128], p=2, dim=-1)  # 3x faster retrieval
emb_64 = F.normalize(full_emb[:, :64], p=2, dim=-1)    # 6x faster retrieval

┌──────────────────────────────────────────────────────────────┐
│                  multi-modal-embed-small                     │
├──────────────────────────────────────────────────────────────┤
│  Text Encoder:  MiniLM-L6-v2           (22M params, 6 layers)│
│  Image Encoder: SigLIP-base-patch16-512 (86M params)         │
│  Audio Encoder: Whisper-tiny encoder    (8M params, 4 layers) │
│  Fusion:        2-layer Transformer                          │
├──────────────────────────────────────────────────────────────┤
│  Output: 384-dim normalized embeddings                       │
│  2DMSE:  Layer 0-5 early exit support                        │
│  MRL:    32, 64, 128, 256, 384 dim truncation                │
└──────────────────────────────────────────────────────────────┘

@misc{multi-modal-embed-small-2026,
  title={multi-modal-embed-small: Compact Multimodal Embeddings with 2DMSE},
  author={Semantic Router Team},
  year={2026},
  url={https://huggingface.co/llm-semantic-router/multi-modal-embed-small}
}