jinaai/jina-embeddings-v5-omni-small-classification-mlx

# 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
  }
}

pip install mlx tokenizers huggingface_hub transformers pillow requests librosa av

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-classification-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}")

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)

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)

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)

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

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)