jina-embeddings-v5-omni-nano-retrieval on AXERA NPU

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

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-nano-retrieval
cd AXERA-TECH/jina-embeddings-v5-omni-nano-retrieval
hf download AXERA-TECH/jina-embeddings-v5-omni-nano-retrieval --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
├── llama_p128_l0_together.axmodel
├── ...
├── llama_p128_l11_together.axmodel
├── llama_post.axmodel
├── jina_v5_omni_nano_vision_256x256.axmodel
├── jina_v5_omni_nano_audio_8s.axmodel
├── jina_v5_omni_nano_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-nano-retrieval-AX650-P128-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-nano-retrieval-AX650-P128-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-nano-retrieval-AX650-P128-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-nano-retrieval-AX650-P128-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-nano-retrieval-AX650-P128-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, 768]	-	0.999655
Text query	red_planet_query	[1, 768]	-	0.999539
Image	vision_sample	[1, 768]	64	0.993549
Audio	audio_test_chunk0_8s_wav	[1, 768]	200	0.992984
Video	video_visual_red_panda_openai_mp4	[1, 768]	192	0.997040

The board-side comparison script does not run the original Hugging Face model. It compares the API output with cached server-generated HF reference embeddings packaged in this repository.

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. Measurements exclude network overhead.

Scenario	Prompt	LLM input tokens	Soft tokens	Output shape	Media prepare	Standalone encoder axmodel	Prepare overhead	LLM prefill	Runtime total
Text document	document	19	-	[1, 768]	-	-	-	120.13 ms	121.30 ms
Text query	query	12	-	[1, 768]	-	-	-	93.95 ms	94.34 ms
Image	query	85	64	[1, 768]	139.05 ms	16.42 ms	122.63 ms	82.75 ms	221.81 ms
Audio (8s, 16kHz mono PCM WAV)	query	223	200	[1, 768]	804.79 ms	210.09 ms	594.70 ms	146.17 ms	950.96 ms
Video (3 frames)	query	213	192	[1, 768]	204.91 ms	49.25 ms	155.66 ms	143.55 ms	348.47 ms

Standalone encoder .axmodel latency was measured with ax_run_model -r 50 -w 10 on AX650:

Encoder axmodel	Output tokens	CMM size	Pulsar2 version	Avg latency
jina_v5_omni_nano_vision_256x256.axmodel	64	109390730 Bytes	6.0 563427a6	16.416 ms
jina_v5_omni_nano_audio_8s.axmodel	200	724467726 Bytes	6.0 563427a6	210.087 ms
jina_v5_omni_nano_audio_30s.axmodel	750	1030304014 Bytes	6.0 563427a6	3792.821 ms

Media prepare includes media loading, preprocessing, encoder execution, tokenizer work, and LLM input assembly. Standalone encoder axmodel is only the bare encoder model latency reported by ax_run_model. For video rows, Standalone encoder axmodel is per-frame vision encoder latency x 3. For the shipped 16kHz mono PCM WAV audio case, runtime profiling shows read_ms=5.316, resample_ms=0.001, log_mel_ms=571.178, and standalone audio encoder avg=210.087 ms.

Runtime Footprint

Minimum board resources for the default config.json profile:

Item	Value
AXERA CMM required	~1.14 GiB
Linux RAM used after startup	~82 MiB
Service ready time	15 s

If both packaged audio profiles are enabled, the board resource profile is:

Item	Value
AXERA CMM required	~2.08 GiB
Linux RAM used after startup	~82 MiB
Service ready time	16 s

CMM means AXERA contiguous multimedia memory. RSS means resident set size in Linux process memory. The package enables release_axmodel_buffer_after_init, so the runtime releases host-side .axmodel file buffers after model initialization.

Token Layout and Static Shapes

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

Default encoder profiles:

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

Optional packaged profile:

Encoder	Static input profile	Soft tokens	Encoder output
jina_v5_omni_nano_audio_30s.axmodel	30.0s, 16kHz, mono PCM WAV, 3000 mel frames	750	[1, 750, 768]

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 = 128
• warm-prefill groups: 128 / 256 / 384 / 512 / 640 / 768 / 896
• effective prefill_max_token_num = 1024
• max_token_len = 2047

Audio Profile Selection

The default config.json only loads:

"filename_audio_encoder_axmodel_short": "jina_v5_omni_nano_audio_8s.axmodel"

The package also includes jina_v5_omni_nano_audio_30s.axmodel. To enable long-audio deployments, add:

"filename_audio_encoder_axmodel_long": "jina_v5_omni_nano_audio_30s.axmodel"

If only one audio profile is configured, only that profile is loaded. If both short and long profiles are configured, the runtime selects the short profile for short audio and falls back to the long profile for longer audio.

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.
• The upstream nano model family contains multiple task adapters. This AX650 package is the retrieval release.

Conversion References

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

If you rebuild the LLM AXModels from the original Hugging Face checkpoint, make sure the build config contains the flattened field text_config.rope_theta = 1000000.0. The upstream nano config stores this value under text_config.rope_parameters.rope_theta, while the AXERA llama build route reads text_config.rope_theta.

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.