lockR/rubert-base-pii-ner-tensorrt

ruBERT PII NER — TensorRT FP16 + CUDA Graphs

An optimized TensorRT FP16 inference export of redmadrobot-rnd/rubert-base-pii-ner, with a runtime that reuses CUDA Graphs for low-latency Russian personal-data and identity-document recognition.

On an NVIDIA RTX 4070 Ti SUPER, the supplied runtime achieved 1.089 ms p50 and 2.944 ms p95 end-to-end latency: 5.87× faster at the median than the remeasured PyTorch FP32 baseline. Exact span micro-F1 changed by −0.039 percentage points on all 2,841 benchmark texts.

The original fine-tuning belongs to Red Mad Robot R&D. This repository adds an inference export and runtime, with no additional training or fine-tuning. The engine uses FP16 with FP32 normalization accumulation; no Unsloth, INT8, or 4-bit quantization was used.

The prebuilt engine targets RTX 4070 Ti SUPER and TensorRT 10.13.3.9. Rebuild it on your deployment GPU using the supplied ONNX file and build script. For an option that does not require building a TensorRT engine, see the sibling ONNX Runtime 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-tensorrt --local-dir rubert-pii-trt
cd rubert-pii-trt
python -m pip install -r requirements.txt

# Build on your deployment GPU when its GPU/software stack differs.
python build_trt.py --onnx model.onnx --output model.fp16.engine

python pii_ner.py 'Иванов Пётр Сергеевич, паспорт 45 11 123456, тел. +7 999 123-45-67' --backend trt-graph

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

from pii_ner import PiiNER

ner = PiiNER(model_dir=".", backend="trt-graph")
ner.warmup()  # Capture all graph buckets before serving requests.
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. backend="trt-fp16" selects the same engine without CUDA Graphs.

The runtime uses 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. Long texts are processed through multiple windows. This postprocessing is part of the reported quality and latency.

What is included

File	Purpose
model.fp16.engine	Prebuilt TensorRT engine, approximately 359 MB
model.onnx	Original FP32 ONNX graph for rebuilding, approximately 711 MB
build_trt.py	FP16 engine builder with FP32 normalization accumulation
trt_backend.py	TensorRT buffers, execution contexts, and CUDA Graphs
pii_ner.py	Tokenization, windowing, and character-span decoding
config.json, tokenizer files	Original label mapping and tokenizer
requirements.txt	Pinned runtime dependencies

The engine's dynamic shape profile is min (1, 2), opt (1, 128), max (32, 512), expressed as (batch, sequence). The build uses a 4 GiB workspace limit and disables TF32. Inputs are input_ids, attention_mask, and token_type_ids; output is logits[batch, sequence, 43].

For batch size 1, the runtime retains CUDA Graphs for sequence-length buckets 32, 64, 128, 256, and 512. Each bucket owns a context and stable buffers, while engine weights are shared. First use creates and captures the bucket; later requests replay it. Larger batches use ordinary dynamic TensorRT execution. An instance serializes GPU calls to protect its reusable contexts and buffers. These are runtime optimizations; graphs are captured on the target machine and are not stored in the engine file.

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×
TensorRT FP16	80.873	82.544	81.700	−0.039	2.198	5.565	2.91×
TensorRT FP16 + CUDA Graphs	80.873	82.544	81.700	−0.039	1.089	2.944	5.87×

With CUDA Graphs, entity spans differed from the FP32 baseline on 10 of 2,841 documents; confidence scores are not included in this count. The largest per-type F1 declines were 0.611 percentage points for COUNTRY and 0.386 points for INN. Aggregate recall changed by −0.018 percentage points. The two TensorRT modes have the same aggregate metrics here but are not guaranteed to produce identical spans 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 baseline. The reported warm latency does not represent first-use graph capture or cold-start latency.

Environment. NVIDIA GeForce RTX 4070 Ti SUPER, compute capability 8.9; Linux under WSL2; Python 3.12.11; PyTorch 2.8.0 with CUDA 12.8; TensorRT 10.13.3.9 (CUDA 12 package); Transformers 4.57.1; ONNX 1.19.1; 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, CUDA Graphs metrics, ordinary TensorRT metrics, and reproduction instructions. The repository includes timing samples, environment metadata, and export/build provenance so that the reported comparison can be audited.

Engine portability. The supplied engine is specific to its build GPU and TensorRT/software environment. Matching compute capability alone does not establish compatibility. Rebuild on the deployment GPU after changing the GPU or TensorRT stack, and remeasure quality and latency. Builder tactics, hardware, and software versions can affect both numerical output and speed. Only the named build environment was measured; no universal 1 ms latency claim is made. CUDA Graphs also require stable, retained execution resources, and first requests incur initialization costs.

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.

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.