kiji-pii-model-onnx

import numpy as np
from onnxruntime import InferenceSession
from transformers import AutoTokenizer

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("DataikuNLP/kiji-pii-model-onnx")
session = InferenceSession("DataikuNLP/kiji-pii-model-onnx/model_quantized.onnx")  # or local path

# Tokenize
text = "Contact John Smith at [email protected] or call +1-555-123-4567."
inputs = tokenizer(text, return_tensors="np", truncation=True, max_length=512)

# Run inference
outputs = session.run(None, dict(inputs))
pii_logits, coref_logits = outputs  # (1, seq_len, 53), (1, seq_len, 7)

# Decode PII predictions
pii_predictions = np.argmax(pii_logits, axis=-1)[0]

# See label_mappings.json for label ID -> label name mapping