fix: allow more flexibility in choosing models used for PII removal

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)" \

backend/app/core/validators/README.md

@@ -176,15 +176,18 @@ 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:
 
 - `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,13 +552,29 @@ 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
 {
   "type": "pii_remover",
   "on_fail": "fix",
   "entity_types": ["PERSON", "PHONE_NUMBER", "IN_AADHAAR"],
+  "nlp_engine_type": "spacy",
   "threshold": 0.6
 }
 ```

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(),
         )

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)
+