Modelpublic

JoramMillenaar/bge-small-zh-v1.5-vocab-quantized

sourceHugging Facemitupdated 3d agoView on Hugging Face

0likes39downloads

Model Card

bge-small-zh-v1.5 — vocab-quantized ONNX for WebGPU

BAAI/bge-small-zh-v1.5 as ONNX for Transformers.js, built for the WebGPU execution provider. Converted from Xenova/bge-small-zh-v1.5.

Only the word-embedding (vocabulary) table is quantized. It is stored as block-quantized uint8 (com.microsoft::GatherBlockQuantized, asymmetric, with zero points, along the hidden axis). The rest of the network (attention, FFN, LayerNorm, position/token-type embeddings) stays in fp32 or fp16. Inputs are int64 and the output last_hidden_state is float32 in every variant.

Files

File	Vocab table	Rest of graph	Size	Load with	Min cosine vs fp32*
`onnx/model_vocab8.onnx`	8-bit (block 128)	fp32	62.8 MB	`model_file_name: 'model_vocab8', dtype: 'fp32'`	0.99998
`onnx/model_vocab8_fp16.onnx`	8-bit (block 128)	fp16	36.9 MB	`model_file_name: 'model_vocab8', dtype: 'fp16'`	0.99998
`onnx/model_vocab4.onnx`	4-bit (block 32)	fp32	58.5 MB	`model_file_name: 'model_vocab4', dtype: 'fp32'`	0.99744
`onnx/model_vocab4_fp16.onnx`	4-bit (block 32)	fp16	32.1 MB	`model_file_name: 'model_vocab4', dtype: 'fp16'`	0.99744
`onnx/model_quantized.onnx` (CPU fallback)	int8 (dynamic)	int8 (dynamic)	24.0 MB	`dtype: 'q8'` (default on `device: 'wasm'`)	0.98559

\* Cosine similarity of the pooled, normalized sentence embedding against the original fp32 model, over 12 test sentences in English, German, French, Spanish, Chinese, Japanese and Russian.

The model_vocab* graphs contain no QuantizeLinear, DynamicQuantizeLinear, QLinearMatMul or MatMulInteger nodes. Checked with onnxruntime-web 1.31 (the version Transformers.js 4.3 ships): in those graphs every compute node is placed on WebGpuExecutionProvider. The only CPU nodes are the 7 small int64 shape nodes (attention-mask Unsqueeze/Cast, position-id Shape/Gather/Slice), which ORT always puts on CPU on purpose. The original fp32 export does the same.

Usage (Transformers.js ≥ 4)

import { pipeline } from '@huggingface/transformers';

const extractor = await pipeline('feature-extraction', '<your-namespace>/bge-small-zh-v1.5', {
  device: 'webgpu',
  model_file_name: 'model_vocab4', // or 'model_vocab8'
  dtype: 'fp16',                   // 'fp16' -> model_vocab4_fp16.onnx, 'fp32' -> model_vocab4.onnx
});

const output = await extractor(['This is a simple test.'], { pooling: 'cls', normalize: true });

Transformers.js builds the path as onnx/${model_file_name}${dtype === 'fp16' ? '_fp16' : ''}.onnx. There is no onnx/model.onnx, so for WebGPU always pass model_file_name. dtype: 'fp16' needs a GPU with the shader-f16 WebGPU feature; if the device lacks it, Transformers.js throws. Fall back to dtype: 'fp32' in that case.

CPU fallback (WASM)

onnx/model_quantized.onnx is the unmodified dynamic-int8 export from Xenova/bge-small-zh-v1.5. It is fast on CPU/WASM but uses DynamicQuantizeLinear + MatMulInteger ops, so do not use it with device: 'webgpu'.

const extractor = await pipeline('feature-extraction', '<your-namespace>/bge-small-zh-v1.5', {
  device: 'wasm', // dtype defaults to 'q8' on wasm -> onnx/model_quantized.onnx
});

Its embeddings drift more from the fp32 model than the model_vocab* files do (see the table above). If you store vectors in one index, embed everything with the same file.