From 3f5d4c5fd53d66099d33353a5eb7027d26a229e6 Mon Sep 17 00:00:00 2001 From: Yurii Havrylko <10372700+yuriihavrylko@users.noreply.github.com> Date: Fri, 12 Jun 2026 21:56:25 +0200 Subject: [PATCH 1/6] feat: add onnxruntime support to HuggingFaceNerRecognizer --- .../ner/huggingface_ner_recognizer.py | 130 +++++++++++-- presidio-analyzer/pyproject.toml | 5 + .../tests/test_huggingface_ner_recognizer.py | 181 +++++++++++++++++- .../test_huggingface_ner_recognizer_e2e.py | 165 ++++++++++++++++ 4 files changed, 454 insertions(+), 27 deletions(-) create mode 100644 presidio-analyzer/tests/test_huggingface_ner_recognizer_e2e.py diff --git a/presidio-analyzer/presidio_analyzer/predefined_recognizers/ner/huggingface_ner_recognizer.py b/presidio-analyzer/presidio_analyzer/predefined_recognizers/ner/huggingface_ner_recognizer.py index 43b17683f1..290c1396ac 100644 --- a/presidio-analyzer/presidio_analyzer/predefined_recognizers/ner/huggingface_ner_recognizer.py +++ b/presidio-analyzer/presidio_analyzer/predefined_recognizers/ner/huggingface_ner_recognizer.py @@ -35,6 +35,11 @@ except ImportError: torch = None +try: + from optimum.pipelines import pipeline as optimum_pipeline +except ImportError: + optimum_pipeline = None + logger = logging.getLogger("presidio-analyzer") @@ -117,7 +122,8 @@ def __init__( tokenizer_name: Optional[str] = None, text_chunker: Optional[BaseTextChunker] = None, label_prefixes: Optional[List[str]] = None, - **kwargs, + backend: str = "torch", + **model_kwargs, ): """Initialize the HuggingFace NER Recognizer. @@ -136,30 +142,63 @@ def __init__( ("simple", "first", "average", "max"). Recommendation: Use "simple" or "first" so that entities are pre-aggregated by the model, preserving performance and alignment. - :param device: Device to use. Accepts: + :param device: Device to use ("torch" backend only). Accepts: - "cpu" or -1 for CPU - "cuda" or "cuda:N" or int N for GPU - None for auto-detection (GPU if available, else CPU) - Defaults to None. + Defaults to None. Ignored by the "ort" backend — select + hardware there via the `provider` model kwarg. :param chunk_overlap: Number of characters to overlap between chunks. :param chunk_size: Maximum number of characters per chunk. :param tokenizer_name: Name of the tokenizer. Defaults to model_name. :param text_chunker: Custom text chunking strategy. If None, uses CharacterBasedTextChunker with provided chunk_size and chunk_overlap. :param label_prefixes: List of label prefixes to strip (e.g., B-, I-). - :raises ImportError: If transformers or torch libraries are not installed. + :param backend: Inference backend to use. + - "torch" (default): PyTorch via transformers pipeline. + Requires: torch, transformers. + - "ort": ONNX Runtime via optimum. + Requires: optimum, optimum-onnx[onnxruntime]. + For NVIDIA GPU, install `onnxruntime-gpu` and pass + `provider="CUDAExecutionProvider"` via model_kwargs. + :param model_kwargs: Additional keyword arguments forwarded to the + underlying model loader. + For the "torch" backend, passed to `transformers.pipeline` as + `model_kwargs=`. For "ort", passed to + `ORTModel.from_pretrained` directly (so they scope to the model + loader only — important for mixed-layout repos where ONNX is + under `onnx/` but the tokenizer/config are at the repo root). + Use this for e.g. `file_name`, `subfolder`, `revision`, + `cache_dir`, `provider`, `provider_options`, `session_options`. + :raises ValueError: If `backend` is not one of "torch" or "ort". + :raises ImportError: If required libraries for the chosen backend + are not installed. """ + if backend not in ("torch", "ort"): + raise ValueError( + f"Unsupported backend: {backend!r}. Expected 'torch' or 'ort'." + ) + self.backend = backend + # Early check for required dependencies if hf_pipeline is None: raise ImportError( "transformers is not installed. Please install it " - "(pip install transformers torch) to use this recognizer." - ) - if torch is None: - raise ImportError( - "torch is not installed. Please install it " - "(pip install torch) to use this recognizer." + "(pip install transformers) to use this recognizer." ) + if self.backend == "torch": + if torch is None: + raise ImportError( + "torch is not installed. Please install it " + "(pip install torch) to use the 'torch' backend." + ) + elif self.backend == "ort": + if optimum_pipeline is None: + raise ImportError( + "optimum is not installed. Please install it " + "(pip install optimum 'optimum-onnx[onnxruntime]') " + "to use the 'ort' backend." + ) self.model_name = model_name self.tokenizer_name = tokenizer_name or model_name @@ -173,16 +212,16 @@ def __init__( "aggregation_strategy='none' may result in fragmented entities " "(e.g., 'B-PER', 'I-PER'). Recommended: 'simple' or 'first'." ) + if self.backend == "ort" and device is not None: + logger.warning( + "The 'device' parameter is ignored by the 'ort' backend. " + "Select hardware via the 'provider' model kwarg instead, " + "e.g. provider='CUDAExecutionProvider'." + ) self.device = self._parse_device(device) self.label_prefixes = label_prefixes or ["B-", "I-", "U-", "L-"] self.ner_pipeline = None - - if kwargs: - logger.warning( - "Ignoring unsupported kwargs in %s: %s", - name, - sorted(kwargs.keys()), - ) + self.model_kwargs = model_kwargs # Derive supported entities from label mapping if supported_entities: @@ -247,8 +286,10 @@ def load(self) -> None: """Load the HuggingFace NER pipeline. This method handles: - 1. Hardware acceleration setup (CUDA validation and fallback) - 2. Lazy-loading of the heavyweight ML pipeline. + 1. Backend selection (torch or ort) + 2. Hardware acceleration setup (CUDA validation and fallback for + the torch backend; provider selection for ort via model_kwargs) + 3. Lazy-loading of the heavyweight ML pipeline. :raises ValueError: If model_name is not set """ @@ -261,6 +302,13 @@ def load(self) -> None: "Pass it to __init__() or set it directly." ) + if self.backend == "torch": + self._load_torch_pipeline() + else: + self._load_ort_pipeline() + + def _load_torch_pipeline(self) -> None: + """Load the NER pipeline using PyTorch backend.""" # Device validation and fallback device = self.device if device >= 0: @@ -275,7 +323,10 @@ def load(self) -> None: ) device = -1 - logger.info(f"Loading HuggingFace model: {self.model_name}, device={device}") + logger.info( + f"Loading HuggingFace model: {self.model_name}, " + f"backend=torch, device={device}" + ) try: self.ner_pipeline = hf_pipeline( @@ -284,12 +335,51 @@ def load(self) -> None: tokenizer=self.tokenizer_name, aggregation_strategy=self.aggregation_strategy, device=device, + model_kwargs=self.model_kwargs or None, ) logger.info(f"Successfully loaded {self.model_name}") except Exception: logger.exception(f"Failed to load model {self.model_name}") raise + def _load_ort_pipeline(self) -> None: + """Load the NER pipeline using optimum's ONNX Runtime backend. + + Pre-loads ``ORTModelForTokenClassification`` explicitly so that + model_kwargs like ``subfolder`` and ``file_name`` are scoped to the + model loader only. Passing them at the pipeline level leaks them + into transformers' config/tokenizer loading, which breaks + mixed-layout repos (e.g. onnx-community/*, Xenova/*) where the ONNX + file lives under ``onnx/`` but config/tokenizer live at the repo + root. + """ + try: + from optimum.onnxruntime import ORTModelForTokenClassification + except ImportError as e: + raise ImportError( + "optimum-onnx is not installed. Please install it " + "(pip install 'optimum-onnx[onnxruntime]') " + "to use the 'ort' backend." + ) from e + + logger.info(f"Loading HuggingFace model: {self.model_name}, backend=ort") + + try: + model = ORTModelForTokenClassification.from_pretrained( + self.model_name, **self.model_kwargs + ) + self.ner_pipeline = optimum_pipeline( + self.DEFAULT_HF_TASK, + model=model, + tokenizer=self.tokenizer_name, + aggregation_strategy=self.aggregation_strategy, + accelerator="ort", + ) + logger.info(f"Successfully loaded {self.model_name} with ort backend") + except Exception: + logger.exception(f"Failed to load model {self.model_name} with ort backend") + raise + def _normalize_label(self, label: str) -> str: """Normalize label by removing prefixes like B-/I-/U-/L-. diff --git a/presidio-analyzer/pyproject.toml b/presidio-analyzer/pyproject.toml index 8f947c7ace..6256bfcf15 100644 --- a/presidio-analyzer/pyproject.toml +++ b/presidio-analyzer/pyproject.toml @@ -69,6 +69,11 @@ langextract = [ "more-itertools (>=10.0.0,<12.0.0)", "jinja2 (>=3.0.0,<4.0.0)", ] +onnxruntime = [ + "optimum (>=2.1.0,<3.0.0)", + "optimum-onnx[onnxruntime] (>=0.1.0,<2.0.0)", + "transformers (>=4.0.0,<5.0.0)", +] [tool.poetry.group.dev.dependencies] pip = "*" diff --git a/presidio-analyzer/tests/test_huggingface_ner_recognizer.py b/presidio-analyzer/tests/test_huggingface_ner_recognizer.py index 7191aa6a43..ef8498fd66 100644 --- a/presidio-analyzer/tests/test_huggingface_ner_recognizer.py +++ b/presidio-analyzer/tests/test_huggingface_ner_recognizer.py @@ -277,6 +277,7 @@ def test_load_invokes_hf_pipeline_with_expected_args(): tokenizer="test-model", aggregation_strategy="simple", device=-1, + model_kwargs=None, ) @@ -472,14 +473,28 @@ def test_hf_recognizer_device_fallback_and_validation(): @pytest.mark.usefixtures("mock_torch_installed") -@patch(HF_PIPELINE_PATH, new=MagicMock()) -def test_hf_recognizer_init_logs_warning_for_extra_kwargs(caplog): - """Test that valid but unsupported kwargs trigger a warning.""" - caplog.set_level(logging.WARNING, logger="presidio-analyzer") - # Passed 'unsupported_arg' which is not in __init__ - HuggingFaceNerRecognizer(model_name="test-model", unsupported_arg="some_value") +def test_hf_recognizer_forwards_extra_kwargs_as_model_kwargs(): + """Test that extra kwargs are forwarded to the pipeline via model_kwargs.""" + with patch(HF_PIPELINE_PATH, new=MagicMock()) as mock_hf_pipeline: + rec = HuggingFaceNerRecognizer( + model_name="test-model", + device=-1, + revision="main", + cache_dir="/tmp/cache", + ) - assert "Ignoring unsupported kwargs" in caplog.text + assert rec.model_kwargs == { + "revision": "main", + "cache_dir": "/tmp/cache", + } + mock_hf_pipeline.assert_called_once_with( + "token-classification", + model="test-model", + tokenizer="test-model", + aggregation_strategy="simple", + device=-1, + model_kwargs={"revision": "main", "cache_dir": "/tmp/cache"}, + ) @pytest.mark.usefixtures("mock_torch_installed") @@ -602,6 +617,158 @@ def test_hf_recognizer_analyze_handles_malformed_pipeline_output( assert rec.analyze("test", entities=["PERSON"]) == [] +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_ort_backend_requires_optimum(): + """ort backend raises ImportError when optimum is not installed.""" + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.optimum_pipeline", + None, + ): + with pytest.raises(ImportError, match="optimum is not installed"): + HuggingFaceNerRecognizer(model_name="test-model", backend="ort") + + +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_unknown_backend_raises(): + """Backends other than 'torch' or 'ort' are rejected at construction.""" + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with pytest.raises(ValueError, match="Unsupported backend"): + HuggingFaceNerRecognizer(model_name="test-model", backend="ov") + + +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_ort_backend_warns_when_device_set(caplog): + """Explicit device with ort backend logs a warning (device is ignored).""" + caplog.set_level(logging.WARNING, logger="presidio-analyzer") + mock_optimum = MagicMock() + mock_ort_model_cls = MagicMock() + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.optimum_pipeline", + mock_optimum, + ): + with patch( + "optimum.onnxruntime.ORTModelForTokenClassification", + mock_ort_model_cls, + ): + HuggingFaceNerRecognizer( + model_name="test-model", backend="ort", device="cuda" + ) + + assert "ignored by the 'ort' backend" in caplog.text + + +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_ort_backend_loads_optimum_pipeline(): + """ort backend pre-loads ORTModel, then hands it to optimum_pipeline.""" + mock_optimum = MagicMock() + mock_ort_model_cls = MagicMock() + mock_model_instance = MagicMock() + mock_ort_model_cls.from_pretrained.return_value = mock_model_instance + + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.optimum_pipeline", + mock_optimum, + ): + with patch( + "optimum.onnxruntime.ORTModelForTokenClassification", + mock_ort_model_cls, + ): + HuggingFaceNerRecognizer(model_name="test-model", backend="ort") + + mock_ort_model_cls.from_pretrained.assert_called_once_with("test-model") + mock_optimum.assert_called_once_with( + "token-classification", + model=mock_model_instance, + tokenizer="test-model", + aggregation_strategy="simple", + accelerator="ort", + ) + + +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_optimum_model_kwargs_scoped_to_model_loader(): + """model_kwargs flow only to ORTModel.from_pretrained, not the pipeline. + + Pipeline-level kwargs would leak into transformers' tokenizer/config + loading and break mixed-layout repos (e.g. onnx-community/* where the + ONNX file is in ``onnx/`` but the tokenizer is at the repo root). + """ + mock_optimum = MagicMock() + mock_ort_model_cls = MagicMock() + mock_model_instance = MagicMock() + mock_ort_model_cls.from_pretrained.return_value = mock_model_instance + + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.optimum_pipeline", + mock_optimum, + ): + with patch( + "optimum.onnxruntime.ORTModelForTokenClassification", + mock_ort_model_cls, + ): + HuggingFaceNerRecognizer( + model_name="test-model", + backend="ort", + subfolder="onnx", + file_name="model_fp16.onnx", + ) + + mock_ort_model_cls.from_pretrained.assert_called_once_with( + "test-model", subfolder="onnx", file_name="model_fp16.onnx" + ) + _, pipeline_kwargs = mock_optimum.call_args + assert "subfolder" not in pipeline_kwargs + assert "file_name" not in pipeline_kwargs + assert "model_kwargs" not in pipeline_kwargs + + +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_torch_backend_no_torch_raises(): + """Test that torch backend raises ImportError when torch is missing.""" + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.torch", + None, + ): + with pytest.raises(ImportError, match="torch is not installed"): + HuggingFaceNerRecognizer(model_name="test-model", backend="torch") + + +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_ort_backend_no_torch_ok(): + """Test that ort backend works without torch installed.""" + mock_optimum = MagicMock() + mock_ort_model_cls = MagicMock() + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.torch", + None, + ): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.optimum_pipeline", + mock_optimum, + ): + with patch( + "optimum.onnxruntime.ORTModelForTokenClassification", + mock_ort_model_cls, + ): + rec = HuggingFaceNerRecognizer( + model_name="test-model", backend="ort" + ) + assert rec.backend == "ort" + + def test_hf_recognizer_loader_supported_entities_filtering(): """Verify if supported_entities survives the RecognizerListLoader logic.""" rec_conf = { diff --git a/presidio-analyzer/tests/test_huggingface_ner_recognizer_e2e.py b/presidio-analyzer/tests/test_huggingface_ner_recognizer_e2e.py new file mode 100644 index 0000000000..97bdf9cfa9 --- /dev/null +++ b/presidio-analyzer/tests/test_huggingface_ner_recognizer_e2e.py @@ -0,0 +1,165 @@ +"""End-to-end tests for HuggingFaceNerRecognizer with real model loading. + +Unlike test_huggingface_ner_recognizer.py (fully mocked), these tests +download real models from the HuggingFace Hub and exercise the actual +transformers/optimum pipelines. They are the only coverage for the +recognizer's load() path against real library behavior. + +Two tiers: + +1. Mechanics tests with hf-internal-testing/tiny-random-bert (~100KB). + The weights are random, so they assert on mechanics (spans, types, + thresholds, torch/ort parity), not on meaningful predictions. +2. Semantic tests with the Stanford de-identifier — the model family used + by conf/onnx.yaml and already downloaded by the engine tests (conftest + references the torch variant). These assert on actual PII detection, + and the ort test covers the mixed-layout repo scenario (ONNX under + onnx/, tokenizer at root) that requires subfolder/file_name scoping. +""" + +import pytest +from presidio_analyzer import RecognizerResult +from presidio_analyzer.predefined_recognizers import ( + HuggingFaceNerRecognizer, +) + +TINY_MODEL = "hf-internal-testing/tiny-random-bert" +TINY_TEXT = "John Smith works at Contoso in Berlin since 2019." +# Random-weight model emits LABEL_0/LABEL_1; map both so output is non-empty. +TINY_LABEL_MAPPING = {"LABEL_0": "PERSON", "LABEL_1": "LOCATION"} + + +def _assert_valid_results(results, text): + assert isinstance(results, list) + assert len(results) > 0 + for r in results: + assert isinstance(r, RecognizerResult) + assert r.entity_type in ("PERSON", "LOCATION") + assert 0 <= r.start < r.end <= len(text) + assert 0.0 <= r.score <= 1.0 + + +def test_hf_recognizer_e2e_torch_backend(): + """Real torch pipeline end-to-end on a tiny random model.""" + pytest.importorskip("torch", reason="torch is not installed") + + rec = HuggingFaceNerRecognizer( + model_name=TINY_MODEL, + backend="torch", + device="cpu", + label_mapping=TINY_LABEL_MAPPING, + threshold=0.0, + ) + results = rec.analyze(TINY_TEXT, entities=["PERSON", "LOCATION"]) + _assert_valid_results(results, TINY_TEXT) + + +def test_hf_recognizer_e2e_ort_backend(): + """Real ONNX Runtime pipeline end-to-end on a tiny random model. + + The repo has no .onnx weights, so this also exercises optimum's + export-on-the-fly path (export=True forwarded via **model_kwargs). + """ + pytest.importorskip( + "optimum.onnxruntime", reason="optimum-onnx is not installed" + ) + + rec = HuggingFaceNerRecognizer( + model_name=TINY_MODEL, + backend="ort", + label_mapping=TINY_LABEL_MAPPING, + threshold=0.0, + export=True, + ) + results = rec.analyze(TINY_TEXT, entities=["PERSON", "LOCATION"]) + _assert_valid_results(results, TINY_TEXT) + + +def test_hf_recognizer_e2e_torch_and_ort_agree_on_spans(): + """Both backends run the same model; spans and scores should match.""" + pytest.importorskip("torch", reason="torch is not installed") + pytest.importorskip( + "optimum.onnxruntime", reason="optimum-onnx is not installed" + ) + + common = dict( + model_name=TINY_MODEL, + label_mapping=TINY_LABEL_MAPPING, + threshold=0.0, + ) + torch_results = HuggingFaceNerRecognizer( + backend="torch", device="cpu", **common + ).analyze(TINY_TEXT, entities=["PERSON", "LOCATION"]) + ort_results = HuggingFaceNerRecognizer( + backend="ort", export=True, **common + ).analyze(TINY_TEXT, entities=["PERSON", "LOCATION"]) + + torch_spans = [(r.entity_type, r.start, r.end) for r in torch_results] + ort_spans = [(r.entity_type, r.start, r.end) for r in ort_results] + assert torch_spans == ort_spans + for tr, orr in zip(torch_results, ort_results): + assert abs(tr.score - orr.score) < 1e-3 + + +DEID_TEXT = ( + "Hi, I'm Dr. Sarah Chen from Mount Sinai Hospital. " + "Please call me at 555-123-4567 on March 5th, 2026." +) +DEID_LABEL_MAPPING = { + "PATIENT": "PERSON", + "HCW": "PERSON", + "HOSPITAL": "ORGANIZATION", + "VENDOR": "ORGANIZATION", + "DATE": "DATE_TIME", + "PHONE": "PHONE_NUMBER", + "ID": "ID", +} +DEID_ENTITIES = ["PERSON", "ORGANIZATION", "DATE_TIME", "PHONE_NUMBER"] + + +def _assert_deid_detections(results, text): + detected = {(r.entity_type, text[r.start : r.end]) for r in results} + assert ("PERSON", "Dr. Sarah Chen") in detected + assert ("ORGANIZATION", "Mount Sinai Hospital") in detected + assert ("PHONE_NUMBER", "555-123-4567") in detected + + +def test_hf_recognizer_e2e_torch_stanford_deidentifier(): + """Real PII detection with the torch backend on the Stanford model.""" + pytest.importorskip("torch", reason="torch is not installed") + + rec = HuggingFaceNerRecognizer( + model_name="StanfordAIMI/stanford-deidentifier-base", + backend="torch", + device="cpu", + label_mapping=DEID_LABEL_MAPPING, + threshold=0.5, + ) + results = rec.analyze(DEID_TEXT, entities=DEID_ENTITIES) + _assert_deid_detections(results, DEID_TEXT) + + +def test_hf_recognizer_e2e_ort_mixed_layout_repo(): + """Real PII detection with ort on a mixed-layout repo. + + onnx-community/stanford-deidentifier-base-ONNX keeps config/tokenizer at + the repo root and ONNX files under onnx/. This is the scenario that + requires subfolder/file_name to be scoped to the model loader only — + regression coverage for the pipeline-level kwarg leak. + """ + pytest.importorskip( + "optimum.onnxruntime", reason="optimum-onnx is not installed" + ) + + rec = HuggingFaceNerRecognizer( + model_name="onnx-community/stanford-deidentifier-base-ONNX", + backend="ort", + subfolder="onnx", + # INT8 variant: same detections as model.onnx on this text, but a + # 105MB download instead of 416MB (matters in CI, no HF cache there). + file_name="model_quantized.onnx", + label_mapping=DEID_LABEL_MAPPING, + threshold=0.5, + ) + results = rec.analyze(DEID_TEXT, entities=DEID_ENTITIES) + _assert_deid_detections(results, DEID_TEXT) From b2c7847e398c1416a99a4a1e090cd817d6cdb668 Mon Sep 17 00:00:00 2001 From: Yurii Havrylko <10372700+yuriihavrylko@users.noreply.github.com> Date: Fri, 12 Jun 2026 21:59:27 +0200 Subject: [PATCH 2/6] fix: exclude None values in recognizer registry configuration validation --- presidio-analyzer/presidio_analyzer/input_validation/schemas.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/presidio-analyzer/presidio_analyzer/input_validation/schemas.py b/presidio-analyzer/presidio_analyzer/input_validation/schemas.py index a256ca2a58..0ac9ae10b5 100644 --- a/presidio-analyzer/presidio_analyzer/input_validation/schemas.py +++ b/presidio-analyzer/presidio_analyzer/input_validation/schemas.py @@ -81,7 +81,7 @@ def validate_recognizer_registry_configuration( # Use Pydantic model for validation validated_config = RecognizerRegistryConfig(**config) # Use model_dump() without exclude_unset to include default values - return validated_config.model_dump(exclude_unset=False) + return validated_config.model_dump(exclude_unset=False, exclude_none=True) except ValidationError as e: raise ValueError("Invalid recognizer registry configuration") from e From bc4f45f8d49791f9823d62268d1da16cd1b7c3e7 Mon Sep 17 00:00:00 2001 From: Yurii Havrylko <10372700+yuriihavrylko@users.noreply.github.com> Date: Fri, 12 Jun 2026 23:02:53 +0200 Subject: [PATCH 3/6] docs: add HuggingFace NER inference backends documentation and update navigation --- docs/analyzer/huggingface_ner_inference.md | 335 ++++++++++++++++++ docs/analyzer/nlp_engines/gpu_usage.md | 1 + docs/analyzer/recognizer_registry_provider.md | 13 + mkdocs.yml | 1 + 4 files changed, 350 insertions(+) create mode 100644 docs/analyzer/huggingface_ner_inference.md diff --git a/docs/analyzer/huggingface_ner_inference.md b/docs/analyzer/huggingface_ner_inference.md new file mode 100644 index 0000000000..fffef66f18 --- /dev/null +++ b/docs/analyzer/huggingface_ner_inference.md @@ -0,0 +1,335 @@ +# HuggingFace NER inference backends + +The `HuggingFaceNerRecognizer` runs a HuggingFace token-classification model +to detect PII. It supports two inference backends, selected by the `backend` +constructor parameter — set it in the recognizer YAML or pass it directly +when constructing the recognizer in Python: + +- `torch` (default) — PyTorch via the `transformers` pipeline. Works on CPU + and NVIDIA GPU (CUDA). +- `ort` — ONNX Runtime via `optimum`. Works on CPU, NVIDIA GPU, AMD GPU, + Intel CPU/iGPU/NPU, and Apple Silicon, by selecting an ONNX Runtime + *execution provider*. + +Any extra keyword arguments (e.g. `file_name`, `subfolder`, `provider`, +`provider_options`, `revision`) — whether passed in Python or as extra +fields in the YAML — are captured into `**model_kwargs` and forwarded to +the underlying loader. You do not need to edit the recognizer to plumb +new knobs. + +The examples below use YAML (the +[recognizer registry configuration](recognizer_registry_provider.md)); +every field shown maps 1:1 to a constructor argument. The Python +equivalent of an ort configuration: + +```python +from presidio_analyzer import AnalyzerEngine +from presidio_analyzer.predefined_recognizers import HuggingFaceNerRecognizer + +recognizer = HuggingFaceNerRecognizer( + model_name="onnx-community/stanford-deidentifier-base-ONNX", + backend="ort", + subfolder="onnx", # forwarded via **model_kwargs + file_name="model_quantized.onnx", # forwarded via **model_kwargs + label_mapping={ + "PATIENT": "PERSON", + "HCW": "PERSON", + "PHONE": "PHONE_NUMBER", + }, + threshold=0.5, +) + +analyzer = AnalyzerEngine() +analyzer.registry.add_recognizer(recognizer) +results = analyzer.analyze(text="Dr. Sarah Chen, 555-123-4567", language="en") +``` + +## Choosing a backend + +| Hardware | Recommended backend | Why | +|----------|---------------------|-----| +| NVIDIA GPU, baseline | `torch` with `device: cuda` | Zero extra install, no model conversion | +| NVIDIA GPU + quantized ONNX | `ort` with `CUDAExecutionProvider` | Runs pre-quantized FP16/INT8 ONNX files that torch can't use directly | +| NVIDIA GPU + high throughput | `ort` with `TensorrtExecutionProvider` | Designed for sustained batch inference; slow first request (engine build) | +| Intel Xeon Scalable (Sapphire Rapids+) | `ort` with `OpenVINOExecutionProvider` | Uses native AVX-512 BF16 / AMX kernels not fully exploited by the default CPU EP | +| Intel iGPU / Arc / NPU | `ort` with `OpenVINOExecutionProvider` (`device_type=GPU` or `NPU`) | Only practical path for Intel accelerators | +| Apple Silicon | `ort` with `CoreMLExecutionProvider`, or CPU | The recognizer's device handling does not support MPS; CoreML EP is the only acceleration path | +| Generic cloud CPU | `torch` with `device: cpu`, or `ort` with default CPU EP | Comparable for FP32; ort additionally runs quantized variants | + +## torch backend + +The default. Uses `transformers.pipeline` directly. Best for "just works" +deployments. + +```yaml +- name: "HF NER (torch)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: dslim/bert-base-NER + backend: torch + # device: cuda # omit to let Presidio's device detector decide + aggregation_strategy: simple + threshold: 0.3 + supported_languages: [en] + supported_entities: [PERSON, LOCATION, ORGANIZATION] +``` + +Device selection precedence: + +1. **Explicit `device`** (`cpu` | `cuda` | `cuda:N` | int index) — hard + override; bypasses the detector *and* `PRESIDIO_DEVICE`. +2. **`device` omitted** — Presidio's `DeviceDetector` singleton decides: + the `PRESIDIO_DEVICE` environment variable if set, otherwise + auto-detection (CUDA if available, else CPU). See + [GPU Acceleration](nlp_engines/gpu_usage.md). + +In most deployments, omit `device` and control hardware via +`PRESIDIO_DEVICE` — that keeps the YAML portable across CPU and GPU +environments. Set `device` explicitly only to pin a recognizer to +specific hardware regardless of environment (e.g. force a heavyweight +model onto `cuda:1`). + +If CUDA is requested but unavailable, the recognizer logs a warning and +falls back to CPU. + +!!! note "MPS (Apple Silicon) is not supported" + Presidio's device detector does not support MPS, and the recognizer + normalizes any non-CUDA device to CPU. On Apple Silicon the torch + backend runs on CPU; for acceleration use the ort backend with the + CoreML execution provider. + +## ort backend + +Uses optimum's `ORTModelForTokenClassification` plus ONNX Runtime. + +### Installation + +For CPU (also covers Apple Silicon via the CoreML EP), use the +`onnxruntime` extra: + +```bash +pip install 'presidio-analyzer[onnxruntime]' +``` + +The extra installs the **CPU** build of ONNX Runtime. For GPU or other +accelerators, skip the extra and install the optimum stack with the +ONNX Runtime build matching your hardware — e.g. in a CUDA Docker image: + +```bash +# NVIDIA GPU: onnxruntime-gpu instead of onnxruntime +pip install presidio-analyzer optimum 'optimum-onnx[onnxruntime-gpu]' 'transformers<5' +``` + +| Hardware | ONNX Runtime build | +|----------|--------------------| +| CPU (default) | `onnxruntime` — via `presidio-analyzer[onnxruntime]` | +| NVIDIA GPU | `optimum-onnx[onnxruntime-gpu]` (pulls `onnxruntime-gpu`) | +| Apple Silicon (CoreML) | `onnxruntime` — CoreML EP ships in the default package | +| Intel CPU/GPU/NPU | `onnxruntime-openvino` (install instead of `onnxruntime`) | +| AMD GPU (ROCm) | `onnxruntime-rocm` (install instead of `onnxruntime`) | + +!!! warning "Do not combine the extra with a GPU build" + `onnxruntime`, `onnxruntime-gpu`, `onnxruntime-openvino`, etc. all + ship the same Python module and shadow each other when several are + installed. If you use a non-CPU build, install it *instead of* the + `[onnxruntime]` extra, not on top of it. To recover from a mixed + install: + + ```bash + pip uninstall onnxruntime onnxruntime-gpu onnxruntime-openvino + pip install onnxruntime-gpu # the one you actually want + ``` + +### YAML — CPU (default ORT) + +```yaml +- name: "HF NER (ORT CPU)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: optimum/bert-base-NER + backend: ort + aggregation_strategy: simple + threshold: 0.3 + supported_languages: [en] + supported_entities: [PERSON, LOCATION, ORGANIZATION] +``` + +For models published under the dual-runtime convention +(`onnx-community/*`, `Xenova/*`) where ONNX files +sit under `onnx/` while config/tokenizer live at the repo root, also +specify `subfolder` and `file_name`: + +```yaml +- name: "HF NER (ORT, FP16 ONNX)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: onnx-community/stanford-deidentifier-base-ONNX + backend: ort + subfolder: onnx + file_name: model_fp16.onnx # or model.onnx, model_quantized.onnx, model_int8.onnx, ... + aggregation_strategy: simple + threshold: 0.3 +``` + +The recognizer pre-loads the ORT model with these kwargs scoped to the +model loader, so they don't leak into transformers' tokenizer/config +loading. + +### YAML — NVIDIA GPU + +Install `onnxruntime-gpu` and select the CUDA provider: + +```yaml +- name: "HF NER (ORT, CUDA)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: optimum/bert-base-NER + backend: ort + provider: CUDAExecutionProvider + provider_options: + - device_id: 0 + aggregation_strategy: simple + threshold: 0.3 +``` + +For TensorRT (higher throughput, slow first request because of engine +build): + +```yaml +- name: "HF NER (ORT, TensorRT)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: optimum/bert-base-NER + backend: ort + provider: TensorrtExecutionProvider + provider_options: + - trt_fp16_enable: true + trt_engine_cache_enable: true + trt_engine_cache_path: /var/cache/trt + aggregation_strategy: simple + threshold: 0.3 +``` + +### YAML — Intel CPU / iGPU / NPU via OpenVINO + +Install `onnxruntime-openvino`: + +```bash +pip uninstall onnxruntime onnxruntime-gpu +pip install onnxruntime-openvino +``` + +Then route through the OpenVINO execution provider: + +```yaml +- name: "HF NER (ORT, OpenVINO)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: optimum/bert-base-NER + backend: ort + provider: OpenVINOExecutionProvider + provider_options: + - device_type: CPU # CPU | GPU | NPU | AUTO + aggregation_strategy: simple + threshold: 0.3 +``` + +OpenVINO is most worthwhile on: + +- Intel Xeon Scalable (Sapphire Rapids and newer) — its kernels use + AVX-512 BF16 / AMX instructions that the default CPU EP does not fully + exploit. Intel publishes substantial speedups for BERT-class models on + this hardware; actual gains depend heavily on model, sequence length, + and batch size. +- Intel Arc / iGPU / Data Center GPU Max — native GPU acceleration without + CUDA. +- Intel NPUs (Meteor Lake, Lunar Lake, Arrow Lake) — currently the only + practical runtime for on-device NPU inference. + +!!! note "Benchmark before adopting" + The throughput figures published for these providers come from vendor + benchmarks under favorable conditions (large batches, long + sequences). Typical Presidio traffic — short texts at low batch + sizes — usually sees smaller gains, and on generic cloud VMs (where + the CPU generation isn't guaranteed) the difference over the default + CPU EP can be negligible. Measure on your target hardware with your + own traffic shape before committing to a provider. + +### Apple Silicon (CoreML) + +The default `onnxruntime` package ships the CoreML execution provider on +macOS: + +```yaml +- name: "HF NER (ORT, CoreML)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: optimum/bert-base-NER + backend: ort + provider: CoreMLExecutionProvider +``` + +On Apple Silicon the torch backend runs on CPU (the recognizer does not +support MPS), so the CoreML EP is the only acceleration path — at the +cost of needing an ONNX model (pre-built or exported). + +### AMD GPU (ROCm) + +```yaml +- name: "HF NER (ORT, ROCm)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: optimum/bert-base-NER + backend: ort + provider: ROCMExecutionProvider +``` + +Requires `onnxruntime-rocm`. + +## Picking an ONNX variant + +When the repo ships multiple quantized ONNX files (common in +`onnx-community/*`), the precision/quantization knob is the `file_name`, +not a runtime flag: + +| File | Precision | Notes | +|------|-----------|-------| +| `model.onnx` | FP32 | Baseline; largest | +| `model_fp16.onnx` | FP16 | ~half the size; benefits from NVIDIA tensor cores, limited CPU support | +| `model_bf16.onnx` | BF16 | Needs BF16-capable hardware (Sapphire Rapids+, recent NVIDIA) | +| `model_int8.onnx` | INT8 (dynamic) | ~quarter the size; aimed at CPU inference | +| `model_quantized.onnx` | INT8 (static) | Pre-calibrated variant of INT8 | +| `model_uint8.onnx` | UINT8 | Similar to INT8 | +| `model_q4.onnx`, `model_q4f16.onnx` | 4-bit | Smallest; expect an accuracy trade-off | + +Smaller is not automatically faster — quantized kernels, hardware +support, and accuracy interact. Validate both latency *and* detection +quality (precision/recall on your own data) when switching variants. + +Pick by setting `file_name:` in the YAML. Inference precision otherwise +follows what's baked into the ONNX file — there is no separate `dtype` +parameter at the recognizer level. + +## Troubleshooting + +- **`FileNotFoundError: Could not find any ONNX files`** — the repo uses a + mixed layout. Set `subfolder: onnx` and a *bare* `file_name` + (e.g. `model.onnx`, not `onnx/model.onnx`). Path-in-`file_name` does not + match optimum's resolver even when the file is listed in the error. +- **`OSError: Could not locate onnx/config.json`** — `subfolder` is + leaking into tokenizer/config loading. Confirm you are on a recognizer + version that pre-loads `ORTModelForTokenClassification` explicitly; on + older versions the `subfolder` is passed at the pipeline level and + breaks mixed-layout repos. +- **`ImportError: cannot import name 'is_offline_mode' from + 'transformers.utils'`** — `optimum-onnx` is incompatible with + `transformers>=5`. Pin `transformers<5`. +- **Provider not available** at runtime — usually means another ORT + package (e.g. `onnxruntime` next to `onnxruntime-gpu`) is shadowing the + one with the provider you want. Uninstall all and reinstall a single + variant. +- **First TensorRT request is slow** — TRT builds an engine on first use. + Persist `trt_engine_cache_path` so subsequent processes reuse it. +- **`TypeError: got an unexpected keyword argument 'foo'`** — a key in + your YAML is being forwarded to the loader and isn't recognized. Stray + kwargs are no longer silently dropped; the loud failure is intentional. diff --git a/docs/analyzer/nlp_engines/gpu_usage.md b/docs/analyzer/nlp_engines/gpu_usage.md index dc11f3b2b5..d7e112acc6 100644 --- a/docs/analyzer/nlp_engines/gpu_usage.md +++ b/docs/analyzer/nlp_engines/gpu_usage.md @@ -145,3 +145,4 @@ Related Presidio documentation: - [Transformer Models](transformers.md) - [spaCy and Stanza Models](spacy_stanza.md) - [Customizing NLP Models](../customizing_nlp_models.md) +- [HuggingFace NER inference backends](../huggingface_ner_inference.md) — ORT execution providers (CUDA, TensorRT, OpenVINO, CoreML, ROCm) and Intel-specific acceleration via `HuggingFaceNerRecognizer` diff --git a/docs/analyzer/recognizer_registry_provider.md b/docs/analyzer/recognizer_registry_provider.md index 4fd39e5c3a..c8ed0d51a2 100644 --- a/docs/analyzer/recognizer_registry_provider.md +++ b/docs/analyzer/recognizer_registry_provider.md @@ -89,6 +89,7 @@ The recognizer list comprises of both the predefined and custom recognizers, for type: "predefined" class_name: "HuggingFaceNerRecognizer" model_name: "dslim/bert-base-NER" + backend: "torch" # or "ort" for ONNX Runtime supported_languages: - en supported_entities: ["PERSON", "LOCATION", "ORGANIZATION"] @@ -96,6 +97,18 @@ The recognizer list comprises of both the predefined and custom recognizers, for device: "cpu" ``` +!!! tip "Optimized inference (GPU, CPU, Intel)" + + `HuggingFaceNerRecognizer` supports a `torch` backend (default) and an + `ort` backend (ONNX Runtime via `optimum`). The `ort` backend can + target NVIDIA GPU (`CUDAExecutionProvider`, + `TensorrtExecutionProvider`), Intel CPU/iGPU/NPU + (`OpenVINOExecutionProvider`), AMD GPU (`ROCMExecutionProvider`), and + Apple Silicon (`CoreMLExecutionProvider`) via execution providers. + See [HuggingFace NER inference backends](huggingface_ner_inference.md) + for install instructions, YAML examples per accelerator, and guidance + on picking quantized ONNX variants. + ### The recognizer parameters - `supported_languages`: A list of supported languages that the analyzer will support. In case this field is missing, a recognizer will be created for each supported language provided to the `AnalyzerEngine`. diff --git a/mkdocs.yml b/mkdocs.yml index e6967be33f..10c72ef620 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -43,6 +43,7 @@ nav: - Tutorial: analyzer/adding_recognizers.md - Best practices: analyzer/developing_recognizers.md - Recognizer registry from file: analyzer/recognizer_registry_provider.md + - HuggingFace NER inference backends: analyzer/huggingface_ner_inference.md - Filtering recognizers by country: analyzer/filtering_by_country.md - Multi-language support: analyzer/languages.md - Customizing the NLP model: From 2f0d172ac8a08f01dced61ea1c2f69437d41b4d9 Mon Sep 17 00:00:00 2001 From: Yurii Havrylko <10372700+yuriihavrylko@users.noreply.github.com> Date: Fri, 12 Jun 2026 23:05:07 +0200 Subject: [PATCH 4/6] feat: add configuration for HuggingFace NER recognizer with ONNX support --- .../presidio_analyzer/conf/hf_ner_onnx.yaml | 41 +++++++++++++++++++ .../test_huggingface_ner_recognizer_e2e.py | 2 +- 2 files changed, 42 insertions(+), 1 deletion(-) create mode 100644 presidio-analyzer/presidio_analyzer/conf/hf_ner_onnx.yaml diff --git a/presidio-analyzer/presidio_analyzer/conf/hf_ner_onnx.yaml b/presidio-analyzer/presidio_analyzer/conf/hf_ner_onnx.yaml new file mode 100644 index 0000000000..ac9c984ae3 --- /dev/null +++ b/presidio-analyzer/presidio_analyzer/conf/hf_ner_onnx.yaml @@ -0,0 +1,41 @@ +supported_languages: + - en +default_score_threshold: 0 +nlp_configuration: + nlp_engine_name: slim + models: + - lang_code: en + model_name: en_core_web_sm + +recognizer_registry: + recognizers: + - name: SpacyRecognizer + type: predefined + enabled: false + + - name: "HuggingFace NER (ONNX)" + type: predefined + class_name: HuggingFaceNerRecognizer + model_name: onnx-community/stanford-deidentifier-base-ONNX + backend: ort + subfolder: onnx + file_name: model.onnx + supported_languages: + - en + supported_entities: + - PERSON + - ORGANIZATION + - DATE_TIME + - PHONE_NUMBER + - ID + label_mapping: + PATIENT: PERSON + HCW: PERSON + HOSPITAL: ORGANIZATION + VENDOR: ORGANIZATION + DATE: DATE_TIME + PHONE: PHONE_NUMBER + ID: ID + aggregation_strategy: simple + threshold: 0.3 + device: cpu diff --git a/presidio-analyzer/tests/test_huggingface_ner_recognizer_e2e.py b/presidio-analyzer/tests/test_huggingface_ner_recognizer_e2e.py index 97bdf9cfa9..1a95440adc 100644 --- a/presidio-analyzer/tests/test_huggingface_ner_recognizer_e2e.py +++ b/presidio-analyzer/tests/test_huggingface_ner_recognizer_e2e.py @@ -11,7 +11,7 @@ The weights are random, so they assert on mechanics (spans, types, thresholds, torch/ort parity), not on meaningful predictions. 2. Semantic tests with the Stanford de-identifier — the model family used - by conf/onnx.yaml and already downloaded by the engine tests (conftest + by conf/hf_ner_onnx.yaml and already downloaded by the engine tests (conftest references the torch variant). These assert on actual PII detection, and the ort test covers the mixed-layout repo scenario (ONNX under onnx/, tokenizer at root) that requires subfolder/file_name scoping. From f2a8ff9f1b6c19d7bca4dc51956dfc991b220cb6 Mon Sep 17 00:00:00 2001 From: Yurii Havrylko <10372700+yuriihavrylko@users.noreply.github.com> Date: Fri, 26 Jun 2026 23:19:07 +0200 Subject: [PATCH 5/6] Revert "fix: exclude None values in recognizer registry configuration validation" This reverts commit b2c7847e398c1416a99a4a1e090cd817d6cdb668. --- presidio-analyzer/presidio_analyzer/input_validation/schemas.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/presidio-analyzer/presidio_analyzer/input_validation/schemas.py b/presidio-analyzer/presidio_analyzer/input_validation/schemas.py index 0ac9ae10b5..a256ca2a58 100644 --- a/presidio-analyzer/presidio_analyzer/input_validation/schemas.py +++ b/presidio-analyzer/presidio_analyzer/input_validation/schemas.py @@ -81,7 +81,7 @@ def validate_recognizer_registry_configuration( # Use Pydantic model for validation validated_config = RecognizerRegistryConfig(**config) # Use model_dump() without exclude_unset to include default values - return validated_config.model_dump(exclude_unset=False, exclude_none=True) + return validated_config.model_dump(exclude_unset=False) except ValidationError as e: raise ValueError("Invalid recognizer registry configuration") from e From dcf5a5b9a63166d29ca9dd6025cb895f83ffa59a Mon Sep 17 00:00:00 2001 From: Yurii Havrylko <10372700+yuriihavrylko@users.noreply.github.com> Date: Fri, 26 Jun 2026 23:30:24 +0200 Subject: [PATCH 6/6] fix: remove device parameter handling for ort backend in HuggingFaceNerRecognizer --- .../presidio_analyzer/conf/hf_ner_onnx.yaml | 1 - .../ner/huggingface_ner_recognizer.py | 20 ++++--- .../tests/test_huggingface_ner_recognizer.py | 52 +++++++++++++++++++ 3 files changed, 65 insertions(+), 8 deletions(-) diff --git a/presidio-analyzer/presidio_analyzer/conf/hf_ner_onnx.yaml b/presidio-analyzer/presidio_analyzer/conf/hf_ner_onnx.yaml index ac9c984ae3..c31537147b 100644 --- a/presidio-analyzer/presidio_analyzer/conf/hf_ner_onnx.yaml +++ b/presidio-analyzer/presidio_analyzer/conf/hf_ner_onnx.yaml @@ -38,4 +38,3 @@ recognizer_registry: ID: ID aggregation_strategy: simple threshold: 0.3 - device: cpu diff --git a/presidio-analyzer/presidio_analyzer/predefined_recognizers/ner/huggingface_ner_recognizer.py b/presidio-analyzer/presidio_analyzer/predefined_recognizers/ner/huggingface_ner_recognizer.py index 290c1396ac..c04d8c4b45 100644 --- a/presidio-analyzer/presidio_analyzer/predefined_recognizers/ner/huggingface_ner_recognizer.py +++ b/presidio-analyzer/presidio_analyzer/predefined_recognizers/ner/huggingface_ner_recognizer.py @@ -212,13 +212,19 @@ def __init__( "aggregation_strategy='none' may result in fragmented entities " "(e.g., 'B-PER', 'I-PER'). Recommended: 'simple' or 'first'." ) - if self.backend == "ort" and device is not None: - logger.warning( - "The 'device' parameter is ignored by the 'ort' backend. " - "Select hardware via the 'provider' model kwarg instead, " - "e.g. provider='CUDAExecutionProvider'." - ) - self.device = self._parse_device(device) + if self.backend == "ort": + if device not in (None, "cpu", -1): + logger.warning( + "The 'device' parameter is ignored by the 'ort' backend. " + "Select hardware via the 'provider' model kwarg instead, " + "e.g. provider='CUDAExecutionProvider'." + ) + # ort selects hardware via the execution provider, not device. + # Skip parsing/auto-detection and keep self.device consistent + # with actual behavior. + self.device = -1 + else: + self.device = self._parse_device(device) self.label_prefixes = label_prefixes or ["B-", "I-", "U-", "L-"] self.ner_pipeline = None self.model_kwargs = model_kwargs diff --git a/presidio-analyzer/tests/test_huggingface_ner_recognizer.py b/presidio-analyzer/tests/test_huggingface_ner_recognizer.py index ef8498fd66..2e06018b35 100644 --- a/presidio-analyzer/tests/test_huggingface_ner_recognizer.py +++ b/presidio-analyzer/tests/test_huggingface_ner_recognizer.py @@ -661,6 +661,58 @@ def test_hf_recognizer_ort_backend_warns_when_device_set(caplog): assert "ignored by the 'ort' backend" in caplog.text +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_ort_backend_ignores_invalid_device(caplog): + """ort backend skips device parsing: an invalid device must not raise.""" + caplog.set_level(logging.WARNING, logger="presidio-analyzer") + mock_optimum = MagicMock() + mock_ort_model_cls = MagicMock() + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.optimum_pipeline", + mock_optimum, + ): + with patch( + "optimum.onnxruntime.ORTModelForTokenClassification", + mock_ort_model_cls, + ): + # "not-a-device" would raise ValueError via _parse_device on + # the torch backend; ort skips parsing and forces CPU. + rec = HuggingFaceNerRecognizer( + model_name="test-model", + backend="ort", + device="not-a-device", + ) + + assert rec.device == -1 + assert "ignored by the 'ort' backend" in caplog.text + + +@pytest.mark.usefixtures("mock_torch_installed") +def test_hf_recognizer_ort_backend_cpu_device_no_warning(caplog): + """device='cpu' on ort is consistent with default behavior; no warning.""" + caplog.set_level(logging.WARNING, logger="presidio-analyzer") + mock_optimum = MagicMock() + mock_ort_model_cls = MagicMock() + with patch(HF_PIPELINE_PATH, new=MagicMock()): + with patch( + "presidio_analyzer.predefined_recognizers.ner." + "huggingface_ner_recognizer.optimum_pipeline", + mock_optimum, + ): + with patch( + "optimum.onnxruntime.ORTModelForTokenClassification", + mock_ort_model_cls, + ): + rec = HuggingFaceNerRecognizer( + model_name="test-model", backend="ort", device="cpu" + ) + + assert rec.device == -1 + assert "ignored by the 'ort' backend" not in caplog.text + + @pytest.mark.usefixtures("mock_torch_installed") def test_hf_recognizer_ort_backend_loads_optimum_pipeline(): """ort backend pre-loads ORTModel, then hands it to optimum_pipeline."""