jonpark0/embeddinggemma-2-AX650

EmbeddingGemma 2 — AX650 / AX8850 (LLM-8850) NPU

google/embeddinggemma-2 (text, images and audio → one 768-d embedding space) converted with Axera Pulsar2 7.0-patch1 to run on the AX650 / AX8850 NPU, tested on the M5Stack LLM-8850 PCIe card (AXCL). The repository is a ready-to-run bundle: compiled axmodels, host-side assets, and a small host runtime that needs only numpy, Pillow and tokenizers (no PyTorch). The same code runs on Windows and Linux hosts.

한국어 안내는 아래에 있습니다.

Accuracy and speed

Measured on an LLM-8850 card (PCIe Gen2 x2) against the float32 model run through its own processor (the sentence-transformers pipeline: mean pooling + L2 normalization).

Input	Cosine to float32 (min)	Notes
Text (STS-B and KorSTS, 1,000 sentences; retrieval pairs)	0.9996	STS-B Spearman 0.9477 (float32 0.9478), KorSTS 0.8980 (0.8984)
Long text (up to 1024 tokens)	0.9997
Images (12 photos)	0.9994	text → photo retrieval top-1 100%
Audio (10 Korean / English clips)	0.9994	transcript → clip retrieval top-1 100%
Text + image + audio in one input	0.9994

Model	NPU time	One input end to end
Text, 128 tokens	15.4 ms	short sentence ≈ 17 ms
Text, 512 tokens	78.4 ms
Text, 1024 tokens	246.7 ms	≈ 255 ms
Vision (three graphs, 2520 patches)	≈ 1.34 s	≈ 1.6 s per image
Audio (up to 11.2 s)	62 ms	5.8 s clip ≈ 0.15 s

selftest.py repeats this check against the float32 references in refs/ and prints PASS.

Contents

Path	Description
assets/eg2_text_l128/l512/l1024.axmodel	Text encoder for 128 / 512 / 1024 tokens; the shortest one that fits is used
assets/eg2_vision_p2520_a/b1/b2.axmodel	Vision encoder, 16 layers split 8 + 4 + 4 (up to 280 soft tokens per image, aspect ratio kept)
assets/eg2_audio_f1120.axmodel	Audio encoder, up to 1120 log-mel frames (11.2 s → 280 soft tokens)
assets/*.npy, embed_tokens.bf16.bin, tokenizer.json, host_meta.json	Host assets (token embeddings, vision position table, projection, mel filters, metadata)
assets/config.json, embeddinggemma2_tokenizer.txt	Config for text-only serving with axllm (see below)
eg2_host.py	Preprocessing, tokenization, multimodal sequence assembly and NPU calls
axcl_session.py	ctypes wrapper of the AXCL runtime (libaxcl_rt.dll on Windows, libaxcl_rt.so on Linux)
eg2_cli.py, eg2_server.py	Command line tool; OpenAI-compatible /v1/embeddings server (text, images, audio)
selftest.py, refs/	Self test against float32 references
install_ubuntu.sh, build_axcl_driver.sh, eg2-embed.service	Ubuntu setup, AXCL driver build for kernel 7.x, systemd unit

Quick start

Requirements: an AX650 / AX8850 AXCL card with the AXCL driver and runtime installed (axcl-smi shows the card), Python 3.10+.

git lfs install
git clone https://huggingface.co/<this repo> && cd embeddinggemma-2-AX650
pip install -r requirements.txt          # numpy, tokenizers, pillow, soundfile, scipy
python3 selftest.py                      # -> PASS

python3 eg2_cli.py --prompt query "text:a photo of a cat" "image:refs/files/cat_0.jpeg" "image:refs/files/dog_0.jpeg"
python3 eg2_server.py --host 0.0.0.0 --port 8010
curl -s localhost:8010/v1/embeddings -H 'Content-Type: application/json' \
  -d '{"input": ["weather in Seoul", "weather in Busan"], "input_type": "query", "dimensions": 256}'

from eg2_host import Eg2
eg = Eg2("assets")
q = eg.encode(text="a red city bus", prompt="query")
d = eg.encode(images=["refs/files/bus.jpg"])
mixed = eg.encode(text="A red city bus <|image|> and its announcement <|audio|>",
                  images=["refs/files/bus.jpg"], audios=["refs/files/en_bus.wav"], prompt="document")
print(float(q @ d))
eg.close()

• prompt / input_type selects the task prefix of the original model card: query, document, sts, classification, clustering, code. Use query for queries and document for documents in retrieval; images and audio take no prefix.
• dim / dimensions: 768, 512, 256 or 128 (Matryoshka; re-normalized after truncation).
• Server inputs: a string, a list of strings, or objects {"text": "...<|image|>...", "image": "<base64>", "audio": "<base64 wav>"}.

Text-only serving with axllm

JonPark0/ax-llm (branch windows) adds model_type: embedding_gemma2 to axllm serve. Point it at assets/:

axllm serve assets --port 8000     # OpenAI-compatible /v1/embeddings, text only

Limitations

• At most 1024 tokens per input (text and soft tokens together); longer inputs are truncated keeping BOS and the final EOS. The original model accepts 8K.
• Audio: up to 11.2 s per clip (the original processor's cap of 280 tokens). Images: up to 280 soft tokens (the original default). Video is not wired up.
• One request at a time per card.
• Image preprocessing uses Pillow's bicubic resize instead of torchvision's; on float32 this alone gives cosine ≥ 0.99997 to the original.
• Tested on Windows 11 with the AXCL Windows driver. Linux hosts use the same code; the kernel 7.x driver build in build_axcl_driver.sh compiles but has not been load-tested on hardware.

How it was converted

• Each encoder was re-implemented in PyTorch with export-friendly operations and checked against the Hugging Face model (cosine 1.0) before export: static shapes, padding masks applied as multiply-and-renormalize after softmax, RMSNorm in rsqrt form, no dynamic shapes or boolean tensors in the graph.
• Pulsar2 7.0-patch1, U16 activations for all layers, NPU3, calibration on 32–64 real inputs (English, Korean, Chinese and code text; photos; speech).
• Host-side parts: token embedding lookup (× √512), vision patch position embeddings, 2-D RoPE tables and the 3×3 pooling matrix (so one graph serves every aspect ratio), the final vision RMSNorm and 768→512 projection (its squares overflow the NPU float path), audio log-mel features, L2 normalization.

License

Apache License 2.0, as the original model. See LICENSE and NOTICE (attribution and list of modifications).

---

한국어

Google embeddinggemma-2(텍스트·이미지·음성 → 하나의 768차원 벡터 공간)를 Axera Pulsar2 7.0-patch1로 변환해 AX650 / AX8850 NPU에서 돌리는 묶음입니다. M5Stack LLM-8850 PCIe 카드(AXCL)에서 시험했습니다. torch 없이 numpy·Pillow·tokenizers만으로 돌아가며, Windows와 Linux에서 같은 코드가 동작합니다.

• 정확도: 원본(float32) 대비 코사인 최소 0.9994(텍스트 0.9996). STS-B Spearman 0.9477(원본 0.9478), KorSTS 0.8980(0.8984). 텍스트·사진·음성 검색 정답률 모두 100%.
• 속도: 짧은 문장 약 17 ms, 1024토큰 약 255 ms, 이미지 1장 약 1.6초, 5.8초 음성 약 0.15초.
• 사용: pip install -r requirements.txt → python3 selftest.py(PASS 확인) → eg2_cli.py 또는 eg2_server.py(OpenAI 호환 /v1/embeddings). 텍스트만 쓸 때는 JonPark0/ax-llm windows 브랜치의 axllm serve assets도 쓸 수 있습니다.
• 작업 접두어: query, document, sts, classification, clustering, code. 검색에서는 질의에 query, 문서에 document를 씁니다. 이미지·음성에는 붙이지 않습니다.
• 제한: 입력 하나당 1024토큰(원본은 8K), 음성은 클립당 11.2초, 이미지는 장당 280토큰, 영상은 미지원, 요청은 한 번에 하나.
• Ubuntu 설치, 커널 7.x용 AXCL 드라이버 빌드, systemd 서비스 등 자세한 한국어 안내는 GUIDE_ko.md에 있습니다.
• 라이선스: 원본과 같은 Apache 2.0(LICENSE, NOTICE).