README.md · jonpark0/embeddinggemma-2-AX650 at main

Add EmbeddingGemma 2 converted for the AX650 / AX8850 NPU

9037c05 verified 21 days ago

8.32 kB

---
license: apache-2.0
base_model: google/embeddinggemma-2
base_model_relation: quantized
pipeline_tag: feature-extraction
language:
- multilingual
- en
- ko
tags:
- embedding
- multimodal-embedding
- sentence-similarity
- matryoshka
- axera
- ax650
- ax8850
- llm-8850
- axcl
- npu
- pulsar2
---
# EmbeddingGemma 2 — AX650 / AX8850 (LLM-8850) NPU

[`google/embeddinggemma-2`](https://huggingface.co/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+.

```bash
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}'
```

```python
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](https://github.com/JonPark0/ax-llm) (branch `windows`) adds `model_type: embedding_gemma2` to `axllm serve`. Point it at `assets/`:

```bash
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`](https://huggingface.co/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](https://github.com/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`](GUIDE_ko.md)에 있습니다.
- 라이선스: 원본과 같은 Apache 2.0(`LICENSE`, `NOTICE`).

	---
	license: apache-2.0
	base_model: google/embeddinggemma-2
	base_model_relation: quantized
	pipeline_tag: feature-extraction
	language:
	- multilingual
	- en
	- ko
	tags:
	- embedding
	- multimodal-embedding
	- sentence-similarity
	- matryoshka
	- axera
	- ax650
	- ax8850
	- llm-8850
	- axcl
	- npu
	- pulsar2
	---

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

	[`google/embeddinggemma-2`](https://huggingface.co/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+.

	```bash
	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}'
	```

	```python
	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](https://github.com/JonPark0/ax-llm) (branch `windows`) adds `model_type: embedding_gemma2` to `axllm serve`. Point it at `assets/`:

	```bash
	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`](https://huggingface.co/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](https://github.com/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`](GUIDE_ko.md)에 있습니다.
	- 라이선스: 원본과 같은 Apache 2.0(`LICENSE`, `NOTICE`).