Spaces:
Runtime error
Runtime error
| 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 |