This project implements a PII (Personally Identifiable Information) detection and masking system using DistilBERT fine-tuned on the ai4privacy/pii-masking-200k dataset. The system exposes a Flask API for uploading text files and returning masked outputs.
To Train the Model run every cell in the file NER_Masking.ipynb
▶️ Run the Application
To run the Flask API:
python app.py
The server will start locally (default: http://127.0.0.1:5000).
📤 API Usage
Nice. You now have two clean portals into your PII engine, like two doors to the same vault, one for raw text, one for files. Here is a concise explanation you can add to your README under an API Endpoints section.
🌐 API Endpoints
1️⃣ /predict
Method:POSTDescription: Performs PII detection on raw text input.
Request Body (JSON):
json
1{2"text":"Your input text here"3}
Response:
json
1{2"masked":"Masked text output",3"highlighted":"<html with highlighted entities>"4}
Calls pii_inference(text)
Returns both masked text and dynamically highlighted HTML output
2️⃣ /upload
Method:POSTDescription: Uploads a .txt file and processes it in batches of 5 lines.
Form-Data Key:
file
Processing Logic:
Reads file
Splits into lines
Processes 5 lines at a time
Runs pii_inference on each batch
Merges results into final masked output
Generates highlighted HTML
Response:
json
1{2"masked":"Final masked text",3"highlighted":"<html with highlighted entities>"4}
🔍 Design Insight
/predict → Low latency, single inference call
/upload → Memory-efficient batch processing
Batch size (5 lines) prevents long-sequence instability in transformer inference
🧠 Model Details
Base Model: DistilBERT
Dataset: ai4privacy/pii-masking-200k
Training Framework: Hugging Face Trainer
Batch Size: 32
Epochs: 10
Mixed Precision (FP16): Enabled
📊 Performance
Precision: 90.86%
Recall: 93.50%
F1 Score: 92.16%
Strong performance on structured PII types such as Email, URL, SSN, and Username.
