PII Removal › Code Implementation

# Install spacy model if not already present
# python -m spacy download en_core_web_sm

@dataclass
class PIISpan:
    start: int
    end: int
    pii_type: str
    text: str


@dataclass
class PIIResult:
    original_text: str
    redacted_text: str
    spans: List[PIISpan] = field(default_factory=list)

# Compiled regex patterns for common PII types
PII_PATTERNS = {
    "EMAIL": re.compile(
        r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b"
    ),
    "PHONE": re.compile(
        r"\b(?:\+1[\s\-]?)?\(?(\d{3})\)?[\s\-]?(\d{3})[\s\-]?(\d{4})\b"
    ),
    "SSN": re.compile(
        r"\b(?!000|666|9\d{2})\d{3}[\s\-](?!00)\d{2}[\s\-](?!0000)\d{4}\b"
    ),
    "CREDIT_CARD": re.compile(
        r"\b(?:4\d{3}[\s\-]?\d{4}[\s\-]?\d{4}[\s\-]?\d{4}|"
        r"5[1-5]\d{2}[\s\-]?\d{4}[\s\-]?\d{4}[\s\-]?\d{4}|"
        r"3[47]\d{2}[\s\-]?\d{6}[\s\-]?\d{5})\b"
    ),
    "IP_ADDRESS": re.compile(
        r"\b(?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}"
        r"(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\b"
    ),
}


def detect_pattern_pii(text: str) -> List[PIISpan]:
    spans = []
    for pii_type, pattern in PII_PATTERNS.items():
        for match in pattern.finditer(text):
            spans.append(
                PIISpan(
                    start=match.start(),
                    end=match.end(),
                    pii_type=pii_type,
                    text=match.group(),
                )
            )
    return spans

# Load spaCy model for NER
# Requires: python -m spacy download en_core_web_sm
try:
    nlp = spacy.load("en_core_web_sm")
    ner_available = True
except OSError:
    ner_available = False


def detect_name_pii(text: str) -> List[PIISpan]:
    if not ner_available:
        return []
    doc = nlp(text)
    spans = []
    for ent in doc.ents:
        if ent.label_ == "PERSON":
            spans.append(
                PIISpan(
                    start=ent.start_char,
                    end=ent.end_char,
                    pii_type="NAME",
                    text=ent.text,
                )
            )
    return spans

def merge_spans(spans: List[PIISpan]) -> List[PIISpan]:
    """Remove overlapping spans, preferring longer matches."""
    if not spans:
        return []
    sorted_spans = sorted(spans, key=lambda s: (s.start, -(s.end - s.start)))
    merged = [sorted_spans[0]]
    for span in sorted_spans[1:]:
        last = merged[-1]
        if span.start < last.end:
            # Overlap: keep longer span (already sorted by length descending)
            continue
        merged.append(span)
    return merged


def redact_text(text: str, spans: List[PIISpan]) -> str:
    """Replace PII spans with placeholder tokens."""
    # Process spans in reverse order to preserve character offsets
    result = list(text)
    for span in sorted(spans, key=lambda s: s.start, reverse=True):
        placeholder = f"[{span.pii_type}]"
        result[span.start : span.end] = list(placeholder)
    return "".join(result)


def process_document(text: str) -> PIIResult:
    """Run full PII detection and redaction pipeline."""
    pattern_spans = detect_pattern_pii(text)
    name_spans = detect_name_pii(text)
    all_spans = merge_spans(pattern_spans + name_spans)
    redacted = redact_text(text, all_spans)
    return PIIResult(
        original_text=text,
        redacted_text=redacted,
        spans=all_spans,
    )

def compute_pii_density(text: str, spans: List[PIISpan]) -> float:
    """Compute fraction of characters covered by PII spans."""
    if not text:
        return 0.0
    pii_chars = sum(span.end - span.start for span in spans)
    return pii_chars / len(text)


def should_remove_document(
    text: str, threshold: float = 0.10
) -> Tuple[bool, float]:
    """Decide whether a document should be removed due to high PII density."""
    result = process_document(text)
    density = compute_pii_density(text, result.spans)
    return density >= threshold, density

def generate_synthetic_examples(
    n_per_type: int = 50, seed: int = 42
) -> List[Dict]:
    """Generate synthetic test cases with known PII ground truth."""
    rng = random.Random(seed)

    # Templates for each category
    email_templates = [
        "Contact {pii} for assistance.",
        "Reach out to {pii} with questions.",
        "The account email is {pii}.",
        "Send invoices to {pii} by Friday.",
    ]
    phone_templates = [
        "Call {pii} during business hours.",
        "The helpline number is {pii}.",
        "Text {pii} to opt out.",
        "Fax documents to {pii}.",
    ]
    ssn_templates = [
        "Enter your SSN: {pii}.",
        "Reference number {pii} on your W-2.",
        "Taxpayer ID {pii} is required.",
    ]

    def rand_email():
        user = "".join(
            rng.choices(string.ascii_lowercase, k=rng.randint(5, 10))
        )
        domain = rng.choice(
            ["gmail.com", "yahoo.com", "example.org", "company.net"]
        )
        return f"{user}@{domain}"

    def rand_phone():
        area = rng.randint(200, 999)
        exchange = rng.randint(200, 999)
        line = rng.randint(1000, 9999)
        return f"({area}) {exchange}-{line}"

    def rand_ssn():
        area = rng.randint(100, 665)
        group = rng.randint(10, 99)
        serial = rng.randint(1000, 9999)
        return f"{area:03d}-{group:02d}-{serial:04d}"

    examples = []
    generators = [
        ("EMAIL", rand_email, email_templates),
        ("PHONE", rand_phone, phone_templates),
        ("SSN", rand_ssn, ssn_templates),
    ]

    for pii_type, gen_fn, templates in generators:
        for _ in range(n_per_type):
            pii_value = gen_fn()
            template = rng.choice(templates)
            text = template.format(pii=pii_value)
            examples.append(
                {
                    "text": text,
                    "pii_type": pii_type,
                    "pii_value": pii_value,
                }
            )

    # Add negative examples (no PII)
    negatives = [
        "The meeting is scheduled for Thursday at noon.",
        "Please review the attached report before Monday.",
        "The quarterly results exceed analyst expectations.",
        "All employees must complete the training module.",
        "The new policy takes effect on the first of next month.",
    ]
    for neg_text in negatives:
        examples.append({"text": neg_text, "pii_type": None, "pii_value": None})

    return examples


test_examples = generate_synthetic_examples(n_per_type=50)

def evaluate_pipeline(examples: List[Dict]) -> Dict:
    """Compute precision and recall per PII type."""
    results = {
        pii_type: {"tp": 0, "fp": 0, "fn": 0}
        for pii_type in ["EMAIL", "PHONE", "SSN"]
    }
    false_positives_on_negatives = 0

    for ex in examples:
        result = process_document(ex["text"])
        detected_types = {s.pii_type for s in result.spans}

        if ex["pii_type"] is None:
            # Negative example: any detection is a false positive
            if detected_types:
                false_positives_on_negatives += len(detected_types)
        else:
            pii_type = ex["pii_type"]
            if pii_type in detected_types:
                results[pii_type]["tp"] += 1
            else:
                results[pii_type]["fn"] += 1
            # Count unexpected detections of other types as FP
            for dtype in detected_types:
                if dtype != pii_type:
                    results.get(dtype, {})["fp"] = (
                        results.get(dtype, {}).get("fp", 0) + 1
                    )

    metrics = {}
    for pii_type, counts in results.items():
        tp, fp, fn = counts["tp"], counts["fp"], counts["fn"]
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = (
            2 * precision * recall / (precision + recall)
            if (precision + recall) > 0
            else 0.0
        )
        metrics[pii_type] = {
            "precision": precision,
            "recall": recall,
            "f1": f1,
            "tp": tp,
            "fp": fp,
            "fn": fn,
        }

    metrics["FALSE_POSITIVES_ON_NEGATIVES"] = false_positives_on_negatives
    return metrics


eval_metrics = evaluate_pipeline(test_examples)