gravitee-io/bert-small-pii-detection 🚀

Token-classification model for PII detection, fine-tuned from prajjwal1/bert-small on gravitee-io/pii-detection-dataset.

Label Set

AGE, COORDINATE, CREDIT_CARD, DATE_TIME, EMAIL_ADDRESS, FINANCIAL, HONORIFIC, IBAN_CODE, IMEI,
IP_ADDRESS, LOCATION, MAC_ADDRESS, NRP, ORGANIZATION, PASSWORD, PERSON, PHONE_NUMBER,
TITLE, URL, US_BANK_NUMBER, US_DRIVER_LICENSE, US_ITIN, US_LICENSE_PLATE, US_PASSPORT, US_SSN

How to Use

Quick start (pipeline)

from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

repo = "gravitee-io/bert-small-pii-detection"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForTokenClassification.from_pretrained(repo)

pipe = pipeline("token-classification", model=model, tokenizer=tok, aggregation_strategy="simple")
text = "Contact John Smith at john@example.com"
pipe(text)

ONNX

pip install transformers onnxruntime huggingface_hub 

from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, AutoConfig
import onnxruntime as ort

model_id = "gravitee-io/bert-small-pii-detection"

tokenizer = AutoTokenizer.from_pretrained(model_id)
id2label = AutoConfig.from_pretrained(model_id).id2label
session = ort.InferenceSession(hf_hub_download(model_id, "model.quant.onnx"))

text = "Contact John Smith at john@example.com"
enc = tokenizer(text, return_tensors="np")
inputs = {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]}
logits = session.run(None, inputs)[0][0]

tokens = tokenizer.convert_ids_to_tokens(enc["input_ids"][0])
labels = [id2label[i] for i in logits.argmax(-1)]

for tok, label in zip(tokens, labels):
  print(f"{tok:<20} {label}")

Intended use

Detect personally identifiable information (PII) spans in english text. Suitable for privacy filtering, redaction pipelines, and data-leak prevention particularly on structured data (JSON, HTML, XML, SQL, Document)

Evaluation

Limitations

• English-focused; other languages will degrade
• Domain drift is real: audit on your own data

---

Benchmarks

External-corpus evaluation (English only), seqeval. Last run: 2026-05-21.

Per-entity breakdown

gretelai/gretel-pii-masking-en-v1:test gretelai/synthetic_pii_finance_multilingual:test DataikuNLP/kiji-pii-training-data:test beki/privy:test beki/privy:test-large

Citation

Data citation are present in the dataset card used for this model. If you use the model, please consider citing the papers:

@misc{bhargava2021generalization,
      title={Generalization in NLI: Ways (Not) To Go Beyond Simple Heuristics}, 
      author={Prajjwal Bhargava and Aleksandr Drozd and Anna Rogers},
      year={2021},
      eprint={2110.01518},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

@article{DBLP:journals/corr/abs-1908-08962,
  author    = {Iulia Turc and
               Ming{-}Wei Chang and
               Kenton Lee and
               Kristina Toutanova},
  title     = {Well-Read Students Learn Better: The Impact of Student Initialization
               on Knowledge Distillation},
  journal   = {CoRR},
  volume    = {abs/1908.08962},
  year      = {2019},
  url       = {http://arxiv.org/abs/1908.08962},
  eprinttype = {arXiv},
  eprint    = {1908.08962},
  timestamp = {Thu, 29 Aug 2019 16:32:34 +0200},
  biburl    = {https://dblp.org/rec/journals/corr/abs-1908-08962.bib},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}