Spaces:
Runtime error
Runtime error
File size: 2,634 Bytes
a08e3d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import re
import spacy
from typing import List, Dict, Tuple
# Load spaCy model for NER (en_core_web_sm)
nlp = spacy.load("en_core_web_sm")
# Regex patterns for PII fields
PATTERNS = {
"email": re.compile(r"[a-zA-Z0-9+._%-]+@[a-zA-Z0-9._%-]+\.[a-zA-Z]{2,6}"),
"phone_number": re.compile(r"\b(?:\+?\d{1,3}[ -]?)?(?:\(\d{3}\)|\d{3})[ -]?\d{3}[ -]?\d{4}\b"),
"dob": re.compile(r"\b\d{2}[/-]\d{2}[/-]\d{4}\b"),
"aadhar_num": re.compile(r"\b\d{4}[ -]?\d{4}[ -]?\d{4}\b"),
"credit_debit_no": re.compile(r"\b\d{4}[ -]?\d{4}[ -]?\d{4}[ -]?\d{4}\b"),
"cvv_no": re.compile(r"\b\d{3}\b"),
"expiry_no": re.compile(r"\b(0[1-9]|1[0-2])\/[0-9]{2}\b"),
}
Entity = Dict[str, object]
def mask_pii(text: str) -> Tuple[str, List[Entity]]:
"""
Masks PII in the text. Returns masked_text and list of entities with their positions and original values.
"""
entities: List[Entity] = []
masked_text = text
# 1) Regex-based masking
for ent_type, pattern in PATTERNS.items():
for match in pattern.finditer(masked_text):
start, end = match.span()
original = match.group()
entities.append({
"position": [start, end],
"classification": ent_type,
"entity": original,
})
masked_text = masked_text[:start] + f"[{ent_type}]" + masked_text[end:]
# 2) spaCy NER for PERSON and DATE
doc = nlp(masked_text)
offset = 0
for ent in doc.ents:
label = ent.label_
if label in ("PERSON", "DATE"):
ent_type = "full_name" if label == "PERSON" else "dob"
start, end = ent.start_char, ent.end_char
original = masked_text[start:end]
entities.append({
"position": [start, end],
"classification": ent_type,
"entity": original,
})
masked_text = masked_text[:start] + f"[{ent_type}]" + masked_text[end:]
# Sort entities by start index
entities.sort(key=lambda x: x["position"][0])
return masked_text, entities
def demask_pii(masked_text: str, entities: List[Entity]) -> str:
"""
Restores original PII values into the masked_text using stored entity list.
"""
result = masked_text
# Apply replacements in reverse order to preserve indices
for ent in sorted(entities, key=lambda x: x["position"][0], reverse=True):
start, end = ent["position"]
original = ent["entity"]
# Replace the [ent_type] placeholder back to original
result = result[:start] + original + result[end:]
return result |