diff --git a/backend/Dockerfile b/backend/Dockerfile index c6f4a36..20a9333 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -28,6 +28,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --frozen --no-install-project ARG SPACY_MODEL_WHEEL_URL="https://github.com/explosion/spacy-models/releases/download/en_core_web_lg-3.8.0/en_core_web_lg-3.8.0-py3-none-any.whl" +ARG SPACY_TOKENIZER_MODEL_WHEEL_URL="https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl" ENV PYTHONPATH=/app RUN apt-get update && apt-get install -y jq && rm -rf /var/lib/apt/lists/* @@ -44,19 +45,26 @@ COPY ./app /app/app RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --frozen -# Install pinned spaCy model in the final environment used at runtime. +# Install pinned spaCy models in the final environment used at runtime. RUN python -m pip install --no-deps "${SPACY_MODEL_WHEEL_URL}" +RUN python -m pip install --no-deps "${SPACY_TOKENIZER_MODEL_WHEEL_URL}" # Set HuggingFace cache directory ENV HF_HOME=/app/hf_cache -# Pre-download HuggingFace model +# Pre-download HuggingFace models RUN --mount=type=secret,id=HF_TOKEN \ HF_TOKEN="$(cat /run/secrets/HF_TOKEN 2>/dev/null || true)" \ /app/.venv/bin/python -c "from transformers import AutoTokenizer, AutoModelForSequenceClassification; \ AutoTokenizer.from_pretrained('textdetox/xlmr-large-toxicity-classifier', cache_dir='/app/hf_cache'); \ AutoModelForSequenceClassification.from_pretrained('textdetox/xlmr-large-toxicity-classifier', cache_dir='/app/hf_cache')" +RUN --mount=type=secret,id=HF_TOKEN \ + HF_TOKEN="$(cat /run/secrets/HF_TOKEN 2>/dev/null || true)" \ + /app/.venv/bin/python -c "from transformers import AutoTokenizer, AutoModelForTokenClassification; \ +AutoTokenizer.from_pretrained('dslim/bert-base-NER-uncased', cache_dir='/app/hf_cache'); \ +AutoModelForTokenClassification.from_pretrained('dslim/bert-base-NER-uncased', cache_dir='/app/hf_cache')" + # Pre-install Guardrails hub validators so container startup is not blocked by downloads RUN --mount=type=secret,id=GUARDRAILS_HUB_API_KEY \ GUARDRAILS_HUB_API_KEY="$(cat /run/secrets/GUARDRAILS_HUB_API_KEY 2>/dev/null || true)" \ diff --git a/backend/app/core/validators/README.md b/backend/app/core/validators/README.md index e31d14b..ab961ee 100644 --- a/backend/app/core/validators/README.md +++ b/backend/app/core/validators/README.md @@ -176,7 +176,9 @@ Recommendation: Parameters / customization: - `entity_types: list[str] | None` (default: all supported types) -- `threshold: float` (default: `0.5`) +- `threshold: float` (default: `0.5` for spaCy, `0.7` for transformers) +- `nlp_engine_type: str` (default: `"spacy"`) — NLP backend to use. `"spacy"` uses a spaCy model for both tokenization and NER; `"transformers"` uses a spaCy model for tokenization only and a HuggingFace token-classification model for NER. +- `model_name: str | None` (default: `None`) — Override the NLP model. For `"spacy"`, a spaCy model name (resolved to `en_core_web_lg`); for `"transformers"`, a HuggingFace model ID (resolved to `dslim/bert-base-NER-uncased`). - `on_fail` Threshold guidance: @@ -184,7 +186,8 @@ Threshold guidance: - `threshold` is the minimum confidence score required for a detected entity to be treated as PII. - Lower threshold -> more detections (higher recall, more false positives/over-masking). - Higher threshold -> fewer detections (higher precision, more false negatives/missed PII). -- Start around `0.5`, then tune using real conversation samples by reviewing both missed PII and unnecessary masking. +- For spaCy mode, start around `0.5`. For transformers mode, start around `0.7` as transformer model scores tend to be higher overall. +- Tune using real conversation samples by reviewing both missed PII and unnecessary masking. - If the product is privacy-critical, prefer a slightly lower threshold and tighter `entity_types`; if readability is primary, prefer a slightly higher threshold. Supported default entity types: @@ -195,8 +198,10 @@ Notes / limitations: - Rule/ML recognizers can under-detect free-text references. - Threshold and entity selection should be tuned per deployment context. -- Runtime requirement: this validator is configured to use spaCy model `en_core_web_lg`. - The model is pre-installed at build time in the Docker image to ensure fast startup and no runtime internet dependency. +- Runtime requirements: + - `"spacy"` mode (default): requires `en_core_web_lg`. Pre-installed in the Docker image at build time for fast startup with no runtime internet dependency. + - `"transformers"` mode: requires `en_core_web_sm` for tokenization (lightweight) and the specified HuggingFace model, which is also pre-installed in the Docker image. +- The `"transformers"` mode uses CoNLL-style label mapping (`PER` → `PERSON`, `LOC` → `LOCATION`). Models that use different label schemes will silently produce no detections — verify label compatibility before switching models. Evidence and evaluation: - Compared approaches: - Custom PII validator (this codebase) @@ -534,7 +539,7 @@ Notes / limitations: ## Example Config Payloads -Example: create validator config (stored shape) +Example: create validator config (stored shape) — spaCy mode (default) ```json { @@ -547,6 +552,21 @@ Example: create validator config (stored shape) } ``` +Example: create validator config — transformers mode with a custom HuggingFace model + +```json +{ + "type": "pii_remover", + "stage": "input", + "on_fail_action": "fix", + "is_enabled": true, + "entity_types": ["PERSON", "PHONE_NUMBER", "LOCATION"], + "nlp_engine_type": "transformers", + "model_name": "dslim/bert-base-NER-uncased", + "threshold": 0.7 +} +``` + Example: runtime guardrail validator object (execution shape) ```json @@ -554,6 +574,7 @@ Example: runtime guardrail validator object (execution shape) "type": "pii_remover", "on_fail": "fix", "entity_types": ["PERSON", "PHONE_NUMBER", "IN_AADHAAR"], + "nlp_engine_type": "spacy", "threshold": 0.6 } ``` diff --git a/backend/app/core/validators/config/pii_remover_safety_validator_config.py b/backend/app/core/validators/config/pii_remover_safety_validator_config.py index c4bb793..bd7fcd6 100644 --- a/backend/app/core/validators/config/pii_remover_safety_validator_config.py +++ b/backend/app/core/validators/config/pii_remover_safety_validator_config.py @@ -7,12 +7,16 @@ class PIIRemoverSafetyValidatorConfig(BaseValidatorConfig): type: Literal["pii_remover"] - entity_types: Optional[List[str]] = None # list of PII entity types to remove - threshold: float = 0.5 # confidence threshold for PII detection + entity_types: Optional[List[str]] = None + threshold: Optional[float] = None + nlp_engine_type: str = "spacy" + model_name: Optional[str] = None def build(self): return PIIRemover( entity_types=self.entity_types, threshold=self.threshold, + nlp_engine_type=self.nlp_engine_type, + model_name=self.model_name, on_fail=self.resolve_on_fail(), ) diff --git a/backend/app/core/validators/pii_remover.py b/backend/app/core/validators/pii_remover.py index efe8ad0..6dc4801 100644 --- a/backend/app/core/validators/pii_remover.py +++ b/backend/app/core/validators/pii_remover.py @@ -1,6 +1,7 @@ from __future__ import annotations import os -from typing import Callable, Optional +from collections.abc import Callable +from typing import Optional from guardrails import OnFailAction from guardrails.validators import ( @@ -10,24 +11,14 @@ ValidationResult, Validator, ) -from presidio_analyzer import AnalyzerEngine -from presidio_analyzer.nlp_engine import NlpEngineProvider +from presidio_analyzer import AnalyzerEngine, EntityRecognizer, RecognizerResult +from presidio_analyzer.nlp_engine import NlpEngineProvider, SpacyNlpEngine from presidio_anonymizer import AnonymizerEngine -from presidio_analyzer.predefined_recognizers.country_specific.india.in_aadhaar_recognizer import ( - InAadhaarRecognizer, -) -from presidio_analyzer.predefined_recognizers.country_specific.india.in_pan_recognizer import ( - InPanRecognizer, -) -from presidio_analyzer.predefined_recognizers.country_specific.india.in_passport_recognizer import ( - InPassportRecognizer, -) -from presidio_analyzer.predefined_recognizers.country_specific.india.in_vehicle_registration_recognizer import ( - InVehicleRegistrationRecognizer, -) -from presidio_analyzer.predefined_recognizers.country_specific.india.in_voter_recognizer import ( - InVoterRecognizer, -) +from presidio_analyzer.predefined_recognizers.country_specific.india.in_aadhaar_recognizer import InAadhaarRecognizer +from presidio_analyzer.predefined_recognizers.country_specific.india.in_pan_recognizer import InPanRecognizer +from presidio_analyzer.predefined_recognizers.country_specific.india.in_passport_recognizer import InPassportRecognizer +from presidio_analyzer.predefined_recognizers.country_specific.india.in_vehicle_registration_recognizer import InVehicleRegistrationRecognizer +from presidio_analyzer.predefined_recognizers.country_specific.india.in_voter_recognizer import InVoterRecognizer os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -49,13 +40,31 @@ "IN_VOTER", ] -CONFIGURATION = { +SPACY_CONFIGURATION = { "nlp_engine_name": "spacy", "models": [{"lang_code": "en", "model_name": "en_core_web_lg"}], } -_GLOBAL_NLP_ENGINE = None -_ANALYZER_CACHE = {} +# Lightweight spaCy model used only for tokenization when the transformers engine handles NER +SPACY_TOKENIZER_CONFIGURATION = { + "nlp_engine_name": "spacy", + "models": [{"lang_code": "en", "model_name": "en_core_web_sm"}], +} + +# Shared label mapping for models using CoNLL-style PER/ORG/LOC/MISC labels +# (e.g. dslim/bert-base-NER, Davlan/xlm-roberta-base-ner-hrl) +CONLL_NER_LABEL_MAPPING = { + "PER": "PERSON", + "ORG": "ORGANIZATION", + "LOC": "LOCATION", + "MISC": "MISC", +} + +DEFAULT_TRANSFORMERS_MODEL = "dslim/bert-base-NER-uncased" +DEFAULT_TRANSFORMERS_THRESHOLD = 0.7 + +_NLP_ENGINE_CACHE: dict = {} +_ANALYZER_CACHE: dict = {} INDIA_RECOGNIZERS = { @@ -67,23 +76,97 @@ } -def _build_analyzer(entity_types: list[str]) -> AnalyzerEngine: - global _GLOBAL_NLP_ENGINE - if _GLOBAL_NLP_ENGINE is None: - provider = NlpEngineProvider(nlp_configuration=CONFIGURATION) - _GLOBAL_NLP_ENGINE = provider.create_engine() - analyzer = AnalyzerEngine(nlp_engine=_GLOBAL_NLP_ENGINE) +class HuggingFaceNERRecognizer(EntityRecognizer): + """Presidio EntityRecognizer backed by a HuggingFace token-classification pipeline.""" + + def __init__(self, model_name: str, label_mapping: dict[str, str], threshold: float = 0.5): + supported_entities = list(set(label_mapping.values())) + super().__init__(supported_entities=supported_entities, name="HuggingFaceNERRecognizer") + from transformers import pipeline as hf_pipeline + + self.ner_pipeline = hf_pipeline( + "token-classification", + model=model_name, + aggregation_strategy="simple", + ) + self.label_mapping = label_mapping + self.threshold = threshold + + def load(self) -> None: + pass + + def analyze(self, text: str, entities: list[str], nlp_artifacts=None) -> list[RecognizerResult]: + results: list[RecognizerResult] = [] + for ent in self.ner_pipeline(text): + presidio_label = self.label_mapping.get(ent["entity_group"]) + if presidio_label is None or presidio_label not in entities: + continue + if ent["score"] < self.threshold: + continue + results.append( + RecognizerResult( + entity_type=presidio_label, + start=ent["start"], + end=ent["end"], + score=float(ent["score"]), + ) + ) + return results + + +def _build_spacy_engine(configuration: dict) -> SpacyNlpEngine: + provider = NlpEngineProvider(nlp_configuration=configuration) + return provider.create_engine() # type: ignore[return-value] + + +def _build_analyzer( + entity_types: list[str], nlp_engine_type: str, model_name: str, threshold: float +) -> AnalyzerEngine: + if nlp_engine_type == "transformers": + # Use en_core_web_sm only for tokenization; BERT handles NER + if "spacy_tokenizer_engine" not in _NLP_ENGINE_CACHE: + _NLP_ENGINE_CACHE["spacy_tokenizer_engine"] = _build_spacy_engine(SPACY_TOKENIZER_CONFIGURATION) + nlp_engine = _NLP_ENGINE_CACHE["spacy_tokenizer_engine"] + else: + if "spacy_engine" not in _NLP_ENGINE_CACHE: + _NLP_ENGINE_CACHE["spacy_engine"] = _build_spacy_engine(SPACY_CONFIGURATION) + nlp_engine = _NLP_ENGINE_CACHE["spacy_engine"] + + analyzer = AnalyzerEngine(nlp_engine=nlp_engine) + + if nlp_engine_type == "transformers": + # Remove SpacyRecognizer so spaCy NER doesn't run alongside BERT + analyzer.registry.remove_recognizer("SpacyRecognizer") + + hf_recognizer_key = ("hf_recognizer", model_name) + if hf_recognizer_key not in _NLP_ENGINE_CACHE: + _NLP_ENGINE_CACHE[hf_recognizer_key] = HuggingFaceNERRecognizer( + model_name=model_name, + label_mapping=CONLL_NER_LABEL_MAPPING, + threshold=threshold, + ) + analyzer.registry.add_recognizer(_NLP_ENGINE_CACHE[hf_recognizer_key]) for entity_type, recognizer_cls in INDIA_RECOGNIZERS.items(): if entity_type in entity_types: analyzer.registry.add_recognizer(recognizer_cls()) + return analyzer -def _get_cached_analyzer(entity_types: list[str]) -> AnalyzerEngine: - recognizer_key = tuple(sorted(t for t in entity_types if t in INDIA_RECOGNIZERS)) +def _get_cached_analyzer( + entity_types: list[str], nlp_engine_type: str, model_name: str, threshold: float +) -> AnalyzerEngine: + recognizer_key = ( + nlp_engine_type, + model_name, + threshold, + tuple(sorted(t for t in entity_types if t in INDIA_RECOGNIZERS)), + ) if recognizer_key not in _ANALYZER_CACHE: - _ANALYZER_CACHE[recognizer_key] = _build_analyzer(entity_types) + _ANALYZER_CACHE[recognizer_key] = _build_analyzer( + entity_types, nlp_engine_type, model_name, threshold + ) return _ANALYZER_CACHE[recognizer_key] @@ -98,15 +181,25 @@ class PIIRemover(Validator): def __init__( self, entity_types=None, - threshold=0.5, + threshold: float | None = None, + nlp_engine_type: str = "spacy", + model_name: str | None = None, on_fail: Optional[Callable] = OnFailAction.FIX, ): super().__init__(on_fail=on_fail) self.entity_types = entity_types or ALL_ENTITY_TYPES - self.threshold = threshold + self.nlp_engine_type = nlp_engine_type + if nlp_engine_type == "transformers": + self.model_name = model_name or DEFAULT_TRANSFORMERS_MODEL + self.threshold = threshold if threshold is not None else DEFAULT_TRANSFORMERS_THRESHOLD + else: + self.model_name = model_name or "en_core_web_lg" + self.threshold = threshold if threshold is not None else 0.5 self.on_fail = on_fail - self.analyzer = _get_cached_analyzer(self.entity_types) + self.analyzer = _get_cached_analyzer( + self.entity_types, self.nlp_engine_type, self.model_name, self.threshold + ) self.anonymizer = AnonymizerEngine() def _validate(self, value: str, metadata: dict | None = None) -> ValidationResult: @@ -122,3 +215,4 @@ def _validate(self, value: str, metadata: dict | None = None) -> ValidationResul error_message="PII detected in the text.", fix_value=anonymized_text ) return PassResult(value=text) + diff --git a/backend/app/tests/validators/test_pii_remover.py b/backend/app/tests/validators/test_pii_remover.py index 1b8f757..1fbf624 100644 --- a/backend/app/tests/validators/test_pii_remover.py +++ b/backend/app/tests/validators/test_pii_remover.py @@ -3,7 +3,8 @@ import pytest from app.core.validators import pii_remover -from app.core.validators.pii_remover import ALL_ENTITY_TYPES, PIIRemover +from app.core.validators.pii_remover import ALL_ENTITY_TYPES, DEFAULT_TRANSFORMERS_MODEL, DEFAULT_TRANSFORMERS_THRESHOLD, PIIRemover +from app.core.validators.config.pii_remover_safety_validator_config import PIIRemoverSafetyValidatorConfig # ------------------------------- # Fixtures @@ -79,22 +80,54 @@ def test_custom_entity_types_override(mock_presidio): assert v.entity_types == ["EMAIL_ADDRESS"] +def test_transformers_engine_uses_correct_defaults(mock_presidio): + v = PIIRemover(nlp_engine_type="transformers") + assert v.model_name == DEFAULT_TRANSFORMERS_MODEL + assert v.threshold == DEFAULT_TRANSFORMERS_THRESHOLD + + +def test_spacy_engine_uses_correct_defaults(mock_presidio): + v = PIIRemover(nlp_engine_type="spacy") + assert v.model_name == "en_core_web_lg" + assert v.threshold == 0.5 + + +def test_transformers_engine_accepts_custom_model(mock_presidio): + v = PIIRemover(nlp_engine_type="transformers", model_name="dslim/bert-base-NER-uncased") + assert v.model_name == "dslim/bert-base-NER-uncased" + + +def test_config_builds_validator_with_nlp_engine_params(mock_presidio): + config = PIIRemoverSafetyValidatorConfig( + type="pii_remover", + nlp_engine_type="transformers", + model_name="dslim/bert-base-NER-uncased", + ) + v = config.build() + assert v.nlp_engine_type == "transformers" + assert v.model_name == "dslim/bert-base-NER-uncased" + + +def test_config_defaults_to_spacy_engine(mock_presidio): + config = PIIRemoverSafetyValidatorConfig(type="pii_remover") + v = config.build() + assert v.nlp_engine_type == "spacy" + assert v.model_name == "en_core_web_lg" + + def test_cached_analyzer_registers_only_requested_indian_recognizers(): with patch( - "app.core.validators.pii_remover.NlpEngineProvider" - ) as mock_provider, patch( + "app.core.validators.pii_remover._build_spacy_engine" + ) as mock_build_engine, patch( "app.core.validators.pii_remover.AnalyzerEngine" ) as mock_analyzer: pii_remover._ANALYZER_CACHE.clear() - pii_remover._GLOBAL_NLP_ENGINE = None + pii_remover._NLP_ENGINE_CACHE.clear() analyzer_instance = mock_analyzer.return_value - pii_remover._get_cached_analyzer(["EMAIL_ADDRESS", "IN_AADHAAR", "IN_PAN"]) - pii_remover._get_cached_analyzer(["EMAIL_ADDRESS", "IN_AADHAAR", "IN_PAN"]) + pii_remover._get_cached_analyzer(["EMAIL_ADDRESS", "IN_AADHAAR", "IN_PAN"], "spacy", "en_core_web_lg", 0.5) + pii_remover._get_cached_analyzer(["EMAIL_ADDRESS", "IN_AADHAAR", "IN_PAN"], "spacy", "en_core_web_lg", 0.5) - mock_provider.assert_called_once_with( - nlp_configuration=pii_remover.CONFIGURATION - ) - mock_provider.return_value.create_engine.assert_called_once() + mock_build_engine.assert_called_once_with(pii_remover.SPACY_CONFIGURATION) mock_analyzer.assert_called_once() assert analyzer_instance.registry.add_recognizer.call_count == 2