ONNX export of dmis-lab/UltraBERT-380M 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 |
Top-1 token match |
| fp32 |
onnx/model.onnx |
1720 MB |
1.0000 |
1.00000 |
100.0% |
| fp16 |
onnx/model_fp16.onnx |
860 MB |
0.9995 |
0.99995 |
99.2% |
Inputs are input_ids and attention_mask, with dynamic batch size and sequence length. Output is last_hidden_state plus logits from the masked-language-model head. 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 mean-pooled last hidden states 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, capped at 4,096 tokens for this model because of memory on our test machine. Top-1 token match is how often the highest-scoring vocabulary token at each position is the same as in PyTorch.
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.926 and the masked-LM predictions stopped matching, 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).
Usage with Transformers.js
import { AutoTokenizer, AutoModel } from '@huggingface/transformers';
const id = 'pastel-labs/UltraBERT-380M-ONNX';
const tokenizer = await AutoTokenizer.from_pretrained(id);
const model = await AutoModel.from_pretrained(id, { dtype: 'fp16' }); // or 'fp32'
const out = await model(tokenizer('Paris is the capital of <|mask|>.'));
// out.logits: masked-language-model scores, [batch, tokens, 49152]
// out.last_hidden_state: token vectors, [batch, tokens, 1024]
Transformers.js v3 has no ultrabert model type, so AutoModel prints "Unknown model class" and falls back to the base class. That works and returns both outputs. AutoModelForMaskedLM and the fill-mask pipeline refuse the type. If you want a masked-LM class, ModernBertForMaskedLM.from_pretrained(id) loads this graph and returns logits (we checked that it predicts " France" for the example above). The architecture is not ModernBERT; that class only wraps the session.
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.
