"""Portable inference for exported RuBERT PII ONNX and TensorRT models. Only the tokenizer is loaded through Transformers. Model weights are executed by ONNX Runtime or TensorRT. The default artifact directory is this file's parent, so a downloaded repository can be used independently of the working directory. TensorRT CUDA graphs are captured lazily for five sequence-length buckets. The first request in each bucket includes graph capture and is slower than warmed inference. Call ``model.warmup()`` before serving latency-sensitive requests. """ from __future__ import annotations import argparse import importlib.util import json from pathlib import Path from threading import Lock import numpy as np from transformers import AutoTokenizer ROOT = Path(__file__).resolve().parent INPUT_NAMES = ('input_ids', 'attention_mask', 'token_type_ids') BACKENDS = ('onnx-fp16', 'onnx-fp32', 'trt-graph', 'trt-fp16') def require_file(path): path = Path(path) if not path.is_file(): raise FileNotFoundError(f'Missing model artifact: {path}') return path class OnnxBackend: def __init__(self, model_dir, precision='fp16'): model_path = require_file(Path(model_dir) / f'model.optimized.{precision}.onnx') # Import torch before ORT so its CUDA/cuDNN libraries are discoverable. try: import torch # noqa: F401 import onnxruntime as ort except ImportError as error: raise RuntimeError( 'ONNX inference requires CUDA-enabled torch and onnxruntime-gpu; ' 'install the repository requirements.' ) from error ort.preload_dlls() if 'CUDAExecutionProvider' not in ort.get_available_providers(): raise RuntimeError('ONNX Runtime CUDA provider is unavailable; install onnxruntime-gpu') options = ort.SessionOptions() options.intra_op_num_threads = 4 options.inter_op_num_threads = 1 self.session = ort.InferenceSession( str(model_path), options, providers=[('CUDAExecutionProvider', {'device_id': 0, 'use_tf32': 0}), 'CPUExecutionProvider'], ) if self.session.get_providers()[0] != 'CUDAExecutionProvider': raise RuntimeError( 'ONNX Runtime failed to initialize CUDA and fell back to CPU. ' 'Check the NVIDIA driver and compatible CUDA/cuDNN libraries.' ) self.names = [item.name for item in self.session.get_inputs()] def run(self, inputs): return self.session.run(['logits'], {k: inputs[k] for k in self.names})[0] def load_trt_backend(model_dir, backend, pad_token_id): engine_path = require_file(model_dir / 'model.fp16.engine') runtime_path = require_file(model_dir / 'trt_backend.py') try: import torch except ImportError as error: raise RuntimeError('TensorRT inference requires CUDA-enabled torch') from error if not torch.cuda.is_available(): raise RuntimeError('TensorRT inference requires an available NVIDIA CUDA GPU') # Resolve the sibling runtime explicitly; no cwd/sys.path modifications. spec = importlib.util.spec_from_file_location('_pii_ner_trt_backend', runtime_path) module = importlib.util.module_from_spec(spec) try: spec.loader.exec_module(module) except ImportError as error: raise RuntimeError( 'TensorRT inference requires tensorrt-cu12 and CUDA-enabled torch; ' 'install the repository requirements matching the engine build.' ) from error if backend == 'trt-graph': return module.TensorRTGraphBackend(engine_path, pad_token_id=pad_token_id) return module.TensorRTBackend(engine_path) def resolve_window_overlaps(entities): """Match HF pipeline overlap policy: longest span, then highest confidence. Overflow windows may disagree about a boundary or label. Resolve those alternatives before merging adjacent same-type fragments into one entity. """ if not entities: return [] ordered = sorted(entities, key=lambda item: item['start']) resolved, previous = [], ordered[0] for item in ordered[1:]: if item['start'] < previous['end']: length = item['end'] - item['start'] previous_length = previous['end'] - previous['start'] if length > previous_length or (length == previous_length and item['score'] > previous['score']): previous = item else: resolved.append(previous) previous = item resolved.append(previous) return resolved def merge_fragments(text, entities): """Merge overlapping / adjacent same-type fragments, preserving character offsets.""" merged = [] for item in sorted(entities, key=lambda x: (x['start'], x['end'])): if merged: previous = merged[-1] gap = text[previous['end']:item['start']] if previous['label'] == item['label'] and ( item['start'] <= previous['end'] or (not gap or gap.isspace()) ): previous['end'] = max(previous['end'], item['end']) previous['score'] = max(previous['score'], item['score']) previous['text'] = text[previous['start']:previous['end']] continue merged.append(dict(item)) return merged def decode_window(text, offsets, logits, id2label, min_confidence=0.3): # Float32 softmax for identical aggregation across all numeric backends. logits = np.asarray(logits, dtype=np.float32) shifted = logits - logits.max(axis=-1, keepdims=True) probabilities = np.exp(shifted) probabilities /= probabilities.sum(axis=-1, keepdims=True) ids = logits.argmax(axis=-1) confidences = probabilities[np.arange(len(ids)), ids] entities, current = [], None def flush(): nonlocal current if current: current['score'] = float(np.mean(current.pop('_scores'))) if current['score'] >= min_confidence: current['text'] = text[current['start']:current['end']] entities.append(current) current = None for (start, end), label_id, score in zip(offsets, ids, confidences): start, end = int(start), int(end) if end <= start: flush() continue label = id2label[int(label_id)] if label == 'O': flush() continue prefix, kind = label.split('-', 1) if current and prefix == 'I' and kind == current['label']: current['end'] = end current['_scores'].append(float(score)) else: flush() current = {'start': start, 'end': end, 'label': kind, '_scores': [float(score)]} flush() return entities class PiiNER: """Recognize PII spans using a local exported model repository. Args: model_dir: Artifact directory; defaults to the directory of this file. backend: ``onnx-fp16``, ``onnx-fp32``, ``trt-fp16`` or ``trt-graph``. If omitted, read ``default_backend`` from runtime_config.json; if that file is absent, default to ``onnx-fp16``. min_confidence: Mean token confidence threshold, applied before merging. batch_size: Number of overflow windows per inference call, from 1 to 32. CUDA graphs apply only to batches of one; larger batches use the dynamic TensorRT backend. ``predict(str)`` returns a list of entity dictionaries. ``predict(list)`` and ``predict_batch(list)`` return one entity list per input text. Offsets are Python character indices with an exclusive end, so the returned text is always ``input_text[entity['start']:entity['end']]``. """ def __init__(self, model_dir=None, backend=None, min_confidence=0.3, batch_size=1): if not 0 <= min_confidence <= 1: raise ValueError('min_confidence must be between 0 and 1') if not isinstance(batch_size, int) or not 1 <= batch_size <= 32: raise ValueError('batch_size must be an integer between 1 and 32') self.model_dir = Path(model_dir).expanduser().resolve() if model_dir is not None else ROOT if backend is None: runtime_path = self.model_dir / 'runtime_config.json' runtime_config = json.loads(runtime_path.read_text(encoding='utf-8')) if runtime_path.is_file() else {} backend = runtime_config.get('default_backend', 'onnx-fp16') if backend not in BACKENDS: raise ValueError(f'Unsupported backend: {backend!r}; choose one of {", ".join(BACKENDS)}') self.backend_name = backend self.min_confidence, self.batch_size = min_confidence, batch_size for filename in ('config.json', 'tokenizer.json', 'tokenizer_config.json'): require_file(self.model_dir / filename) self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir, local_files_only=True) config = json.loads((self.model_dir / 'config.json').read_text(encoding='utf-8')) self.id2label = {int(k): v for k, v in config['id2label'].items()} if backend.startswith('onnx-'): self.backend = OnnxBackend(self.model_dir, backend.split('-')[1]) else: self.backend = load_trt_backend(self.model_dir, backend, self.tokenizer.pad_token_id) # GPU contexts / reusable buffers must not be shared by concurrent calls. self._lock = Lock() def warmup(self): """Initialize kernels and capture every CUDA graph bucket before serving. This performs actual GPU inference, retains graph contexts/buffers, and can take several seconds. First-use graph capture is otherwise charged to the first request at each new bucket. Call once after construction. Dynamic ONNX/TensorRT modes are exercised at the same five lengths; their latency can still vary for previously unseen dynamic shapes. """ # Trigger the tokenizer's own first-call initialization as well. self.tokenizer('Привет', add_special_tokens=True) buckets = getattr(self.backend, 'buckets', (32, 64, 128, 256, 512)) batch_sizes = [1] if self.backend_name == 'trt-graph' else [self.batch_size] if self.backend_name == 'trt-graph' and self.batch_size != 1: batch_sizes.append(self.batch_size) with self._lock: for batch_size in batch_sizes: for length in buckets: ids = np.full((batch_size, length), self.tokenizer.pad_token_id, dtype=np.int64) ids[:, 0] = self.tokenizer.cls_token_id ids[:, 1] = self.tokenizer.sep_token_id attention_mask = np.zeros_like(ids) attention_mask[:, :2] = 1 output = self.backend.run({ 'input_ids': ids, 'attention_mask': attention_mask, 'token_type_ids': np.zeros_like(ids), }) if not np.isfinite(output).all(): raise RuntimeError('Non-finite logits returned by inference backend during warmup') def predict(self, texts): if isinstance(texts, str): return self.predict_batch([texts])[0] return self.predict_batch(texts) def predict_batch(self, texts): texts = list(texts) if not texts: return [] if any(not isinstance(t, str) for t in texts): raise TypeError('Every input must be a string') result = [[] for _ in texts] encoded = self.tokenizer( texts, truncation=True, max_length=512, stride=128, return_overflowing_tokens=True, return_offsets_mapping=True, padding=False, ) mapping = encoded['overflow_to_sample_mapping'] # Sort windows by length to avoid excess padding in bulk processing. order = sorted(range(len(mapping)), key=lambda i: len(encoded['input_ids'][i])) with self._lock: for start in range(0, len(order), self.batch_size): indices = order[start:start + self.batch_size] features = [{k: encoded[k][i] for k in INPUT_NAMES if k in encoded} for i in indices] batch = self.tokenizer.pad(features, padding=True, pad_to_multiple_of=8, return_tensors='np') inputs = {k: np.ascontiguousarray(batch[k], dtype=np.int64) for k in INPUT_NAMES if k in batch} if 'token_type_ids' not in inputs: inputs['token_type_ids'] = np.zeros_like(inputs['input_ids']) output = self.backend.run(inputs) if not np.isfinite(output).all(): raise RuntimeError('Non-finite logits returned by inference backend') for row, i in enumerate(indices): document = mapping[i] offsets = encoded['offset_mapping'][i] result[document].extend(decode_window( texts[document], offsets, output[row, :len(offsets)], self.id2label, self.min_confidence, )) return [merge_fragments(text, resolve_window_overlaps(spans)) for text, spans in zip(texts, result)] if __name__ == '__main__': parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('text', help='Text to analyze') parser.add_argument('--model-dir', type=Path, help="Local artifact directory; defaults to this file's directory") parser.add_argument('--backend', choices=BACKENDS, help='Override runtime_config.json default_backend') parser.add_argument('--min-confidence', type=float, default=0.3) parser.add_argument('--warmup', action='store_true', help='Warm all sequence buckets before prediction') args = parser.parse_args() model = PiiNER(args.model_dir, args.backend, min_confidence=args.min_confidence) if args.warmup: model.warmup() print(json.dumps(model.predict(args.text), ensure_ascii=False, indent=2))