lockR/rubert-base-pii-ner-onnx

ruBERT PII NER — ONNX Runtime

Optimized FP16 and FP32 ONNX inference exports of redmadrobot-rnd/rubert-base-pii-ner, a Russian model for detecting personal information and identity-document numbers.

On an NVIDIA RTX 4070 Ti SUPER, the FP16 export achieved 2.58× lower median end-to-end latency than the remeasured PyTorch FP32 baseline, with an exact span micro-F1 change of −0.007 percentage points on all 2,841 benchmark texts. These are measurements for the supplied runtime and test environment, not a guarantee for other hardware or workloads.

The original fine-tuning belongs to Red Mad Robot R&D. This repository changes the inference representation and runtime; it introduces no additional training or fine-tuning. FP16 is a floating-point conversion; no Unsloth, INT8, or 4-bit quantization was used.

For the lowest measured NVIDIA GPU latency, see the sibling TensorRT + CUDA Graphs export.

Quickstart

Use Python 3.12 and an NVIDIA GPU with a driver compatible with the CUDA 12.8 PyTorch runtime. The pinned dependencies are in requirements.txt.

python -m pip install "huggingface_hub==0.36.2"
hf download lockR/rubert-base-pii-ner-onnx --local-dir rubert-pii-onnx
cd rubert-pii-onnx
python -m pip install -r requirements.txt
python pii_ner.py 'Иванов Пётр Сергеевич, паспорт 45 11 123456, тел. +7 999 123-45-67' --backend onnx-fp16

From Python, with the downloaded repository on the import path:

from pii_ner import PiiNER

ner = PiiNER(model_dir=".", backend="onnx-fp16")
entities = ner.predict("Иванов Пётр, паспорт 45 11 123456")
print(entities)

results = ner.predict_batch([
    "Почта: elena@pochta.ru",
    "Телефон: +7 999 123-45-67",
])

Each entity contains label, start, end, text, and score. Offsets refer to the original Python string; end is exclusive. Create the instance once and reuse it. Choose backend="onnx-fp32" to use the FP32 export.

The included runtime applies 512-token sliding windows with 128-token overlap, a fixed confidence threshold of 0.3, and batch size 1 by default. It resolves overlapping window predictions and joins adjacent same-type fragments, including fragments separated only by whitespace. This postprocessing is part of the reported quality and latency. Long texts are processed through multiple windows.

What is included

File	Purpose
model.optimized.fp16.onnx	Recommended FP16 graph, approximately 356 MB
model.optimized.fp32.onnx	FP32 graph, approximately 711 MB
config.json, tokenizer files	Original label mapping and tokenizer
pii_ner.py	CUDA inference, windowing, and character-span decoding
requirements.txt	Pinned runtime dependencies

The graphs use ONNX opset 17 and ONNX Runtime's BERT fusions, including com.microsoft operators. They are intended for ONNX Runtime with CUDAExecutionProvider; compatibility with other ONNX engines is not assumed. Both graphs accept input_ids, attention_mask, and token_type_ids as int64[batch, sequence], and return FP32 logits[batch, sequence, 43]. Batch and sequence dimensions are dynamic; a single window is limited to 512 tokens. Tokenization and entity decoding are provided by the Python wrapper.

Speed and quality

Quality uses all 2,841 test texts and 5,614 annotated entities in redmadrobot-rnd/pii_benchmark. Precision, recall, and F1 below are percentages; a correct prediction must match both character boundaries and one of the 21 original entity types.

Runtime	Precision	Recall	Exact micro-F1	F1 change, pp	p50, ms	p95, ms	p50 speedup
PyTorch FP32 baseline	80.932	82.561	81.739	—	6.387	18.346	1.00×
ONNX Runtime FP32	80.932	82.561	81.739	0.000	3.283	8.534	1.95×
ONNX Runtime FP16	80.918	82.561	81.732	−0.007	2.477	6.229	2.58×

ONNX FP32 reproduced the baseline entity spans on every benchmark document. ONNX FP16 changed the entity spans on 7 of 2,841 documents; confidence scores are not included in this agreement count. Its largest per-type F1 decline was 0.132 percentage points for INN. These results describe observed agreement on this dataset, not numerical equivalence on every input.

Timing protocol. Warm, sequential, local document requests at batch size 1: 256 texts sampled with seed 42, 30 warm-up requests, then three repetitions (768 timings per backend). End-to-end latency includes tokenization, every sliding window, CPU↔GPU transfers, synchronized GPU completion, and entity decoding. It excludes loading, HTTP, queueing, and concurrent serving. The three repetitions reuse the same 256 texts. Long-document windows run sequentially. TF32 was disabled for the FP32 comparisons.

Environment. NVIDIA GeForce RTX 4070 Ti SUPER; Linux under WSL2; Python 3.12.11; PyTorch 2.8.0 with CUDA 12.8; Transformers 4.57.1; ONNX 1.19.1; ONNX Runtime GPU 1.23.2; NumPy 2.5.3; Tokenizers 0.22.2.

Evaluation protocol. All backends share the tokenizer, decoder, confidence threshold, and inputs. Gold BIO boundaries and original texts are preserved; no threshold was tuned on the test split. Two rows contain an annotated WWW token for original text www; an explicitly recorded, equal-length case-insensitive alignment preserves their original offsets. No rows were excluded.

The upstream card reports 83.6 F1 for a folded 14-category exact-match protocol and 94.7 F1 for PERSON/LOCATION overlap matching. Those scores use different evaluation scopes and must not be directly compared with the 21-type exact-match scores above. The quality-preservation reference here is the remeasured upstream PyTorch FP32 model under the same protocol.

Entity types

Architecture: BertForTokenClassification, 12 layers, hidden size 768, 12 attention heads, and 177,749,803 source parameters.

The original 43-label BIO head is retained: O, plus B- and I- for each of these 21 types.

Group	Labels
Names	FIRST_NAME, LAST_NAME, MIDDLE_NAME
Addresses	COUNTRY, REGION, DISTRICT, CITY, STREET, HOUSE
Contacts and network identifiers	EMAIL, PHONE, URL, IP_ADDRESS
Document and payment identifiers	PASSPORT, INN, SNILS, OMS, CREDIT_CARD, DRIVER_LICENSE, MILITARY_ID, BIRTH_CERTIFICATE

Origin, evaluation artifacts, and limitations

The upstream model card describes fine-tuning ai-forever/ruBert-base on redmadrobot-rnd/pii_train: 17,137 annotated Russian sentences with 39,687 spans. The original authors report 10 epochs, learning rate 3e-5, batch size 16, and a maximum sequence length of 512. These are upstream training details; this export used no training data or calibration set. pii_benchmark was used only for evaluation.

The source model is pinned to c802e8cd26f85d1cf920973ea6f83965a0618d63, and the benchmark to f77ea831274daf980cc45c61a93c226be9d978d6. See the full comparison, FP16 metrics, FP32 metrics, and reproduction instructions. The repository includes timing samples, environment metadata, and export provenance so that the reported comparison can be audited.

This is a Russian PII recognizer, not a complete anonymization pipeline. It can miss entities or mark non-PII text as PII; a high aggregate F1 does not guarantee complete removal of personal data. The upstream model's weak IP_ADDRESS handling is inherited. Types outside the 21-label head, such as DATE_TIME, BANK_ACCOUNT, BIK, and TELEGRAM, require other components. Performance on other languages, domains, GPUs, and concurrency patterns has not been established by this benchmark. The supplied optimized FP16 graph and wrapper target NVIDIA CUDA inference, not CPU deployment.

License and credit: Apache-2.0, retaining the original model license. Credit Red Mad Robot R&D for the model, training data, and benchmark, and cite the original model repository alongside this export when reporting results.