You can deploy this model and start earning money today!

jinaai/
jina-embeddings-v5-omni-small-text-matching-mlx

Name: jina-embeddings-v5-omni-small-text-matching-mlx
Availability: InStock

Views

No views yet

Jina AI: Your Search Foundation, Supercharged!

jina-embeddings-v5-omni-small-text-matching-mlx: Text-Matching-Targeted Omni Embedding (Small) — MLX

ArXiv | Blog

Average score vs. parameter count for open-weight omni embedding models

Average score vs. parameter count across image (MIEB-Lite), video (MMEB-V), and audio (MAEB) benchmarks — jina-v5-omni-nano and jina-v5-omni-small define the open-weight frontier (Table 1 in the ArXiv report).

Model Overview

MLX-native build of jinaai/jina-embeddings-v5-omni-small-text-matching for Apple Silicon (M1/M2/M3/M4) inference. Accepts text, images, video, and audio and produces 1024-dim embeddings in the same vector space as the torch reference and as jinaai/jina-embeddings-v5-text-small-text-matching at the same task — index with text and query with any modality, no reindexing. For a more compact alternative, see jinaai/jina-embeddings-v5-omni-nano-text-matching-mlx.

This is the text-matching-targeted variant of the jina-embeddings-v5-omni-small MLX family.

Feature	Value
Parameters	~2.0B (text + vision + audio towers)
Embedding Dimension	1024
Supported Tasks	`text-matching`
Max Sequence Length	32768
Pooling Strategy	Last-token
Supported Inputs	text, image, video, audio
Supported File Types	images: `.jpg`, `.jpeg`, `.png`, `.gif`, `.webp`, `.bmp`, `.tif`, `.tiff`, `.avif`, `.heic`, `.svg`; video: `.mp4`, `.avi`, `.mov`, `.mkv`, `.webm`, `.flv`, `.wmv`; audio: `.wav`, `.mp3`, `.flac`, `.ogg`, `.m4a`, `.opus`; documents: `.pdf`
Matryoshka Dimensions	32, 64, 128, 256, 512, 768, 1024
Precision	bf16 (vision + audio), fp32 (`language_model`)

Via Elastic Inference Service

The fastest way to use v5-omni in production. Elastic Inference Service (EIS) provides managed embedding inference with built-in scaling, so you can generate embeddings directly within your Elastic deployment.

bash

# Retrieve the configuration of the preconfigured omni-small inference endpoint
GET /_inference/embedding/.jina-embeddings-v5-omni-small

# Generate an embedding for a single piece of text using the predefined endpoint
POST _inference/embedding/.jina-embeddings-v5-omni-small
{
  "input": [
    "This is a test"
  ]
}

# Fuse a text description and an image into a single embedding via a multimodal content block
POST _inference/embedding/.jina-embeddings-v5-omni-small
{
  "input": [
    {
      "content": [
        { "type": "text",  "value": "A small blue square" },
        { "type": "image", "format": "base64", "value": "<BASE64_IMAGE_DATA>" }
      ]
    }
  ]
}

# Create a custom endpoint that truncates omni-small embeddings to 32 dimensions
PUT _inference/embedding/jina-omni-small-32d
{
  "service": "elastic",
  "service_settings": {
    "model_id": "jina-embeddings-v5-omni-small",
    "dimensions": 32
  }
}

See the Elastic Inference Service documentation for setup details.

Install MLX

pip install mlx tokenizers huggingface_hub transformers pillow requests librosa av

Apple Silicon only. mlx >= 0.23 recommended; transformers >= 4.57 for the processors used in the image / video / audio quickstarts; av for video decoding.

Quickstart — text

python

import json
from pathlib import Path

import mlx.core as mx
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer

repo_dir = Path(snapshot_download("jinaai/jina-embeddings-v5-omni-small-text-matching-mlx"))

import sys
sys.path.insert(0, str(repo_dir))
from model import JinaOmniSmallEmbeddingModel, OmniSmallConfig  # type: ignore

cfg = OmniSmallConfig.from_dict(json.loads((repo_dir / "config.json").read_text()))
model = JinaOmniSmallEmbeddingModel(cfg)
model.load_weights(str(repo_dir / "model.safetensors"))
mx.eval(model.parameters())

tok = Tokenizer.from_file(str(repo_dir / "tokenizer.json"))


def embed_text(text: str):
    enc = tok.encode(text)
    input_ids = mx.array([enc.ids])
    attn = mx.array([enc.attention_mask])
    return model.encode_text(input_ids, attn)


q = embed_text("a photo of a cat sitting on a couch")
d = embed_text("an image of a feline lounging on furniture")
cos = float((q[0] * d[0]).sum() / (mx.linalg.norm(q[0]) * mx.linalg.norm(d[0])))
print(f"cos = {cos:.4f}")

No prefix convention. Inputs are embedded verbatim — no Query: / Document: prefixes are needed (unlike the retrieval variant). Both sides of any text pair go in unprefixed.

No custom pooling or padding code needed — the MLX model exposes per-modality encode_* methods that internally apply last-token pooling and L2 normalization; min_pixels / max_pixels / temporal_patch_size come from the bundled processor and model.safetensors metadata.

Quickstart — text + image

python

from io import BytesIO
import requests
from PIL import Image
from transformers import AutoProcessor

proc = AutoProcessor.from_pretrained(str(repo_dir), trust_remote_code=True)


def embed_image(image):
    inputs = proc(images=[image], text="<|vision_start|><|image_pad|><|vision_end|>", return_tensors="pt")
    pixel_values = mx.array(inputs["pixel_values"].numpy())
    grid_thw     = mx.array(inputs["image_grid_thw"].numpy())
    input_ids    = mx.array(inputs["input_ids"].numpy())
    attn         = mx.array(inputs["attention_mask"].numpy())
    return model.encode_image(pixel_values, grid_thw, input_ids, attn)


url = "https://upload.wikimedia.org/wikipedia/commons/thumb/0/02/OSIRIS_Mars_true_color.jpg/800px-OSIRIS_Mars_true_color.jpg"
image = Image.open(BytesIO(requests.get(url).content)).convert("RGB")
image_emb = embed_image(image)

Quickstart — text + video

Video is encoded as a sequence of temporal-pair frames (Qwen3VL's temporal_patch_size=2 Conv3d patch_embed). Decode the clip, group consecutive frames into pairs, and feed them through the processor's video path:

python

import av  # pip install av

def decode_video(path):
    container = av.open(path)
    return [frame.to_image().convert("RGB") for frame in container.decode(video=0)]

frames = decode_video("clip.mp4")
inputs = proc(text="<|vision_start|><|video_pad|><|vision_end|>", videos=frames, return_tensors="pt")
pixel_values = mx.array(inputs["pixel_values_videos"].numpy())
grid_thw     = mx.array(inputs["video_grid_thw"].numpy())     # [T, H, W] with T = num_frames // 2
input_ids    = mx.array(inputs["input_ids"].numpy())
attn         = mx.array(inputs["attention_mask"].numpy())

# Same vision forward as image; encode_video scatters at video_token_id
# (distinct from image_token_id in small). The Conv3d patch_embed consumes the temporal axis.
video_emb = model.encode_video(pixel_values, grid_thw, input_ids, attn)

Quickstart — text + audio

python

import librosa
from transformers import WhisperFeatureExtractor

audio_path = "speech.wav"
wav, _ = librosa.load(audio_path, sr=16000, mono=True)

fx = WhisperFeatureExtractor(feature_size=128, sampling_rate=16000)
feats = fx(wav, sampling_rate=16000, return_tensors="np")
mel = feats["input_features"][0]
feat_mask = feats.get("attention_mask")
feat_len = int(feat_mask[0].sum()) if feat_mask is not None else mel.shape[-1]

input_features = mx.array(mel)
feature_lens = mx.array([feat_len])

aftercnn, output_lens = model.audio_tower.feat_extract_output_lengths(feature_lens)
# avg-pool downsamples 2x; one <|audio_pad|> token per output position
n_audio_tokens = int(output_lens.sum().item())
AUDIO_TOKEN_ID = 151669
input_ids = mx.array([[AUDIO_TOKEN_ID] * n_audio_tokens])
attn = mx.ones_like(input_ids)

audio_emb = model.encode_audio(input_features, feature_lens, input_ids, attn)

WAV / MP3 / FLAC are accepted; audio is resampled to 16kHz mono. Feature extraction stays torch-side (WhisperFeatureExtractor) — only the audio tower forward is ported to MLX.

Selective modality loading (text / vision / audio / omni)

Each modality has its own forward method (encode_text, encode_image, encode_video, encode_audio) on the same model.safetensors. Call only the methods you need — there's no separate modality= flag because the model is loaded once and the unused encoders simply aren't invoked. This mirrors the HF modality= argument from a usage standpoint without requiring a second weights file or a per-modality build.

Matryoshka (truncating embeddings)

Any prefix of the output vector is itself a valid embedding once L2-renormalized. Supported prefix dims: {32, 64, 128, 256, 512, 768, 1024}. Matryoshka is a property of the trained projection head — verified end-to-end through the GGUF F16 path with 0.0000 prefix-vs-full drift on the same 7-input reference set; MLX uses the same weights at bf16, so the structure is preserved by construction.

python

import numpy as np
full = np.array(q[0].astype(mx.float32))
truncated = full[:256]
truncated /= np.linalg.norm(truncated)

Batching

MLX batches by stacking inputs along the leading axis. For text, pad to a common length and stack input_ids / attention_mask:

python

import mlx.core as mx

texts = ["first sentence", "second sentence"]
encs = [tok.encode(t) for t in texts]
max_len = max(len(e.ids) for e in encs)
pad = tok.token_to_id("<|endoftext|>") or 0
input_ids = mx.array([e.ids + [pad] * (max_len - len(e.ids)) for e in encs])
attn      = mx.array([e.attention_mask + [0] * (max_len - len(e.attention_mask)) for e in encs])
embs = model.encode_text(input_ids, attn)   # (2, 1024)

Multimodal inputs are forwarded per-sample (one call per image / video / audio). For high-throughput multimodal serving on non-Mac hardware, prefer the vLLM path on the torch base model.

Multimodal parity vs torch (cos ≥ 0.99 numerical bar)

Verified on Apple Silicon against the same 6-modality fixture set used for the GGUF parity tables (text, image car/cat, audio JFK 11s, PDF 2-page fused, video 4-frame 512²). Torch reference: model.embed() at fp32. MLX side: bf16 vision/audio + fp32 language_model (upcast at load time).

Modality	small-text-matching
Text	7/7 inputs ≥ 0.999
Image (car)	0.9985
Image (cat)	0.9993
Audio (JFK 11s)	1.0000
PDF (2-page fused)	0.9996
Video (4-frame, 512²)	0.9973

To reproduce, run scripts/omni/mlx/test_mlx_full_parity.py --family small --variant text-matching on Apple Silicon (script + torch reference JSON shipped from the v5-omni training repo).

Compatibility

Same vector space as:

jinaai/jina-embeddings-v5-text-small-text-matching — text-only
jinaai/jina-embeddings-v5-omni-small-text-matching — multimodal (transformers / ST / vLLM)
jinaai/jina-embeddings-v5-omni-small-text-matching-GGUF — multimodal (llama.cpp)

Notes

Last-token pooling is used throughout (matches the torch reference).
Audio feature extraction stays torch-side (WhisperFeatureExtractor) and is not ported to MLX. Pass the extracted mel features in.
Precision note. language_model weights are upcast to fp32 at load time (the safetensors on disk stay bf16); vision and audio towers stay at bf16. This matches torch torch_dtype=torch.float32 and keeps short-multilingual text above the cos ≥ 0.99 floor (bf16 alone drifts ~0.97 on inputs like "Bonjour, comment ça va?" on some variants due to the 7-bit mantissa). Vision tokens still see kernel-level fp differences accumulating over long sequences (~0.998-0.9999 cos vs torch fp32).

License

CC BY-NC 4.0. For commercial use, contact us.