PII Detection System

Production-ready PII detection system using a fine-tuned DeBERTa-v3-base token classification model.

Research Papers

• Dataset : PIIBench: A Unified Multi-Source Benchmark Corpus for Personally Identifiable Information Detection

• Models : Fine-Tuning Over Architectural Complexity: PII Detection on PIIBench with DeBERTa

Project Structure

pii-detector/
├── data/                          # Train/val/test splits + label mapping
│   ├── train.jsonl
│   ├── val.jsonl
│   ├── test.jsonl
│   └── label_mapping.json
├── models/
│   ├── deberta-v3-base/           # Downloaded base model weights
│   ├── checkpoints/               # Training checkpoints
│   └── best_model/                # Final fine-tuned model
├── pii_datasets/                  # Raw downloaded datasets
│   └── consolidated/
│       ├── consolidated.jsonl
│       └── entity_types.json
├── src/
│   ├── download_datasets.py       # Downloads 10 datasets from HuggingFace
│   ├── consolidate_pii_datasets.py # Normalises all sources into BIO format
│   ├── data_preparation.py        # Splits, rare-entity dropping, label mapping
│   ├── download_model.py          # Downloads microsoft/deberta-v3-base
│   ├── train.py                   # Fine-tuning with HuggingFace Trainer
│   ├── inference.py               # PIIDetector and FastPIIDetector classes
│   ├── api.py                     # Flask REST API
│   └── exceptions.py              # Custom exception hierarchy
├── run_data_pipeline.py           # Orchestrates steps 1-3 of data prep
├── run_training_pipeline.py       # Orchestrates model download + fine-tuning
├── run_benchmarking.py            # Benchmarks our model vs spaCy vs Presidio
├── run_streaming_model_benchmark.py # Full-test evaluation for a trained model
├── test_detector.py               # Test suite
├── example_client.py              # API client demo
├── Makefile                       # Convenience targets
├── Dockerfile
└── requirements.txt

Datasets

Ten datasets are downloaded from HuggingFace and consolidated into a unified BIO-tagged format:

finer-139 is capped at 150k records during data preparation. Entity types with fewer than 500 B- mentions globally are collapsed to O.

Quick Start

1. Install dependencies

pip install -r requirements.txt

2. Run the data pipeline

python run_data_pipeline.py

This downloads all datasets, normalises them into BIO format, drops rare entity types, and writes stratified 80/10/10 train/val/test splits.

By default, prepared records contain tokens, labels, and source. To also include source text for downstream evaluation, redaction, and dataset-card publishing, run:

python run_data_pipeline.py --include-text

When --include-text is used during consolidation, source-native text is preserved where the constituent dataset provides it; token-only sources use a deterministic whitespace-joined token fallback.

Skip steps you have already run:

python run_data_pipeline.py --skip-download
python run_data_pipeline.py --skip-download --skip-consolidate

3. Run the training pipeline

python run_training_pipeline.py

This downloads microsoft/deberta-v3-base and fine-tunes it. The best checkpoint is saved to ./models/best_model.

Skip the model download if already done:

python run_training_pipeline.py --skip-download

4. Run tests

python test_detector.py

5. Start the API server

python src/api.py --model-path ./models/best_model

Or via Make:

make api

Training Configuration

Default hyperparameters (run_training_pipeline.py):

Override any parameter:

python run_training_pipeline.py \
    --batch-size 8 \
    --grad-accum 8 \
    --epochs 5 \
    --max-length 256 \
    --gradient-checkpointing

Resume from a checkpoint:

python run_training_pipeline.py --skip-download --resume-from-checkpoint

Novel Training

The integrated novelty trainer keeps the corrected prepared splits and weighted loss, then optionally adds source conditioning, curriculum learning, and a coarse-to-fine hierarchical head.

Quick sanity run:

python run_training_pipeline.py \
    --skip-download \
    --novel \
    --source-cond \
    --hierarchical \
    --max-length 256 \
    --batch-size 6 \
    --eval-batch-size 8 \
    --grad-accum 10 \
    --max-steps 500 \
    --eval-steps 100 \
    --save-steps 100 \
    --logging-steps 25 \
    --skip-final-eval

Longer checkpointed run:

python run_training_pipeline.py \
    --skip-download \
    --novel \
    --source-cond \
    --hierarchical \
    --novel-output-dir ./models/full_novel \
    --max-length 256 \
    --batch-size 6 \
    --eval-batch-size 8 \
    --grad-accum 10 \
    --epochs 3 \
    --eval-steps 1000 \
    --save-steps 1000 \
    --logging-steps 50 \
    --eval-accumulation-steps 1 \
    --gradient-checkpointing \
    --skip-final-eval

The saved novelty model is written to models/full_novel/final_model. Hierarchical novelty models include metadata so PIIDetector can load the custom head.

Full novelty run with source conditioning, curriculum learning, and the hierarchical head:

export TOKENIZERS_PARALLELISM=false

python run_training_pipeline.py \
    --novel \
    --source-cond \
    --curriculum \
    --hierarchical \
    --novel-output-dir ./models/full_novel_curriculum \
    --max-length 256 \
    --batch-size 6 \
    --eval-batch-size 8 \
    --grad-accum 10 \
    --eval-steps 1000 \
    --save-steps 1000 \
    --logging-steps 50 \
    --eval-accumulation-steps 1 \
    --gradient-checkpointing \
    --skip-final-eval

Curriculum mode trains one epoch in each of three ordered source-family phases (general NER -> synthetic PII -> financial PII); --epochs and --max-steps do not control this schedule. Checkpoints are written under models/full_novel_curriculum/phase_1, phase_2, and phase_3. Repeat the same command with --skip-download --resume-from-checkpoint after an interruption.

Inference

PIIDetector (single-text / sequential batch)

from src.inference import PIIDetector

detector = PIIDetector(
    model_path="./models/best_model",
    confidence_threshold=0.5,
)

result = detector.detect("My email is [email protected] and SSN is 123-45-6789")
print(result.has_pii)        # True
print(result.redacted_text)  # "My email is [REDACTED] and SSN is [REDACTED]"
print(result.pii_types)      # ["EMAIL", "SSN"]
print(result.entities)
# [
#   {"text": "[email protected]", "type": "EMAIL", "start": 12, "end": 29, "confidence": 0.97},
#   {"text": "123-45-6789", "type": "SSN", "start": 39, "end": 50, "confidence": 0.95},
# ]

detect() always returns a PIIResult dataclass. has_pii is False and error is set on per-item inference failures, so batch callers can continue processing.

FastPIIDetector (padded-batch inference)

from src.inference import FastPIIDetector

detector = FastPIIDetector(
    model_path="./models/best_model",
    confidence_threshold=0.5,
    batch_size=32,
)

texts = ["Contact Alice at [email protected]", "The weather is nice"]
results = detector.batch_detect(texts)
stats = detector.get_pii_statistics(results)

FastPIIDetector pads inputs and runs a single forward pass per batch. Use this for high-throughput pipelines.

PIIResult fields

REST API

Endpoints

GET /health

curl http://localhost:5000/health

{"status": "healthy", "model_loaded": true}

GET /info

Returns model path, device, confidence threshold, supported PII types, num_labels, batch_size.

POST /detect

curl -X POST http://localhost:5000/detect \
  -H "Content-Type: application/json" \
  -d '{"text": "Email me at [email protected]"}'

{
  "has_pii": true,
  "redacted_text": "Email me at [REDACTED]",
  "pii_types": ["EMAIL"],
  "entities": [{"text": "[email protected]", "type": "EMAIL", "start": 12, "end": 28, "confidence": 0.97}],
  "error": null
}

POST /detect/batch

curl -X POST http://localhost:5000/detect/batch \
  -H "Content-Type: application/json" \
  -d '{"texts": ["My name is John", "SSN: 123-45-6789"], "return_stats": true}'

Returns {"results": [...], "statistics": {...}}. Every entry in results is a PIIResult dict — never null.

POST /detect/file

Accepts .txt, .csv, .xlsx, .xls uploads (max 50 MB).

curl -X POST http://localhost:5000/detect/file \
  -F "[email protected]" \
  -F "columns=name,email"

For CSV/Excel, columns is an optional comma-separated list of column names to process. Omit to process all columns. The response includes per-cell results and a fully redacted copy of the records.

For plain text, each non-blank line is processed separately.

All error responses use the shape: {"error": "<ExceptionClassName>", "message": "...", "details": {...}}.

Deployment

# 4-worker gunicorn
make api-prod

# Docker
make docker-build
make docker-run
make docker-run-gpu

API client demo

python example_client.py --host http://localhost:5000

Demonstrates all four endpoints including file upload for .txt and .csv.

Benchmarking

For paper comparisons, generate the fixed corrected held-out benchmark subset. It contains exactly 5,000 records sampled from data/test.jsonl with the current source-stratified preparation method and seed 42.

# Create or recreate the paper-comparison subset
python create_evaluation_subset.py

# Evaluate a trained model, the spaCy model used in PIIBench, and Presidio
python run_benchmarking.py \
  --test-file ./data/test_5k.jsonl \
  --model-path ./models/best_model \
  --system-name "Direct Fine-tuned DeBERTa" \
  --confidence-threshold 0.0 \
  --max-length 256 \
  --device cuda \
  --spacy-model en_core_web_lg \
  --output-dir ./benchmark_results/corrected_test_5k/direct_spacy_presidio

# Evaluate the six public neural baselines from PIIBench
python run_existing_models_benchmark.py \
  --test-path ./data/test_5k.jsonl \
  --output-path ./benchmark_results/corrected_test_5k/public_models.json \
  --device cuda

# After also evaluating the non-curriculum novelty model, compile the required ten rows
python compile_comparative_results.py

# After evaluating the curriculum variant, include it as an eleventh row
python compile_comparative_results.py \
  --curriculum-results ./benchmark_results/corrected_test_5k/curriculum/benchmark_results.json

Entity type normalisation maps are defined in run_benchmarking.py (SPACY_LABEL_MAP, PRESIDIO_LABEL_MAP) to align external systems to our label set before evaluation.

Novelty models are evaluated by running run_benchmarking.py again with each exported model and a distinct --system-name. Completed results and interpretation are in docs/novelty-finetuning.md. compile_comparative_results.py will fail unless Microsoft Presidio, spaCy, the six public neural comparators, and both trained models are all present on the identical held-out subset. The curriculum result is optional and is validated against the same held-out subset when supplied.

Full Test Evaluation For A Trained Model

Use the streaming evaluator for the complete prepared test split. It keeps only 5,000 input records in host memory at a time and uses smaller GPU minibatches within each chunk, which is appropriate for an 8 GB RTX 4070. If a CUDA out-of-memory error occurs, the current chunk is retried with a smaller inference minibatch.

python run_streaming_model_benchmark.py \
  --test-file ./data/test.jsonl \
  --model-path ../cloud_runs/baseline/models/best_model \
  --system-name "Direct Fine-tuned DeBERTa" \
  --confidence-threshold 0.0 \
  --max-length 256 \
  --device cuda \
  --chunk-size 5000 \
  --batch-size 8 \
  --output-dir ./benchmark_results/full_test/direct_deberta

python run_streaming_model_benchmark.py \
  --test-file ./data/test.jsonl \
  --model-path ../cloud_runs/novelty/best_model \
  --system-name "Source-conditioned Hierarchical DeBERTa" \
  --confidence-threshold 0.0 \
  --max-length 256 \
  --device cuda \
  --chunk-size 5000 \
  --batch-size 6 \
  --output-dir ./benchmark_results/full_test/source_conditioned_hierarchical

python analyze_full_test_comparison.py

The script writes cumulative progress after every chunk to streaming_progress.json and writes the final comparable metrics to benchmark_results.json. The analysis command validates the shared test hash and produces entity-level and hierarchical-group comparison tables.

Makefile targets

make install            Install dependencies
make download-deps      Download spaCy models for benchmarking
make data               Full data pipeline (download + consolidate + prepare)
make data SKIP=1        Skip download, re-run consolidation + preparation
make train              Download base model + fine-tune
make train SKIP=1       Skip model download, fine-tune only
make pipeline           data + train
make test               Run test suite
make benchmark          Benchmark all three systems
make benchmark N=500    Benchmark on first 500 records
make api                Start API server (dev mode)
make api-prod           Start API via gunicorn (4 workers)
make docker-build       Build Docker image
make docker-run         Run Docker container
make clean              Remove __pycache__ and .pyc files
make clean-models       Remove trained model checkpoints
make clean-data         Remove data splits
make clean-all          Remove everything generated

Custom exceptions

All exceptions inherit from PIIDetectionError and expose a to_dict() method used by the API error handler.

PIIDetectionError
├── ModelError
│   ├── ModelNotFoundError
│   ├── ModelLoadError
│   └── ModelInferenceError
├── InputError
│   ├── EmptyInputError
│   ├── InputTooLargeError
│   └── InvalidInputTypeError
├── FileParsingError
│   ├── UnsupportedFileTypeError
│   ├── FileReadError
│   ├── TextFileParsingError
│   ├── CSVParsingError
│   └── ExcelParsingError
└── APIError
    ├── MissingFieldError
    └── InvalidFieldError

Dataset on HuggingFace

• PII-Bench: https://huggingface.co/datasets/Pritesh-2711/pii-bench

Fine-tuned models on Huggingface

• Direct DeBERTa: https://huggingface.co/Pritesh-2711/piibench-deberta-base
• Source-conditioned hierarchical DeBERTa: https://huggingface.co/Pritesh-2711/piibench-deberta-sch

Support

If this work helps your research or engineering efforts, consider supporting ongoing open-source development and experimentation:

GitHub Sponsors: https://github.com/sponsors/pritesh-2711

Author

Pritesh Jha

License

Apache License version 2.0