johngunerli/pii-detector

PII Detector: Named Entity Recognition (NER) for Personally Identifiable Information (PII)

Overview

This project implements a Named Entity Recognition (NER) model to detect Personally Identifiable Information (PII) using a fine-tuned DistilBERT model. It identifies various types of PII such as names, emails, usernames, ID numbers, phone numbers, URLs, and addresses in text.

Features

• Synthetic data generation for PII-related entities.
• Token classification using BIO tagging format.
• Fine-tuning of DistilBERT for PII detection.
• Model training with Hugging Face's Trainer API.
• Inference pipeline for real-time PII detection.

Dataset

The dataset is synthetically generated using the Faker library. It includes:

• Student names
• Emails
• Usernames
• ID numbers
• Phone numbers
• Personal URLs
• Street addresses

Each sentence is labeled with corresponding entities in BIO tagging format.

Installation

To set up the project, install the necessary dependencies:

pip install torch transformers datasets faker

Usage

1. Generate Synthetic Data

Run the generate_synthetic_data function to create labeled text samples with PII entities.

2. Tokenize and Align Labels

The function tokenize_and_align_labels tokenizes input text and aligns the entity labels using Hugging Face's tokenizer.

3. Train the Model

Execute the training pipeline using:

trainer.train()

This will fine-tune DistilBERT on the labeled dataset.

4. Save the Model

The trained model is saved using:

trainer.save_model("./pii_detector")

5. Run Inference

To detect PII in a given text, use:

pii_detection("Sample text with PII information")

This will return identified entities along with their labels.

Model Configuration

• Base Model: distilbert-base-uncased
• Tokenizer: AutoTokenizer
• Training Parameters:
  • Batch size: 16
  • Number of epochs: 3
  • Evaluation strategy: Per epoch
  • Device: CUDA (if available)

Output

The model returns a list of detected PII entities with their respective labels and positions in the text.

License

This project is open-source and can be used for educational and research purposes.