privacy-filter.py

privacy-filter.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "transformers>=4.50",
#   "torch",
# ]
# ///
"""
Run openai/privacy-filter over text passed as the first argument.

Usage: ./privacy_filter.py "My name is Stephen and my phone is 555-1234"
"""
import sys
from transformers import pipeline

if len(sys.argv) < 2:
    sys.exit("usage: privacy_filter.py TEXT")

text = sys.argv[1]

classifier = pipeline(
    task="token-classification",
    model="openai/privacy-filter",
    aggregation_strategy="simple",
)

spans = classifier(text)

print(f"Input: {text}\n")

if not spans:
    print("No PII detected.")
    sys.exit(0)

print(f"Detected {len(spans)} span(s):")
for s in spans:
    span_text = text[s["start"]:s["end"]]
    print(f"  [{s['entity_group']}] {span_text!r}  (score={s['score']:.3f})")

redacted = text
for s in sorted(spans, key=lambda x: x["start"], reverse=True):
    redacted = redacted[:s["start"]] + f"[{s['entity_group'].upper()}]" + redacted[s["end"]:]

print(f"\nRedacted: {redacted}")