Add validated ruBERT PII inference export, model card and benchmarks

73be581 verified about 1 month ago

11.2 kB

---
license: apache-2.0
language:
- ru
base_model: redmadrobot-rnd/rubert-base-pii-ner
pipeline_tag: token-classification
datasets:
- redmadrobot-rnd/pii_train
tags:
- tensorrt
- cuda
- cuda-graphs
- fp16
- pii
- ner
- named-entity-recognition
- russian
model-index:
- name: lockR/rubert-base-pii-ner-tensorrt (FP16 + CUDA Graphs)
  results:
  - task:
      type: token-classification
      name: Russian PII named entity recognition
    dataset:
      type: redmadrobot-rnd/pii_benchmark
      name: Russian PII Benchmark — all 21 entity types
      split: test
      revision: f77ea831274daf980cc45c61a93c226be9d978d6
    metrics:
    - type: precision
      name: Exact character-span micro-precision (21 types)
      value: 0.8087260034904014
    - type: recall
      name: Exact character-span micro-recall (21 types)
      value: 0.8254364089775561
    - type: f1
      name: Exact character-span micro-F1 (21 types)
      value: 0.8169957686882934
---
# ruBERT PII NER — TensorRT FP16 + CUDA Graphs

An optimized **TensorRT FP16 inference export** of
[redmadrobot-rnd/rubert-base-pii-ner](https://huggingface.co/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](https://huggingface.co/lockR/rubert-base-pii-ner-onnx).

## 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`.

```bash
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:

```python
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](https://huggingface.co/datasets/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](https://huggingface.co/redmadrobot-rnd/rubert-base-pii-ner)
describes fine-tuning `ai-forever/ruBert-base` on
[redmadrobot-rnd/pii_train](https://huggingface.co/datasets/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`](https://huggingface.co/redmadrobot-rnd/rubert-base-pii-ner/tree/c802e8cd26f85d1cf920973ea6f83965a0618d63),
and the benchmark to
[`f77ea831274daf980cc45c61a93c226be9d978d6`](https://huggingface.co/datasets/redmadrobot-rnd/pii_benchmark/tree/f77ea831274daf980cc45c61a93c226be9d978d6).
See the [full comparison](results/summary.md),
[CUDA Graphs metrics](results/trt-graph.json),
[ordinary TensorRT metrics](results/trt-fp16.json), and
[reproduction instructions](reproduce/README.md). 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](https://huggingface.co/redmadrobot-rnd/rubert-base-pii-ner)
alongside this export when reporting results.

	---
	license: apache-2.0
	language:
	- ru
	base_model: redmadrobot-rnd/rubert-base-pii-ner
	pipeline_tag: token-classification
	datasets:
	- redmadrobot-rnd/pii_train
	tags:
	- tensorrt
	- cuda
	- cuda-graphs
	- fp16
	- pii
	- ner
	- named-entity-recognition
	- russian
	model-index:
	- name: lockR/rubert-base-pii-ner-tensorrt (FP16 + CUDA Graphs)
	results:
	- task:
	type: token-classification
	name: Russian PII named entity recognition
	dataset:
	type: redmadrobot-rnd/pii_benchmark
	name: Russian PII Benchmark — all 21 entity types
	split: test
	revision: f77ea831274daf980cc45c61a93c226be9d978d6
	metrics:
	- type: precision
	name: Exact character-span micro-precision (21 types)
	value: 0.8087260034904014
	- type: recall
	name: Exact character-span micro-recall (21 types)
	value: 0.8254364089775561
	- type: f1
	name: Exact character-span micro-F1 (21 types)
	value: 0.8169957686882934
	---

	# ruBERT PII NER — TensorRT FP16 + CUDA Graphs

	An optimized TensorRT FP16 inference export of
	[redmadrobot-rnd/rubert-base-pii-ner](https://huggingface.co/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](https://huggingface.co/lockR/rubert-base-pii-ner-onnx).

	## 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`.

	```bash
	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:

	```python
	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](https://huggingface.co/datasets/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](https://huggingface.co/redmadrobot-rnd/rubert-base-pii-ner)
	describes fine-tuning `ai-forever/ruBert-base` on
	[redmadrobot-rnd/pii_train](https://huggingface.co/datasets/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`](https://huggingface.co/redmadrobot-rnd/rubert-base-pii-ner/tree/c802e8cd26f85d1cf920973ea6f83965a0618d63),
	and the benchmark to
	[`f77ea831274daf980cc45c61a93c226be9d978d6`](https://huggingface.co/datasets/redmadrobot-rnd/pii_benchmark/tree/f77ea831274daf980cc45c61a93c226be9d978d6).
	See the [full comparison](results/summary.md),
	[CUDA Graphs metrics](results/trt-graph.json),
	[ordinary TensorRT metrics](results/trt-fp16.json), and
	[reproduction instructions](reproduce/README.md). 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](https://huggingface.co/redmadrobot-rnd/rubert-base-pii-ner)
	alongside this export when reporting results.