aria-bert-onnx

AriaBERT (ONNX, feature-extraction export)

An ONNX export of ViraIntelligentDataMining/AriaBERT's base encoder, for CPU inference via ONNX Runtime with no PyTorch dependency at runtime.

What this is

• The base RoBERTa encoder only — exported with task=feature-extraction, so the output is last_hidden_state, not masked-language-model logits. The original checkpoint is a RobertaForMaskedLM; this repo drops the MLM head, since only the encoder is needed to produce embeddings.
• Converted with 🤗 Optimum's ONNX exporter (optimum.exporters.onnx, --task feature-extraction), fp32.
• Everything else — weights, tokenizer, vocabulary — is unchanged from the original checkpoint.

What this is not

AriaBERT was pretrained as a general-purpose masked-language model — it was never fine-tuned with a sentence-similarity objective (contrastive/triplet loss, the way SBERT-style or E5-style embedding models are). Raw mean-pooled embeddings from a model like this measure surface-level similarity less reliably than a model built for it. If you need a Persian sentence-embedding model and don't specifically need this one, consider a model trained for that task instead (e.g. PartAI/Tooka-SBERT-V2-Small, intfloat/multilingual-e5-small).

Usage

No sentence-transformers, no PyTorch required — only onnxruntime and tokenizers:

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
tokenizer = Tokenizer.from_file("tokenizer.json")
pad_id = tokenizer.token_to_id("<pad>")
tokenizer.enable_padding(pad_id=pad_id, pad_token="<pad>")
tokenizer.enable_truncation(max_length=512)

def embed(texts: list[str]) -> np.ndarray:
    encodings = tokenizer.encode_batch(texts)
    input_ids = np.array([e.ids for e in encodings], dtype=np.int64)
    attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)

    (last_hidden_state,) = session.run(
        ["last_hidden_state"],
        {"input_ids": input_ids, "attention_mask": attention_mask},
    )

    # Mean pooling over real (non-pad) tokens, then L2-normalize.
    mask = attention_mask[:, :, None].astype(np.float32)
    summed = (last_hidden_state * mask).sum(axis=1)
    counts = mask.sum(axis=1).clip(min=1e-9)
    pooled = summed / counts
    return pooled / np.linalg.norm(pooled, axis=1, keepdims=True)

embeddings = embed(["سلام دنیا", "این یک متن نمونه است"])
print(embeddings.shape)  # (2, 768)

No task prefix is prepended to the input (unlike E5-family models) — AriaBERT was never trained with one, so adding "query: "/"passage: " here would just be noise, not signal.

Model details

Credits

All weights and training are from the original authors — this repo only changes the export format:

AriaBERT: A Pre-trained Persian BERT Model for Natural Language Understanding Arash Ghafouri, Mohammad Hossein Elahimanesh, Mohammad Reza Hassanpour Charmchi, Mohammad Amin Abbasi, Mahdi Firouzmandi — Vira Intelligent Data Mining Paper: https://www.researchsquare.com/article/rs-3558473/v1 Original model: https://huggingface.co/ViraIntelligentDataMining/AriaBERT