jinaai/jina-embeddings-v5-omni-nano-retrieval-mlx

<br><br>

<p align="center"> <img src="https://huggingface.co/datasets/jinaai/documentation-images/resolve/main/logo.webp" alt="Jina AI: Your Search Foundation, Supercharged!" width="150px"> </p>

jina-embeddings-v5-omni-nano-retrieval-mlx: Retrieval-Targeted Omni Embedding (Nano) — MLX

ArXiv | Blog

<p align="center"> <img src="omni_frontier.png" alt="Average score vs. parameter count for open-weight omni embedding models" width="520px"> </p>

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](https://arxiv.org/abs/2605.08384)).

Model Overview

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

This is the retrieval-targeted variant of the jina-embeddings-v5-omni-nano MLX family.

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.

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

# Generate an embedding for a single piece of text using the predefined endpoint
POST _inference/embedding/.jina-embeddings-v5-omni-nano
{
  "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-nano
{
  "input": [
    {
      "content": [
        { "type": "text",  "value": "A small blue square" },
        { "type": "image", "format": "base64", "value": "<BASE64_IMAGE_DATA>" }
      ]
    }
  ]
}

# Create a custom endpoint that truncates omni-nano embeddings to 32 dimensions
PUT _inference/embedding/jina-omni-nano-32d
{
  "service": "elastic",
  "service_settings": {
    "model_id": "jina-embeddings-v5-omni-nano",
    "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

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-nano-retrieval-mlx"))

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

cfg = OmniNanoConfig.from_dict(json.loads((repo_dir / "config.json").read_text()))
model = JinaOmniNanoEmbeddingModel(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("Query: Which planet is known as the Red Planet?")
d = embed_text("Document: Mars is often referred to as the Red Planet.")
cos = float((q[0] * d[0]).sum() / (mx.linalg.norm(q[0]) * mx.linalg.norm(d[0])))
print(f"cos = {cos:.4f}")

Prefix convention. For retrieval, prepend Query: to query-side text and Document: to document-side text — these prefixes are required to match the torch reference (they correspond to encode_query() / encode_document() in the HF transformers / sentence-transformers integrations).

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

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="<image>", 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

Nano shares its image token (<image>) for video. The processor's video path returns pixel_values_videos and video_grid_thw; rename them to the image keys (the canonical customst flow) so the same `encodeimage path handles temporal pairs (qwen3vl's temporalpatchsize=2` Conv3d):

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="<image>", 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())

video_emb = model.encode_image(pixel_values, grid_thw, input_ids, attn)

Quickstart — text + audio

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, _ = model.audio_tower.feat_extract_output_lengths(feature_lens)
# avg-pool downsamples 2x; one <|AUDIO|> token per output position
n_audio_tokens = int(aftercnn.sum().item()) // 2
AUDIO_TOKEN_ID = 128256
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_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}. 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.

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:

import mlx.core as mx

texts = ["Query: query 1", "Query: query 2"]
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, 768)

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).

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

Compatibility

Same vector space as:

• —`jinaai/jina-embeddings-v5-text-nano-retrieval` — text-only
• —`jinaai/jina-embeddings-v5-omni-nano-retrieval` — multimodal (transformers / ST / vLLM)
• —`jinaai/jina-embeddings-v5-omni-nano-retrieval-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.