jina-embeddings-v5-omni-small on AXERA NPU

Ready-to-run AX650 package for jinaai/jina-embeddings-v5-omni-small.

This repository contains the compiled AX650 .axmodel files, tokenizer files, embedding weight binary, sample assets, helper scripts, and an axllm runtime binary. Runtime inference does not require the original Hugging Face safetensors files.

The validated API is OpenAI-compatible /v1/embeddings for:

• Text embedding
• Single-image embedding
• 8-second audio embedding
• Frame-directory video embedding

Supported Platform

• AX650 / NPU3

Download

mkdir -p AXERA-TECH/jina-embeddings-v5-omni-small
cd AXERA-TECH/jina-embeddings-v5-omni-small
hf download AXERA-TECH/jina-embeddings-v5-omni-small --local-dir .

Package Layout

.
├── README.md
├── config.json
├── bin/axllm
├── jina_v5_omni_tokenizer/
├── jina_v5_omni_tokenizer.txt
├── model.embed_tokens.weight.bfloat16.bin
├── qwen3_p256_l0_together.axmodel
├── ...
├── qwen3_p256_l27_together.axmodel
├── qwen3_post.axmodel
├── jina_v5_omni_vision_256x256.axmodel
├── jina_v5_omni_audio_8s.axmodel
├── jina_v5_omni_vision_448x448.axmodel
├── jina_v5_omni_audio_30s.axmodel
├── python/
└── assets/

Start the Service

Run on the AX650 board from the package root:

chmod +x ./bin/axllm
export LD_LIBRARY_PATH=/soc/lib:${LD_LIBRARY_PATH:-}
./bin/axllm serve . --port 8000

Health checks:

curl http://127.0.0.1:8000/health
curl http://127.0.0.1:8000/v1/models

Expected model id:

AXERA-TECH/jina-embeddings-v5-omni-small-AX650-P256-MG-CTX2047

OpenAI-Compatible Examples

Text:

python3 python/openai_embedding_demo.py \
  --api-url http://127.0.0.1:8000/v1 \
  --model AXERA-TECH/jina-embeddings-v5-omni-small-AX650-P256-MG-CTX2047 \
  --prompt-name query \
  --input "Which planet is known as the Red Planet?"

Image:

python3 python/openai_multimodal_embedding_demo.py \
  --api-url http://127.0.0.1:8000/v1 \
  --model AXERA-TECH/jina-embeddings-v5-omni-small-AX650-P256-MG-CTX2047 \
  --prompt-name query \
  --media-type image \
  --media-path assets/sample.png

Audio:

python3 python/openai_multimodal_embedding_demo.py \
  --api-url http://127.0.0.1:8000/v1 \
  --model AXERA-TECH/jina-embeddings-v5-omni-small-AX650-P256-MG-CTX2047 \
  --prompt-name query \
  --media-type audio \
  --media-path assets/audio_test_chunk0_8s.wav

Video:

python3 python/openai_multimodal_embedding_demo.py \
  --api-url http://127.0.0.1:8000/v1 \
  --model AXERA-TECH/jina-embeddings-v5-omni-small-AX650-P256-MG-CTX2047 \
  --prompt-name query \
  --media-type video \
  --media-path assets/red-panda-openai.frames

The validated video path is a directory of pre-extracted frames. If you want to use a video file, extract frames first and pass the frame directory to the API.

Board Precision

The table below compares board-side axllm serve embeddings with the packaged Hugging Face reference embeddings under python/testdata/service_cases/*/torch_embedding.npy.

Run the packaged validation script after starting the service:

python3 python/compare_openai_api_vs_hf_multimodal.py \
  --api-url http://127.0.0.1:8000/v1 \
  --api-package-root .

Modality	Case	Output shape	Soft tokens	Cosine vs HF
Text document	embedding_doc	[1, 1024]	-	0.998966
Text query	red_planet_query	[1, 1024]	-	0.999144
Image	vision_sample	[1, 1024]	64	0.995698
Audio	audio_test_chunk0_8s_wav	[1, 1024]	200	0.996746
Video	video_visual_red_panda_openai_mp4	[1, 1024]	192	0.996037

Performance

This model returns embeddings and does not run a token-by-token decode loop. The useful runtime metric is media preparation plus LLM prefill.

The table below was measured on AX650 with the default config.json profile: 256x256 vision encoder and 8s audio encoder. The audio row uses the shipped assets/audio_test_chunk0_8s.wav, which is 16kHz mono PCM WAV. This package validates and recommends 16kHz mono PCM WAV for audio input.

Scenario	Prompt	LLM input tokens	Soft tokens	Output shape	Media prepare	LLM prefill	Runtime total
Text document	document	18	-	[1, 1024]	-	607.85 ms	608.51 ms
Text query	query	11	-	[1, 1024]	-	593.91 ms	594.16 ms
Image	query	74	64	[1, 1024]	167.24 ms	582.25 ms	749.49 ms
Audio (8s, 16kHz mono PCM WAV)	query	210	200	[1, 1024]	804.51 ms	577.96 ms	1382.48 ms
Video (3 frames)	query	202	192	[1, 1024]	300.79 ms	577.28 ms	878.08 ms

Standalone encoder latency measured with ax_run_model -r 50 -w 10:

Encoder axmodel	Output tokens	Avg latency
jina_v5_omni_vision_256x256.axmodel	64	51.600 ms
jina_v5_omni_audio_8s.axmodel	200	209.881 ms
jina_v5_omni_vision_448x448.axmodel	196	222.618 ms
jina_v5_omni_audio_30s.axmodel	750	3794.500 ms

Media prepare includes media loading, preprocessing, encoder execution, tokenizer work, and LLM input assembly. Standalone encoder latency is only the bare encoder model latency reported by ax_run_model. For video rows, standalone encoder latency is per-frame vision encoder latency x frame_count.

Optional packaged profiles are included for 448x448 vision and 30s audio. They are not loaded by the default config.json.

Runtime Footprint

Minimum board resources for the default config.json profile:

Item	Value
AXERA CMM required	~2.00 GiB
Linux RAM peak needed during startup	~707 MiB
Linux RAM used after startup	~80 MiB
Service ready time	15 s
Runtime files used by default	2.6 GiB
Full package size	4.1 GiB

The package enables release_axmodel_buffer_after_init, which reduces steady-state Linux RAM after model initialization. The board still needs enough temporary Linux RAM during startup to load and initialize the .axmodel files.

Token Layout and Static Shapes

The final embedding output is always [1, 1024].

Default encoder profiles:

Input	Static input profile	Soft tokens	Encoder output
Image	256x256	64	[1, 64, 1024]
Audio	8.0s, 16kHz, mono PCM WAV, 800 mel frames	200	[1, 200, 1024]
Video	frame directory, 256x256 per frame	64 x frame_count	[frame_count, 64, 1024] logically

Optional packaged profiles:

Encoder	Static input profile	Soft tokens	Encoder output
jina_v5_omni_vision_448x448.axmodel	448x448	196	[1, 196, 1024]
jina_v5_omni_audio_30s.axmodel	30.0s, 16kHz, mono	750	[1, 750, 1024]

The shipped video validation case uses 3 frames, so it contributes 192 visual soft tokens. Choose the frame count according to your application and the compiled prefill budget.

The packaged text backbone is compiled with:

• prefill_len = 256
• prefill_max_token_num = 768
• max_token_len = 2047

Notes

• This package uses static shapes. Arbitrary image resolution, arbitrary audio duration, or arbitrary video token budgets require rebuilding the corresponding encoder or LLM configuration.
• document and query are different retrieval prompt modes. Use document for corpus texts and query for search queries.
• Audio input must be 16kHz mono PCM WAV for this AX650 package. Convert audio offline if needed, for example: ffmpeg -i input.wav -ac 1 -ar 16000 -sample_fmt s16 output_16k_mono.wav.
• The default 8s audio HF reference is generated with 800 mel frames and 200 audio soft tokens. A 30s reference or 30s audio axmodel must use 3000 mel frames and 750 soft tokens.
• The packaged runtime supports frame-directory video embedding. Extract video frames before sending a video request.

Conversion References

• Upstream model: https://huggingface.co/jinaai/jina-embeddings-v5-omni-small
• AXERA runtime: https://github.com/AXERA-TECH/ax-llm

Discussion

This package is intended for AX650 deployment and validation. End users only need the files in this repository; the original Hugging Face checkpoints are not required at runtime.