ElXreno/LaBSE-en-ru-onnx

LaBSE-en-ru ONNX (Dense layer baked in, last_hidden_state output)

ONNX export of cointegrated/LaBSE-en-ru producing embeddings bit-identical to the sentence-transformers pipeline, unlike plain encoder exports.

Why this export exists

The sentence-transformers pipeline for LaBSE is Transformer → Pooling(CLS) → Dense(768→768, tanh) → Normalize. Standard ONNX exports (e.g. via optimum-cli export onnx) contain only the Transformer module — the trained 2_Dense projection is silently dropped, and any runtime that applies its own pooling over last_hidden_state produces vectors from a different embedding space (cosine vs the original pipeline ≈ 0, i.e. incompatible with stored LaBSE vectors).

This export bakes the Dense layer into the graph, applied per token:

last_hidden_state[b, t, :] = tanh(W · encoder(b, t) + b)

Because CLS pooling is just "take token 0", Dense(h[0]) == Dense_per_token(h)[0] — so a runtime that CLS-pools this graph's last_hidden_state and L2-normalizes reproduces the sentence-transformers output exactly (measured pairwise cosine = 1.0000 on ru/en test sets). The extra cost of applying Dense to all tokens is ~1% of encoder FLOPs.

Note: this equivalence holds for CLS pooling only (tanh is nonlinear, so it does not commute with mean pooling). LaBSE uses CLS pooling.

Files

• model.onnx — encoder + per-token Dense, fp32, opset 17, dynamic batch/seq axes; inputs input_ids, attention_mask, token_type_ids; output last_hidden_state [batch, seq, 768]
• tokenizer files identical to the base model

Usage with infinity

infinity_emb v2 \
  --model-id <this-repo> \
  --engine optimum \
  --pooling-method cls

--pooling-method cls is required (see above).

Usage with onnxruntime directly

import numpy as np, onnxruntime as ort
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("<this-repo>")
sess = ort.InferenceSession("model.onnx")
enc = tok(["пример текста", "example text"], padding=True, return_tensors="np")
out = sess.run(["last_hidden_state"], {k: enc[k].astype(np.int64) for k in ("input_ids", "attention_mask", "token_type_ids")})[0]
emb = out[:, 0, :]                                # CLS pooling
emb /= np.linalg.norm(emb, axis=1, keepdims=True) # normalize

Provenance

Exported by export.py (included in this repo) from cointegrated/LaBSE-en-ru; the script asserts cosine parity > 0.99999 against SentenceTransformer.encode before writing. Base model by cointegrated (David Dale), distilled from Google's LaBSE (apache-2.0).