ONNX export of dmis-lab/UltraBERT-Embed-150M for ONNX Runtime and Transformers.js. All credit for the model goes to dmis-lab. The weights are theirs, under the same Apache-2.0 license. We only converted the format.
Files
| Precision |
File |
Size |
Min cosine |
Mean cosine |
| fp32 |
onnx/model.onnx |
605 MB |
1.0000 |
1.00000 |
| fp16 |
onnx/model_fp16.onnx |
303 MB |
0.9999 |
0.99997 |
Inputs are input_ids and attention_mask, with dynamic batch size and sequence length. Output is last_hidden_state. The sliding-window (128 tokens) and full-attention layers are built from the attention mask inside the graph, so padding and any length work.
How we checked it
We ran 20 English texts through the original PyTorch model and through each ONNX file, then compared the sentence embedding (first-token pooling, as in the sentence-transformers config) with cosine similarity. Fifteen texts are short (3 to 25 tokens: questions, reviews, code, SQL, news, medical notes). Five are long book excerpts of 2,000 to 8,000 tokens.
Pass bar: cosine of at least 0.99 on every text. fp32 and fp16 pass.
We left out a 4-bit (q4) file. Weight-only 4-bit quantization (ONNX Runtime MatMulNBits, block size 32) gave a minimum cosine of 0.905, below the bar. Asymmetric, HQQ and block-16 variants did not reach 0.99 either.
Transformers.js 3 loaded each file in Node, produced the same token ids as the Python tokenizer, and reproduced the ONNX Runtime vectors (cosine above 0.999).
We also ran it in headless Chrome with the WebAssembly backend: fp32 and fp16 reproduced the ONNX Runtime vectors (cosine 1.0000 and 1.0000).
Usage with Transformers.js
import { AutoTokenizer, AutoModel } from '@huggingface/transformers';
const id = 'pastel-labs/UltraBERT-Embed-150M-ONNX';
const tokenizer = await AutoTokenizer.from_pretrained(id);
const model = await AutoModel.from_pretrained(id, { dtype: 'fp16' }); // or 'fp32'
async function embed(text) {
const { last_hidden_state } = await model(tokenizer(text));
const v = last_hidden_state.data.slice(0, 768); // first token
const n = Math.hypot(...v);
return v.map((x) => x / n);
}
Transformers.js v3 has no ultrabert model type, so it prints "Unknown model class" and falls back to the base class, which works. Pooling is the first token, as in 1_Pooling/config.json. Use AutoModel as above rather than the feature-extraction pipeline, so the pooling is explicit.
Notes
- fp16 marks the RMSNorm and softmax ops to stay in fp32 to avoid overflow; matrix multiplies run in fp16.
- Exported with the PyTorch TorchScript exporter, opset 17, from the original
modeling_ultrabert.py (eager attention).
- English only, like the original.
