from veildata.core import Module
from veildata.revealers import TokenStore
try:
import spacy
except ImportError as e:
raise ImportError(
"spaCy is not installed. Install with: `pip install veildata[spacy]`"
) from e
[docs]
class SpacyNERRedactor(Module):
"""Redact named entities in text using a spaCy model, with optional reversible tracking."""
def __init__(
self,
model: str = "en_core_web_sm",
entities: list[str] | None = None,
redaction_token: str = "[REDACTED_{counter}]",
store: TokenStore | None = None,
) -> None:
super().__init__()
self.model_name = model
self.entities = set(entities or ["PERSON", "ORG", "GPE", "EMAIL", "PHONE"])
self.redaction_token = redaction_token
self.store = store
self._load_model()
self.counter = 0
def _load_model(self) -> None:
try:
self.nlp = spacy.load(self.model_name, disable=["parser", "tagger"])
except OSError:
raise RuntimeError(
f"spaCy model '{self.model_name}' not found. "
f"Run: python -m spacy download {self.model_name}"
)
[docs]
def forward(self, text: str) -> str:
doc = self.nlp(text)
redacted = text
for ent in reversed(doc.ents):
if ent.label_ in self.entities:
self.counter += 1
token = self.redaction_token.format(counter=self.counter)
if self.store:
self.store.record(token, ent.text)
redacted = redacted[: ent.start_char] + token + redacted[ent.end_char :]
return redacted