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