diff --git a/README.md b/README.md index b1200bbdc..20698e516 100644 --- a/README.md +++ b/README.md @@ -261,6 +261,20 @@ override it to use a different endpoint. When using the CLI or Python SDK, set `NSS_INFERENCE_KEY` (and `NSS_INFERENCE_ENDPOINT` only if not using the default) so column classification can run. +Column classification can also run locally with a Hugging Face causal LM: + +```yaml +replace_pii: + globals: + classify: + backend: local_hf + model: HuggingFaceTB/SmolLM3-3B # optional default; can also be a local model path +``` + +The local backend does not use `NSS_INFERENCE_KEY`. It loads the default 3B +model from the Hugging Face cache or downloads it when online; in HF offline +mode, pre-download the model or set `model` to a complete local directory. + ### Local Endpoint To point to a locally hosted LLM, add the variables to `.env.local` (git-ignored, auto-loaded by mise): diff --git a/docs/tutorials/safe-synthesizer-101.ipynb b/docs/tutorials/safe-synthesizer-101.ipynb index 825029248..0a73860be 100644 --- a/docs/tutorials/safe-synthesizer-101.ipynb +++ b/docs/tutorials/safe-synthesizer-101.ipynb @@ -217,7 +217,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": ".venv", "language": "python", "name": "python3" }, @@ -231,7 +231,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.9" + "version": "3.13.12" } }, "nbformat": 4, diff --git a/docs/user-guide/configuration.md b/docs/user-guide/configuration.md index a2b6283c8..63b418029 100644 --- a/docs/user-guide/configuration.md +++ b/docs/user-guide/configuration.md @@ -193,7 +193,9 @@ Key config parameters: | Field | Default | Description | Guidance | |-------|---------|-------------|----------| -| `replace_pii.globals.classify.enable_classify` | `true` | Enable LLM-based PII column classification | When using the CLI, set `NSS_INFERENCE_KEY` (and optionally `NSS_INFERENCE_ENDPOINT`); set to `false` if no LLM endpoint is available | +| `replace_pii.globals.classify.enable_classify` | `true` | Enable LLM-based PII column classification | For the `api` backend, set `NSS_INFERENCE_KEY` (and optionally `NSS_INFERENCE_ENDPOINT`); set to `false` to skip LLM classification | +| `replace_pii.globals.classify.backend` | `api` | Column classification backend: `api` or `local_hf` | Use `api` for an OpenAI-compatible endpoint; use `local_hf` for an in-process Hugging Face model with no inference API key | +| `replace_pii.globals.classify.model` | `null` | Local Hugging Face model ID or local path for `local_hf` | When unset with `local_hf`, defaults to `HuggingFaceTB/SmolLM3-3B`; ignored by `api`, which uses `NSS_INFERENCE_MODEL` | | `replace_pii.globals.classify.entities` | (see default list) | Entity types used for LLM-based column classification. Defaults to 15 types covering names, addresses, phone numbers, emails, SSN, national/tax IDs, and credit/debit cards -- see [PII Replacement](../product-overview/pii_replacement.md) and [`PiiReplacerConfig`][nemo_safe_synthesizer.config.replace_pii.PiiReplacerConfig] | Override to add or remove entity types from classification | | `replace_pii.globals.ner.ner_threshold` | `0.3` | GLiNER confidence threshold for NER detection | Lower to catch more entities (more false positives); raise to reduce false positives | diff --git a/docs/user-guide/docker.md b/docs/user-guide/docker.md index b1868b813..db24a7134 100644 --- a/docs/user-guide/docker.md +++ b/docs/user-guide/docker.md @@ -124,8 +124,8 @@ docker run --gpus all --shm-size=1g \ | Variable | Required | Purpose | |----------|----------|---------| | `HF_TOKEN` | For gated models | Hugging Face token for downloading gated models (Llama, Mistral, etc.). Get one at [hf.co/settings/tokens](https://huggingface.co/settings/tokens) | -| `NSS_INFERENCE_KEY` | For PII classification | API key for `NSS_INFERENCE_ENDPOINT`. Set when using the CLI/SDK for column classification | -| `NSS_INFERENCE_ENDPOINT` | For PII classification | NIM/OpenAI-compatible endpoint URL (default: `https://integrate.api.nvidia.com/v1`). Override for a custom endpoint | +| `NSS_INFERENCE_KEY` | For API PII classification | API key for `NSS_INFERENCE_ENDPOINT`. Not used by `replace_pii.globals.classify.backend: local_hf` | +| `NSS_INFERENCE_ENDPOINT` | For API PII classification | NIM/OpenAI-compatible endpoint URL (default: `https://integrate.api.nvidia.com/v1`). Override for a custom endpoint | | `WANDB_API_KEY` | For experiment tracking | WandB API key. Only needed when `--wandb-mode online` is used | If `HF_TOKEN` is already stored in your HF cache (`~/.cache/huggingface/token`), diff --git a/docs/user-guide/environment.md b/docs/user-guide/environment.md index 11a8d80a5..f116c94b0 100644 --- a/docs/user-guide/environment.md +++ b/docs/user-guide/environment.md @@ -48,8 +48,8 @@ Grouped by the `Category` column -- `nss`-native settings first, then | `NSS_WANDB_MODE` | nss | `--wandb-mode` | WandB | `disabled` | WandB run mode | Alias for `WANDB_MODE` | | `NSS_WANDB_PROJECT` | nss | `--wandb-project` | WandB | -- | WandB project name | Alias for `WANDB_PROJECT` | | `NSS_INFERENCE_ENDPOINT` | nss | `--inference-endpoint-url` | PII column classifier | NVIDIA integrate URL | OpenAI-compatible endpoint for column classification | [PII appendix](#pii-ner-and-column-classification) | -| `NSS_INFERENCE_KEY` | nss | `--inference-api-key` | PII column classifier | -- | API key for `NSS_INFERENCE_ENDPOINT` | Required for LLM column classification | -| `NSS_INFERENCE_MODEL` | nss | `--inference-model-id` | PII column classifier | `qwen/qwen3-next-80b-a3b-instruct` | Model ID sent to the inference endpoint | [PII appendix](#pii-ner-and-column-classification) | +| `NSS_INFERENCE_KEY` | nss | `--inference-api-key` | PII column classifier | -- | API key for `NSS_INFERENCE_ENDPOINT` | Required only for the `api` classifier backend | +| `NSS_INFERENCE_MODEL` | nss | `--inference-model-id` | PII column classifier | `qwen/qwen3-next-80b-a3b-instruct` | Model ID sent to the API inference endpoint | [PII appendix](#pii-ner-and-column-classification) | | `NSS_PII_REPLACER_CPU_COUNT` | nss | `--cpu-count` | NER worker pool | `max(1, cpu_count - 1)` | CPU processes for PII NER | [PII appendix](#pii-ner-and-column-classification) | | `NEMO_TELEMETRY_ENABLED` | telemetry | `--emit_telemetry` | telemetry | `true` | Enable anonymous usage telemetry | Also `emit_telemetry` in YAML; see [Telemetry](#telemetry) | | `HF_HOME` | third-party | -- | Hugging Face Hub | platform cache dir | Root directory for HF downloads | [HF appendix](#hugging-face-cache-and-offline) | @@ -177,8 +177,9 @@ replacement. For setup examples and NER-only fallback behavior, see ### `NSS_INFERENCE_ENDPOINT` and `NSS_INFERENCE_KEY` -OpenAI-compatible endpoint and API key for column classification. The endpoint -defaults to `https://integrate.api.nvidia.com/v1` when unset. +OpenAI-compatible endpoint and API key for the default `api` column +classification backend. The endpoint defaults to +`https://integrate.api.nvidia.com/v1` when unset. ```bash export NSS_INFERENCE_ENDPOINT="https://your-llm-inference-endpoint" @@ -188,14 +189,23 @@ export NSS_INFERENCE_KEY="your-api-key" # pragma: allowlist secret On the CLI, can also use `--inference-api-key` and optionally `--inference-endpoint-url` instead of exporting these variables. +For local Hugging Face classification, set +`replace_pii.globals.classify.backend: local_hf` in YAML or the SDK. That mode +does not use `NSS_INFERENCE_KEY`; it loads +`replace_pii.globals.classify.model` or the default +`HuggingFaceTB/SmolLM3-3B` from the Hugging Face cache, online Hub access, or a +complete local model path. + To disable column classification entirely, set `replace_pii.globals.classify.enable_classify: false` in YAML or use the SDK. See [Configuration Reference -- Replacing PII](configuration.md#replacing-pii). ### `NSS_INFERENCE_MODEL` -Model ID sent to the inference endpoint. Defaults to +Model ID sent to the API inference endpoint. Defaults to `qwen/qwen3-next-80b-a3b-instruct`. Override with `--inference-model-id`. +This environment variable is ignored by the `local_hf` backend; use +`replace_pii.globals.classify.model` for local Hugging Face classification. ### `NSS_PII_REPLACER_CPU_COUNT` diff --git a/docs/user-guide/evaluating-data.md b/docs/user-guide/evaluating-data.md index 57d2c58a9..2c542f31e 100644 --- a/docs/user-guide/evaluating-data.md +++ b/docs/user-guide/evaluating-data.md @@ -89,10 +89,10 @@ or Could not perform classify, falling back to default entities. ``` -When `NSS_INFERENCE_KEY` is not set, the same log line is followed by guidance to set it (and a note that `NSS_INFERENCE_ENDPOINT` is optional with the default API). When the key is set, a traceback may be included to show the underlying API error. +When `NSS_INFERENCE_KEY` is not set for the default `api` backend, the same log line is followed by guidance to set it (and a note that `NSS_INFERENCE_ENDPOINT` is optional with the default API). When the key is set, a traceback may be included to show the underlying API error. For local Hugging Face classification, use `replace_pii.globals.classify.backend: local_hf` instead of API environment variables. -Fix: set entity types explicitly in your config, or when using the CLI ensure -`NSS_INFERENCE_KEY` is set (and `NSS_INFERENCE_ENDPOINT` if not using the default). PII classify config is deeply nested -- use YAML or SDK: +Fix: set entity types explicitly in your config, configure the `api` backend with +`NSS_INFERENCE_KEY` (and `NSS_INFERENCE_ENDPOINT` if not using the default), or use the `local_hf` backend with a cached/downloadable model. PII classify config is deeply nested -- use YAML or SDK: === "Config reference" diff --git a/docs/user-guide/running.md b/docs/user-guide/running.md index 07f84c42e..4f1caa1f9 100644 --- a/docs/user-guide/running.md +++ b/docs/user-guide/running.md @@ -282,7 +282,7 @@ execute in order (`config` → `dataframe` → `metadata` → `advisory`). | Check name | Stage | What it validates | |-------|-------|-------------------| | `gpu.cuda` | config | PyTorch is importable and a CUDA GPU is visible | -| `env.inference` | config | Inference config for PII classification: `NSS_INFERENCE_KEY` is set, `NSS_INFERENCE_MODEL` is non-empty, and `NSS_INFERENCE_ENDPOINT` is a valid http(s) URL (warnings only) | +| `env.inference` | config | Inference config for PII classification: API backend env vars are usable, or the `local_hf` classifier model reference is usable locally or fetchable from Hugging Face | | `env.hf_model_availability` | config | The pretrained model reference is usable locally or can be fetched from Hugging Face; warns about a missing HF token only when online HF access may be needed | | `dataset.size` | dataframe | Training split meets the hard minimum row count | | `columns.groupby` | dataframe | `group_training_examples_by` column is present and has no nulls | @@ -621,22 +621,37 @@ default in both the CLI and SDK. PII on by default means no config flag is neede ### LLM Column Classification -To enable LLM-based PII column classification (optional), set the API key -before running the pipeline. The endpoint defaults to -`https://integrate.api.nvidia.com/v1`; override `NSS_INFERENCE_ENDPOINT` for a -custom OpenAI-compatible endpoint. +LLM-based PII column classification is optional and can use either the default +OpenAI-compatible API backend or a local Hugging Face model. -When using the CLI, set both for column classification: +For API classification, set the API key before running the pipeline. The +endpoint defaults to `https://integrate.api.nvidia.com/v1`; override +`NSS_INFERENCE_ENDPOINT` for a custom OpenAI-compatible endpoint. ```bash export NSS_INFERENCE_ENDPOINT="https://integrate.api.nvidia.com/v1" # optional; this is the default export NSS_INFERENCE_KEY="your-api-key" # pragma: allowlist secret (required for column classification with the inference endpoint) ``` -PII column classification requires `NSS_INFERENCE_KEY` (and optionally `NSS_INFERENCE_ENDPOINT` if not using the default). -When `NSS_INFERENCE_KEY` is unset, the classification step is attempted but -falls back to NER-only detection (with an error log). No environment -variables are required for NER-only PII replacement. +For local classification, configure the `local_hf` backend. No inference API +key is required: + +```yaml +replace_pii: + globals: + classify: + backend: local_hf + model: HuggingFaceTB/SmolLM3-3B # optional default; can also be a local model path +``` + +The local backend loads the default 3B model from the Hugging Face cache or +downloads it when online. In HF offline mode (`HF_HUB_OFFLINE=1` or +`TRANSFORMERS_OFFLINE=1`), pre-download the complete model snapshot or set +`replace_pii.globals.classify.model` to a complete local model directory. + +When the API key is unset for the `api` backend, classification falls back to +NER-only detection after logging an error. No environment variables are required +for NER-only PII replacement. See [Configuration Reference -- Replacing PII](configuration.md#replacing-pii) for the full parameter reference. diff --git a/docs/user-guide/troubleshooting.md b/docs/user-guide/troubleshooting.md index 6f6f34d2d..d11a97547 100644 --- a/docs/user-guide/troubleshooting.md +++ b/docs/user-guide/troubleshooting.md @@ -517,10 +517,17 @@ check of its own. | `inference_key_missing` | warning | `env.inference` | `NSS_INFERENCE_KEY` not set; PII classification degraded | | `inference_model_blank` | warning | `env.inference` | `NSS_INFERENCE_MODEL` set but empty; the blank value is ignored and the default model id is used | | `inference_endpoint_invalid` | error | `env.inference` | `NSS_INFERENCE_ENDPOINT` set but not a valid http(s) URL; classification requests will fail | -| `hf_token_missing` | warning | `env.hf_model_availability` | Neither `HF_TOKEN` nor `HUGGING_FACE_HUB_TOKEN` set, and model loading may need online Hugging Face access | +| `classify_model_ref_empty` | error | `env.inference` | `replace_pii.globals.classify.model` is empty for the `local_hf` classifier backend | +| `classify_model_ref_invalid` | error | `env.inference` | Local classifier model value is neither an existing path nor a valid Hugging Face model ID | +| `classify_local_model_missing` | error | `env.inference` | Local classifier model is path-like, but the path does not exist | +| `classify_local_model_not_directory` | error | `env.inference` | Local classifier model path exists but is not a directory | +| `classify_local_model_incomplete` | error | `env.inference` | Local classifier model directory is missing required config, tokenizer, weights, or shards | +| `classify_hf_model_not_cached` | warning/error | `env.inference` | Local classifier Hugging Face model is not present in the local cache; severity is error when HF offline mode is enabled | +| `classify_hf_model_cache_incomplete` | warning/error | `env.inference` | Cached local classifier model snapshot is missing required files; severity is error when HF offline mode is enabled | +| `hf_token_missing` | warning | `env.hf_model_availability` / `env.inference` | Neither `HF_TOKEN` nor `HUGGING_FACE_HUB_TOKEN` set, and model loading may need online Hugging Face access | | `hf_model_not_cached` | warning/error | `env.hf_model_availability` | Hugging Face model is not present in the local cache; severity is error when HF offline mode is enabled | | `hf_model_cache_incomplete` | warning/error | `env.hf_model_availability` | Cached Hugging Face model snapshot is missing required config, tokenizer, weights, or shards; severity is error when HF offline mode is enabled | -| `hf_remote_code_not_cached` | warning/error | `env.hf_model_availability` | Trusted model references remote code that is not cached locally; severity is error when HF offline mode is enabled | +| `hf_remote_code_not_cached` | warning/error | `env.hf_model_availability` / `env.inference` | Trusted model references remote code that is not cached locally; severity is error when HF offline mode is enabled | | `preflight.check_crash` | error | (crashing check) | A check raised an unexpected exception; the issue's `check` field names the crashing check and other checks continued running | | `column_not_found` | error | `columns.groupby` / `columns.orderby` | Required column missing from dataset, or input DataFrame uses unsupported MultiIndex columns | | `column_nulls` | error | `columns.groupby` | Required column contains null values | diff --git a/src/nemo_safe_synthesizer/config/replace_pii.py b/src/nemo_safe_synthesizer/config/replace_pii.py index a5aebef57..446eea552 100644 --- a/src/nemo_safe_synthesizer/config/replace_pii.py +++ b/src/nemo_safe_synthesizer/config/replace_pii.py @@ -4,7 +4,8 @@ from __future__ import annotations import os -from typing import Annotated, Any, Self +import warnings +from typing import Annotated, Any, Literal, Self from faker.config import AVAILABLE_LOCALES from pydantic import Field, field_validator, model_validator @@ -183,12 +184,34 @@ class ClassifyConfig(NSSBaseModel): num_samples: int | None = Field(description="Number of column values to sample for classification.", default=3) + backend: Literal["api", "local_hf"] = Field( + default="api", + description="Column classification backend. Use 'api' for an OpenAI-compatible endpoint or 'local_hf' for an in-process Hugging Face model.", + ) + + model: str | None = Field( + default=None, + description="Model name or local path for column classification. For the local_hf backend, defaults to HuggingFaceTB/SmolLM3-3B.", + ) + classify_model_provider: str | None = Field( default=None, description="Name of the model provider in the Inference Gateway for column classification. " "The job compiler will resolve this to the appropriate endpoint URL.", ) + @model_validator(mode="after") + def warn_api_model_ignored(self) -> Self: + """Warn when a local-HF-only model setting is provided for API classification.""" + if self.backend != "local_hf" and self.model is not None: + warnings.warn( + "`replace_pii.globals.classify.model` is only used when " + "`replace_pii.globals.classify.backend` is 'local_hf'. " + "For the api backend, set `NSS_INFERENCE_MODEL` or use `--inference-model-id`.", + stacklevel=2, + ) + return self + class Globals(NSSBaseModel): """Global settings for the PII replacer including locales, seed, NER, and classification.""" diff --git a/src/nemo_safe_synthesizer/defaults.py b/src/nemo_safe_synthesizer/defaults.py index fdfd232c2..2e343c020 100644 --- a/src/nemo_safe_synthesizer/defaults.py +++ b/src/nemo_safe_synthesizer/defaults.py @@ -73,6 +73,7 @@ # default LLM inference endpoint for PII column classification. DEFAULT_NSS_INFERENCE_ENDPOINT = "https://integrate.api.nvidia.com/v1" +DEFAULT_PII_CLASSIFY_LOCAL_MODEL = "HuggingFaceTB/SmolLM3-3B" # training + parameters DEFAULT_BASE_SEQ_LENGTH = 2048 diff --git a/src/nemo_safe_synthesizer/llm/model_host.py b/src/nemo_safe_synthesizer/llm/model_host.py new file mode 100644 index 000000000..72fb982ab --- /dev/null +++ b/src/nemo_safe_synthesizer/llm/model_host.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed lifecycle contract for components that own a local model.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +ModelT = TypeVar("ModelT") +TokenizerT = TypeVar("TokenizerT") + + +class ModelHost(ABC, Generic[ModelT, TokenizerT]): + """Own a local language model and tokenizer through teardown. + + This contract deliberately stops at model ownership. Tasks such as + synthetic-record generation and column classification retain their own + prompt construction, batching, and response parsing. + """ + + @property + @abstractmethod + def model(self) -> ModelT | None: + """Return the hosted model, or ``None`` before initialization.""" + + @property + @abstractmethod + def tokenizer(self) -> TokenizerT | None: + """Return the hosted tokenizer, or ``None`` before initialization.""" + + @abstractmethod + def initialize(self) -> None: + """Load the model and any resources required to use it.""" + + @abstractmethod + def teardown(self) -> None: + """Release the model and its resources; this must be idempotent.""" diff --git a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py index 2bba807ee..c536a7fad 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py +++ b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py @@ -9,9 +9,10 @@ from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass from itertools import chain, islice +from pathlib import Path from time import monotonic from timeit import default_timer as timer -from typing import Optional +from typing import TYPE_CHECKING, Any, Optional import json_repair import pandas as pd @@ -20,6 +21,9 @@ from openai import OpenAI from pydantic import ConfigDict, TypeAdapter, ValidationError +from ...defaults import DEFAULT_PII_CLASSIFY_LOCAL_MODEL +from ...llm.model_host import ModelHost +from ...llm.utils import ModelRef, cleanup_memory from ...observability import get_logger from ...utils import hf_offline_enabled from ..ner import ner_mp @@ -29,6 +33,12 @@ logger = get_logger(__name__) +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizerBase +else: + PreTrainedModel = Any + PreTrainedTokenizerBase = Any + class DefaultLLMConfig: """Default settings for the LLM used in column classification. @@ -49,6 +59,7 @@ class DefaultLLMConfig: SYSTEM_PROMPT = "You are a helpful AI that annotates columns in datasets with their respective types. " MAX_OUTPUT_TOKENS = 2048 TEMPERATURE = 0.2 + LOCAL_HF_CONFIG_ID = DEFAULT_PII_CLASSIFY_LOCAL_MODEL @classmethod def config_id(cls) -> str: @@ -259,7 +270,24 @@ def classify_columns( formatted_prompt = _format_prompt(df, entities, num_samples) if not formatted_prompt: return {} + if client is None: + raise RuntimeError("InferenceAPI classifier not initialized. Use get_classifier() method.") + + entities_str = _classify_prompt_with_openai( + formatted_prompt=formatted_prompt, + client=client, + logger=logger, + ) + return _filter_entities(entities_str, entities, on_validation_error) + +def _classify_prompt_with_openai( + *, + formatted_prompt: str, + client: OpenAI, + logger: logging.Logger, +) -> str: + """Send one formatted classification prompt to an OpenAI-compatible backend.""" llm_start = timer() response = client.chat.completions.create( model=DefaultLLMConfig.config_id(), @@ -281,6 +309,15 @@ def classify_columns( }, ) + return entities_str + + +def _filter_entities( + entities_str: str, + entities: set[str], + on_validation_error: Callable[[], None], +) -> dict[str, Optional[str]]: + """Parse classifier output and map labels outside the valid entity set to ``none``.""" col_entities = _try_extract_entities(entities_str, on_validation_error) return {col: ent if ent in entities else UNKNOWN_ENTITY for col, ent in col_entities.items()} @@ -333,21 +370,37 @@ def _try_extract_entities( class ColumnClassifier(ABC): """Abstract column-type classifier; implementations may use LLM, VertexAI, or other backends.""" - @abstractmethod + _num_samples: Optional[int] + def detect_types(self, df: pd.DataFrame, entities: Optional[set[str]]) -> dict[str, Optional[str]]: """Classify each column into one of the given entity types. - Implementations may sample column values and use an LLM, lookup table, or - other backend to assign exactly one entity type per column. Columns that - cannot be classified or are not in ``entities`` should be mapped to - ``UNKNOWN_ENTITY``. - Args: df: DataFrame whose columns are to be classified. entities: Set of valid entity type names to assign; may be ``None`` for implementations that use a fixed or default set. """ - ... + valid_entities = entities if entities is not None else DEFAULT_ENTITIES + formatted_prompt = _format_prompt(df, valid_entities, self._num_samples) + if not formatted_prompt: + return {} + + entities_str = self._classify_prompt(formatted_prompt) + return _filter_entities(entities_str, valid_entities, self._on_validation_error) + + @abstractmethod + def _classify_prompt(self, formatted_prompt: str) -> str: + """Return raw JSON-ish classifier output for a formatted prompt.""" + + def _on_validation_error(self) -> None: + raise RuntimeError( + "There was an error performing classification: " + "the classifier LLM failed to return valid JSON. " + "Please reach out to support if the error recurs." + ) + + def close(self) -> None: + """Release backend resources after classification.""" class ColumnClassifierNoop(ColumnClassifier): @@ -356,6 +409,9 @@ class ColumnClassifierNoop(ColumnClassifier): def detect_types(self, df: pd.DataFrame, entities: Optional[set[str]] = None) -> dict[str, Optional[str]]: return {col: UNKNOWN_ENTITY for col in df.columns} + def _classify_prompt(self, formatted_prompt: str) -> str: + raise NotImplementedError("ColumnClassifierNoop.detect_types does not classify prompts.") + @dataclass class IAPIClassifierConfig: @@ -391,15 +447,27 @@ def detect_types(self, df: pd.DataFrame, entities: set[str]) -> dict[str, Option if self._llm is None: raise Exception("InferenceAPI classifier not initialized. Use get_classifier() method.") - return classify_columns( - df=df, - entities=entities, - num_samples=self._num_samples, + return super().detect_types(df, entities) + + def _classify_prompt(self, formatted_prompt: str) -> str: + if self._llm is None: + raise RuntimeError("InferenceAPI classifier not initialized. Use get_classifier() method.") + return _classify_prompt_with_openai( + formatted_prompt=formatted_prompt, client=self._llm, - on_validation_error=self._on_validation_error, logger=logger, ) + def close(self) -> None: + if self._llm is None: + return + try: + self._llm.close() + except Exception: + logger.debug("OpenAI client cleanup failed during column classifier teardown", exc_info=True) + finally: + self._llm = None + def _on_validation_error(self) -> None: raise RuntimeError( "There was an error performing classification: " @@ -408,6 +476,175 @@ def _on_validation_error(self) -> None: ) +class ColumnClassifierHF(ColumnClassifier, ModelHost[PreTrainedModel, PreTrainedTokenizerBase]): + """Classify column types with an in-process Hugging Face causal LM.""" + + _model_name_or_path: str + _num_samples: Optional[int] + _model: PreTrainedModel | None + _tokenizer: PreTrainedTokenizerBase | None + + def __init__(self, model_name_or_path: str, num_samples: Optional[int]): + self._model_name_or_path = model_name_or_path + self._num_samples = num_samples + self._model = None + self._tokenizer = None + + @property + def model(self) -> PreTrainedModel | None: + """Return the local classifier model, if loaded.""" + return self._model + + @property + def tokenizer(self) -> PreTrainedTokenizerBase | None: + """Return the local classifier tokenizer, if loaded.""" + return self._tokenizer + + def _classify_prompt(self, formatted_prompt: str) -> str: + self.initialize() + model = self.model + tokenizer = self.tokenizer + if model is None or tokenizer is None: + raise RuntimeError("Local Hugging Face classifier failed to initialize.") + + input_ids, attention_mask = self._encode_prompt(tokenizer, model, formatted_prompt) + + pad_token_id = tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = tokenizer.eos_token_id + + llm_start = timer() + with torch.no_grad(): + output_ids = model.generate( + input_ids=input_ids, + attention_mask=attention_mask, + max_new_tokens=DefaultLLMConfig.MAX_OUTPUT_TOKENS, + do_sample=False, + pad_token_id=pad_token_id, + eos_token_id=tokenizer.eos_token_id, + ) + generated_ids = output_ids[0, input_ids.shape[-1] :] + entities_str = tokenizer.decode(generated_ids, skip_special_tokens=True) + llm_elapsed = timer() - llm_start + logger.info( + f"Local HF column classification took {llm_elapsed} seconds.", + extra={ + "ctx": { + "llm_elapsed": llm_elapsed, + "model": self._model_name_or_path, + }, + }, + ) + return entities_str + + def _encode_prompt( + self, + tokenizer: PreTrainedTokenizerBase, + model: PreTrainedModel, + formatted_prompt: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + messages = [ + {"role": "system", "content": DefaultLLMConfig.SYSTEM_PROMPT}, + {"role": "user", "content": formatted_prompt}, + ] + encoded: Any + if getattr(tokenizer, "chat_template", None): + encoded = tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + return_tensors="pt", + ) + else: + prompt = self._plain_prompt(formatted_prompt) + encoded = tokenizer(prompt, return_tensors="pt") + + try: + input_ids = encoded["input_ids"] + attention_mask = encoded.get("attention_mask", None) + except (KeyError, TypeError): + input_ids = encoded + attention_mask = None + input_ids = input_ids.to(model.device) + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) + else: + attention_mask = attention_mask.to(model.device) + + return input_ids, attention_mask + + @staticmethod + def _plain_prompt(formatted_prompt: str) -> str: + return f"{DefaultLLMConfig.SYSTEM_PROMPT}\n\n{formatted_prompt}" + + def initialize(self) -> None: + if self._model is not None and self._tokenizer is not None: + return + + from transformers import AutoModelForCausalLM, AutoTokenizer + + model_ref = ModelRef.parse(self._model_name_or_path) + load_target = self._load_target(model_ref) + load_kwargs = { + "trust_remote_code": model_ref.trust_remote_code, + "local_files_only": hf_offline_enabled(), + } + if torch.cuda.is_available(): + load_kwargs["device_map"] = "auto" + load_kwargs["dtype"] = torch.bfloat16 + + logger.info("Loading local column classification model: %s", self._model_name_or_path) + self._tokenizer = AutoTokenizer.from_pretrained( + load_target, + trust_remote_code=model_ref.trust_remote_code, + local_files_only=hf_offline_enabled(), + ) + if getattr(self._tokenizer, "pad_token_id", None) is None: + self._tokenizer.pad_token = self._tokenizer.eos_token + self._model = AutoModelForCausalLM.from_pretrained(load_target, **load_kwargs) + if not torch.cuda.is_available(): + self._model = self._model.to("cpu") + self._model.eval() + + def _load(self) -> None: + """Backward-compatible alias for older tests and callers.""" + self.initialize() + + @staticmethod + def _load_target(model_ref: ModelRef) -> str | Path: + """Prefer the repo ID over an incomplete online cache snapshot. + + Transformers treats a snapshot directory as an authoritative local + model. If the HF cache only has metadata or tokenizer files, passing + that path prevents Transformers from downloading the missing files. + """ + if ( + model_ref.repo_id is not None + and model_ref.local_path is not None + and not hf_offline_enabled() + and ModelRef.missing_required_components(model_ref.local_path) + ): + return model_ref.repo_id + return model_ref.target() + + def teardown(self) -> None: + self._model = None + self._tokenizer = None + try: + cleanup_memory() + except Exception: + logger.debug("cleanup_memory failed during column classifier teardown", exc_info=True) + + def close(self) -> None: + self.teardown() + + def _on_validation_error(self) -> None: + raise RuntimeError( + "There was an error performing classification: " + "the local classifier LLM failed to return valid JSON. " + "Try the api backend or a different local classification model if the error recurs." + ) + + @dataclass class ClassifyConfig: """Configuration for column classification and NER (entities, thresholds, GLiNER, regex).""" diff --git a/src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py b/src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py index 9710bcc16..471b45f4b 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py +++ b/src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py @@ -22,7 +22,10 @@ DEFAULT_ENTITIES, UNKNOWN_ENTITY, ClassifyConfig, + ColumnClassifier, + ColumnClassifierHF, ColumnClassifierLLM, + DefaultLLMConfig, EntityExtractor, EntityExtractorGliner, EntityExtractorMulti, @@ -132,8 +135,13 @@ def _get_classify_endpoint_url() -> str: return url -def _column_classify_failure_remediation(exc: BaseException) -> str: +def _column_classify_failure_remediation(exc: BaseException, config: PiiReplacerConfig | None = None) -> str: """Extra log text after column classifier init or classify failures.""" + if config is not None and config.globals.classify.backend == "local_hf": + return ( + " Check the local Hugging Face column classification model configuration and cache. " + f"({type(exc).__name__}: {exc})" + ) if not has_inference_key(): return ( " Please set NSS_INFERENCE_KEY in the environment. Get an API key at https://build.nvidia.com/settings/api-keys. " @@ -144,10 +152,18 @@ def _column_classify_failure_remediation(exc: BaseException) -> str: ) -def get_column_classifier() -> ColumnClassifierLLM: - """Return a column classifier backed by the NSS inference endpoint (``NSS_INFERENCE_ENDPOINT``, ``NSS_INFERENCE_KEY``).""" +def get_column_classifier(config: PiiReplacerConfig | None = None) -> ColumnClassifier: + """Return the configured column classifier backend.""" + pii_config = config or PiiReplacerConfig.get_default_config() + classify_config = pii_config.globals.classify + num_samples = classify_config.num_samples + + if classify_config.backend == "local_hf": + model = classify_config.model or DefaultLLMConfig.LOCAL_HF_CONFIG_ID + return ColumnClassifierHF(model_name_or_path=model, num_samples=num_samples) + classifier = ColumnClassifierLLM() - classifier._num_samples = 5 + classifier._num_samples = num_samples endpoint = _get_classify_endpoint_url() @@ -289,12 +305,12 @@ def classify_df(self, df: pd.DataFrame) -> list[ColumnClassification]: # Try to initialize the column classifier try: - column_classifier = get_column_classifier() + column_classifier = get_column_classifier(self.pii_replacer_config) except Exception as exc: logging.error( "Could not initialize column classifier, PII replacement will run in degraded mode. NER Falling back to default entities. No replacement done except for text columns. %s", - _column_classify_failure_remediation(exc), - exc_info=has_inference_key(), + _column_classify_failure_remediation(exc, self.pii_replacer_config), + exc_info=has_inference_key() or self.pii_replacer_config.globals.classify.backend == "local_hf", ) # Try to perform classification if we successfully got a classifier @@ -313,9 +329,12 @@ def classify_df(self, df: pd.DataFrame) -> list[ColumnClassification]: except Exception as exc: logging.error( "Could not initialize column classifier, PII replacement will run in degraded mode. NER Falling back to default entities. No replacement done except for text columns. %s", - _column_classify_failure_remediation(exc), - exc_info=has_inference_key(), + _column_classify_failure_remediation(exc, self.pii_replacer_config), + exc_info=has_inference_key() + or self.pii_replacer_config.globals.classify.backend == "local_hf", ) + finally: + column_classifier.close() else: logging.info("Column classification is disabled (enable_classify=False), skipping classify call.") finally: diff --git a/src/nemo_safe_synthesizer/preflight/checks/environment.py b/src/nemo_safe_synthesizer/preflight/checks/environment.py index 8a2bdd2b7..74ecc6117 100644 --- a/src/nemo_safe_synthesizer/preflight/checks/environment.py +++ b/src/nemo_safe_synthesizer/preflight/checks/environment.py @@ -13,6 +13,7 @@ from typing_extensions import override +from ...defaults import DEFAULT_PII_CLASSIFY_LOCAL_MODEL from ...llm.utils import ModelRef from ...observability import get_logger from ...utils import hf_offline_enabled @@ -463,16 +464,210 @@ def _is_valid_http_url(value: str | None) -> bool: return parsed.scheme in ("http", "https") and bool(parsed.netloc) +@dataclass(frozen=True) +class _HFModelIssueCodes: + empty: str + missing_local: str + invalid_ref: str + local_not_directory: str + local_incomplete: str + not_cached: str + cache_incomplete: str + remote_code_not_cached: str = "hf_remote_code_not_cached" + token_missing: str = "hf_token_missing" + + +@dataclass(frozen=True) +class _HFModelCheckSpec: + field_path: str + model_label: str + cached_model_label: str + local_model_label: str + offline_failure: str + incomplete_cache_online: str + codes: _HFModelIssueCodes + + +_TRAINING_HF_MODEL_SPEC = _HFModelCheckSpec( + field_path="training.pretrained_model", + model_label="Hugging Face model", + cached_model_label="Hugging Face model", + local_model_label="model", + offline_failure="model loading will fail", + incomplete_cache_online="Model loading will contact Hugging Face unless the full model snapshot is pre-downloaded.", + codes=_HFModelIssueCodes( + empty="model_ref_empty", + missing_local="local_model_missing", + invalid_ref="model_ref_invalid", + local_not_directory="local_model_not_directory", + local_incomplete="local_model_incomplete", + not_cached="hf_model_not_cached", + cache_incomplete="hf_model_cache_incomplete", + ), +) + + +_CLASSIFY_HF_MODEL_SPEC = _HFModelCheckSpec( + field_path="replace_pii.globals.classify.model", + model_label="Column classification Hugging Face model", + cached_model_label="column classification model", + local_model_label="column classification model", + offline_failure="local column classification will fail", + incomplete_cache_online=( + "Model loading will use the Hugging Face model ID to fetch missing files unless " + "the full model snapshot is pre-downloaded." + ), + codes=_HFModelIssueCodes( + empty="classify_model_ref_empty", + missing_local="classify_local_model_missing", + invalid_ref="classify_model_ref_invalid", + local_not_directory="classify_local_model_not_directory", + local_incomplete="classify_local_model_incomplete", + not_cached="classify_hf_model_not_cached", + cache_incomplete="classify_hf_model_cache_incomplete", + ), +) + + +def _check_hf_model_reference(model_name: str, collector: IssueCollector, *, spec: _HFModelCheckSpec) -> None: + """Validate a HF model reference, local path, cache snapshot, and token readiness.""" + if not model_name: + collector.error(spec.codes.empty, f"`{spec.field_path}` must not be empty.") + return + + model_ref = ModelRef.parse(model_name) + if model_ref.local_path is not None and Path(model_name).resolve(strict=False) == model_ref.local_path.resolve( + strict=False + ): + _check_hf_local_path(model_ref.local_path, collector, model_ref=model_ref, spec=spec) + return + + if _is_missing_local_path(model_name): + collector.error(spec.codes.missing_local, f"`{spec.field_path}` points to missing local path '{model_name}'.") + return + + if model_ref.repo_id is None: + collector.error( + spec.codes.invalid_ref, + ( + f"`{spec.field_path}` value '{model_name}' is neither an existing local path " + "nor a valid Hugging Face model ID." + ), + ) + return + + snapshot_path = model_ref.local_path or model_ref.partial_cached_snapshot() + if snapshot_path is None: + _report_missing_hf_cache(model_ref, collector, spec=spec) + return + + missing = ModelRef.missing_required_components(snapshot_path) + if missing: + _report_incomplete_hf_cache(model_ref, snapshot_path, missing, collector, spec=spec) + _report_missing_remote_code(model_ref, snapshot_path, collector, spec=spec) + + +def _check_hf_local_path( + model_path: Path, + collector: IssueCollector, + *, + model_ref: ModelRef, + spec: _HFModelCheckSpec, +) -> None: + if not model_path.is_dir(): + collector.error( + spec.codes.local_not_directory, + f"`{spec.field_path}` points to '{model_path}', but local models must be directories.", + ) + return + + missing = ModelRef.missing_required_components(model_path) + if missing: + collector.error( + spec.codes.local_incomplete, + f"Local {spec.local_model_label} directory '{model_path}' is missing {', '.join(missing)}.", + ) + _report_missing_remote_code(model_ref, model_path, collector, spec=spec) + + +def _report_missing_hf_cache(model_ref: ModelRef, collector: IssueCollector, *, spec: _HFModelCheckSpec) -> None: + message = f"{spec.model_label} '{model_ref.repo_id}' is not present in the local cache at '{model_ref.cache_root}'." + if hf_offline_enabled(): + collector.error( + spec.codes.not_cached, + f"{message} Offline Hugging Face mode is enabled, so {spec.offline_failure}.", + ) + return + collector.warning( + spec.codes.not_cached, + f"{message} Model loading will contact Hugging Face unless the model is pre-downloaded.", + ) + _report_missing_hf_token(collector, code=spec.codes.token_missing) + + +def _report_incomplete_hf_cache( + model_ref: ModelRef, + snapshot_path: Path, + missing: list[str], + collector: IssueCollector, + *, + spec: _HFModelCheckSpec, +) -> None: + message = ( + f"Cached {spec.cached_model_label} '{model_ref.repo_id}' at '{snapshot_path}' is missing {', '.join(missing)}." + ) + if hf_offline_enabled(): + collector.error( + spec.codes.cache_incomplete, + f"{message} Offline Hugging Face mode is enabled, so {spec.offline_failure}.", + ) + return + collector.warning( + spec.codes.cache_incomplete, + f"{message} {spec.incomplete_cache_online}", + ) + _report_missing_hf_token(collector, code=spec.codes.token_missing) + + +def _report_missing_remote_code( + model_ref: ModelRef, + model_path: Path, + collector: IssueCollector, + *, + spec: _HFModelCheckSpec, +) -> None: + if not model_ref.trust_remote_code: + return + + missing = ModelRef.missing_remote_code_components(model_path) + if not missing: + return + + message = ( + f"Trusted Hugging Face model '{model_ref.repo_id}' at '{model_path}' references remote code " + f"that is not cached locally: {', '.join(missing)}." + ) + if hf_offline_enabled(): + collector.error( + spec.codes.remote_code_not_cached, + f"{message} Offline Hugging Face mode is enabled, so Transformers cannot fetch it.", + ) + return + collector.warning( + spec.codes.remote_code_not_cached, + f"{message} Model loading may contact Hugging Face to fetch it.", + ) + _report_missing_hf_token(collector, code=spec.codes.token_missing) + + class InferenceModelCheck(ConfigCheck): """Validate the inference configuration used for PII column classification. - When classification is enabled, the runtime calls an OpenAI-compatible + When API classification is enabled, the runtime calls an OpenAI-compatible inference endpoint configured by ``NSS_INFERENCE_KEY``, - ``NSS_INFERENCE_MODEL``, and ``NSS_INFERENCE_ENDPOINT`` (set directly or via - the matching CLI flags, which are propagated to the environment before - preflight runs). This check reads those env vars -- not ``config`` -- because - the inference settings live in ``CLISettings``/the environment rather than in - ``SafeSynthesizerParameters``. + ``NSS_INFERENCE_MODEL``, and ``NSS_INFERENCE_ENDPOINT``. When local Hugging + Face classification is enabled, the runtime loads + ``replace_pii.globals.classify.model`` or the local classifier default. The body uses a single-dispatch ``match`` over ``(model, key, endpoint)``, so at most one finding is emitted per run -- the highest-priority problem. @@ -493,6 +688,12 @@ def check(self, ctx: ConfigView, collector: IssueCollector) -> None: if config.replace_pii is None or config.replace_pii.globals.classify.enable_classify is False: return + classify_config = config.replace_pii.globals.classify + if classify_config.backend == "local_hf": + model_name = classify_config.model or DEFAULT_PII_CLASSIFY_LOCAL_MODEL + _check_hf_model_reference(model_name, collector, spec=_CLASSIFY_HF_MODEL_SPEC) + return + model = os.environ.get("NSS_INFERENCE_MODEL") key = os.environ.get("NSS_INFERENCE_KEY") endpoint = os.environ.get("NSS_INFERENCE_ENDPOINT") @@ -540,124 +741,16 @@ class HFModelAvailabilityCheck(ConfigCheck): @override def check(self, ctx: ConfigView, collector: IssueCollector) -> None: model_name = ctx.config.training.pretrained_model - if not model_name: - collector.error("model_ref_empty", "`training.pretrained_model` must not be empty.") - return - - model_ref = ModelRef.parse(model_name) - if model_ref.local_path is not None and Path(model_name).resolve(strict=False) == model_ref.local_path.resolve( - strict=False - ): - self._check_local_path(model_ref.local_path, collector, model_ref=model_ref) - return - - if _is_missing_local_path(model_name): - collector.error( - "local_model_missing", - f"`training.pretrained_model` points to missing local path '{model_name}'.", - ) - return - - if model_ref.repo_id is None: - collector.error( - "model_ref_invalid", - ( - f"`training.pretrained_model` value '{model_name}' is neither an existing local path " - "nor a valid Hugging Face model ID." - ), - ) - return - - snapshot_path = model_ref.local_path or model_ref.partial_cached_snapshot() - if snapshot_path is None: - self._report_missing_cache(model_ref, collector) - return - - missing = ModelRef.missing_required_components(snapshot_path) - if missing: - message = ( - f"Cached Hugging Face model '{model_ref.repo_id}' at '{snapshot_path}' is missing {', '.join(missing)}." - ) - if hf_offline_enabled(): - collector.error( - "hf_model_cache_incomplete", - f"{message} Offline Hugging Face mode is enabled, so model loading will fail.", - ) - return - collector.warning( - "hf_model_cache_incomplete", - f"{message} Model loading will contact Hugging Face unless the full model snapshot is pre-downloaded.", - ) - self._report_missing_hf_token(collector) - self._report_missing_remote_code(model_ref, snapshot_path, collector) - - @staticmethod - def _check_local_path(model_path: Path, collector: IssueCollector, *, model_ref: ModelRef) -> None: - if not model_path.is_dir(): - collector.error( - "local_model_not_directory", - f"`training.pretrained_model` points to '{model_path}', but local models must be directories.", - ) - return - - missing = ModelRef.missing_required_components(model_path) - if missing: - collector.error( - "local_model_incomplete", - f"Local model directory '{model_path}' is missing {', '.join(missing)}.", - ) - HFModelAvailabilityCheck._report_missing_remote_code(model_ref, model_path, collector) - - @staticmethod - def _report_missing_cache(model_ref: ModelRef, collector: IssueCollector) -> None: - message = ( - f"Hugging Face model '{model_ref.repo_id}' is not present in the local cache at '{model_ref.cache_root}'." - ) - if hf_offline_enabled(): - collector.error( - "hf_model_not_cached", - f"{message} Offline Hugging Face mode is enabled, so model loading will fail.", - ) - return - collector.warning( - "hf_model_not_cached", - f"{message} Model loading will contact Hugging Face unless the model is pre-downloaded.", - ) - HFModelAvailabilityCheck._report_missing_hf_token(collector) - - @staticmethod - def _report_missing_remote_code(model_ref: ModelRef, model_path: Path, collector: IssueCollector) -> None: - if not model_ref.trust_remote_code: - return - - missing = ModelRef.missing_remote_code_components(model_path) - if not missing: - return - - message = ( - f"Trusted Hugging Face model '{model_ref.repo_id}' at '{model_path}' references remote code " - f"that is not cached locally: {', '.join(missing)}." - ) - if hf_offline_enabled(): - collector.error( - "hf_remote_code_not_cached", - f"{message} Offline Hugging Face mode is enabled, so Transformers cannot fetch it.", - ) - return - collector.warning( - "hf_remote_code_not_cached", - f"{message} Model loading may contact Hugging Face to fetch it.", - ) - HFModelAvailabilityCheck._report_missing_hf_token(collector) - - @staticmethod - def _report_missing_hf_token(collector: IssueCollector) -> None: - if _has_hf_token(): - return - collector.warning( - "hf_token_missing", - ( - "HF_TOKEN is not set. Model downloads from gated repos will fail. " - "Set HF_TOKEN or HUGGING_FACE_HUB_TOKEN in your environment." - ), - ) + _check_hf_model_reference(model_name, collector, spec=_TRAINING_HF_MODEL_SPEC) + + +def _report_missing_hf_token(collector: IssueCollector, *, code: str) -> None: + if _has_hf_token(): + return + collector.warning( + code, + ( + "HF_TOKEN is not set. Model downloads from gated repos will fail. " + "Set HF_TOKEN or HUGGING_FACE_HUB_TOKEN in your environment." + ), + ) diff --git a/tests/config/test_nss_config.py b/tests/config/test_nss_config.py index a3fc694d7..a916d898d 100644 --- a/tests/config/test_nss_config.py +++ b/tests/config/test_nss_config.py @@ -13,6 +13,7 @@ SafeSynthesizerParameters, TimeSeriesParameters, ) +from nemo_safe_synthesizer.config.replace_pii import ClassifyConfig from nemo_safe_synthesizer.configurator.parameters import Parameters from nemo_safe_synthesizer.configurator.validators import ValueValidator @@ -124,6 +125,24 @@ def test_create_default(self): params = PiiReplacerConfig.get_default_config() assert params.globals.ner.ner_threshold == 0.3 + def test_api_backend_warns_when_model_is_set(self): + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ClassifyConfig(backend="api", model="local/model") + + assert any("NSS_INFERENCE_MODEL" in str(warning.message) for warning in caught) + + def test_local_hf_backend_accepts_model_without_warning(self): + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ClassifyConfig(backend="local_hf", model="local/model") + + assert caught == [] + class TestSafeSynthesizerParameters: @pytest.mark.parametrize( diff --git a/tests/pii_replacer/test_column_classification_eval.py b/tests/pii_replacer/test_column_classification_eval.py new file mode 100644 index 000000000..5e854923e --- /dev/null +++ b/tests/pii_replacer/test_column_classification_eval.py @@ -0,0 +1,433 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Column-classification golden tests and optional live benchmark. + +Fast local checks: + + uv run --frozen pytest tests/pii_replacer/test_column_classification_eval.py \ + -k "gold_fixture or accuracy_metric or benchmark_matrix or summary_report" -n 0 -vv + +Live benchmark: + + export NSS_INFERENCE_KEY= + # Optional for local HF if models are not cached or require gated access: + export HF_TOKEN= + + NSS_RUN_CLASSIFICATION_BENCHMARK=1 \ + NSS_CLASSIFICATION_BENCHMARK_OUTPUT=local_runs/classification-benchmark.json \ + uv run --frozen pytest tests/pii_replacer/test_column_classification_eval.py -m slow -s -n 0 + +The live benchmark runs the API baseline and local HF models as separate pytest +cases. Use comma-separated filters to narrow runs: + + NSS_CLASSIFICATION_BENCHMARK_BACKENDS=local_hf + NSS_CLASSIFICATION_BENCHMARK_API_MODELS=default + NSS_CLASSIFICATION_BENCHMARK_LOCAL_MODELS=smollm3,mistral,tinyllama + +Reports: +- Per-case JSON: ``column-classification-benchmark-{backend}-{model}.json`` +- Combined JSON: ``column-classification-benchmark-summary.json`` +- Combined table: ``column-classification-benchmark-summary.md`` + +When ``NSS_CLASSIFICATION_BENCHMARK_OUTPUT`` is set, the per-case and summary +files are written next to that path using its stem. +""" + +# ruff: noqa: E402 +from __future__ import annotations + +import json +import os +from collections.abc import Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Literal +from unittest.mock import patch + +import pandas as pd +import pytest + +pytest.importorskip("torch", reason="torch is required for these tests (install with: uv sync --extra cpu)") + +from nemo_safe_synthesizer.config.replace_pii import PiiReplacerConfig +from nemo_safe_synthesizer.pii_replacer.data_editor.detect import UNKNOWN_ENTITY, DefaultLLMConfig +from nemo_safe_synthesizer.pii_replacer.nemo_pii import classify_config_from_params, get_column_classifier + +LOCAL_HF_CLASSIFICATION_BENCHMARK_MODELS = { + "smollm3": "HuggingFaceTB/SmolLM3-3B", + "mistral": "mistralai/Mistral-7B-Instruct-v0.3", + "tinyllama": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", +} +API_CLASSIFICATION_BENCHMARK_MODELS = { + "default": DefaultLLMConfig.DEFAULT_CONFIG_ID, +} +CLASSIFICATION_BENCHMARK_BACKENDS = ("api", "local_hf") +CLASSIFICATION_BENCHMARK_LOCK_PATH = Path(".pytest_cache") / "column-classification-benchmark.lock" +CLASSIFICATION_BENCHMARK_CASE_IDS = ( + "pii_dataset_structured", + "credit_card_transactions", + "adult_negative_control", +) + + +def _load_gold(path: Path) -> list[dict[str, Any]]: + return json.loads(path.read_text())["cases"] + + +def _column_accuracy(predicted: Mapping[str, str | None], expected: Mapping[str, str]) -> float: + if not expected: + return 1.0 + _validate_prediction_columns(predicted, expected) + correct = sum(_normalize_label(predicted.get(column)) == label for column, label in expected.items()) + return correct / len(expected) + + +def _positive_recall(predicted: Mapping[str, str | None], expected: Mapping[str, str]) -> float: + _validate_prediction_columns(predicted, expected) + positives = {column: label for column, label in expected.items() if label != UNKNOWN_ENTITY} + if not positives: + return 1.0 + correct = sum(_normalize_label(predicted.get(column)) == label for column, label in positives.items()) + return correct / len(positives) + + +def _validate_prediction_columns(predicted: Mapping[str, str | None], expected: Mapping[str, str]) -> None: + missing = set(expected) - set(predicted) + if missing: + raise AssertionError(f"Classifier did not return predictions for columns: {sorted(missing)}") + + +def _normalize_label(label: str | None) -> str: + return label if label else UNKNOWN_ENTITY + + +def _expected_positive_labels(cases: list[dict[str, Any]]) -> set[str]: + return {label for case in cases for label in case["expected_entities"].values() if label != UNKNOWN_ENTITY} + + +def _parse_csv_env(name: str, default: tuple[str, ...]) -> tuple[str, ...]: + value = os.environ.get(name) + if value is None: + return default + return tuple(item.strip() for item in value.split(",") if item.strip()) + + +def _benchmark_matrix() -> list[tuple[str, str, str]]: + backends = _parse_csv_env("NSS_CLASSIFICATION_BENCHMARK_BACKENDS", CLASSIFICATION_BENCHMARK_BACKENDS) + model_names_by_backend = { + "api": _parse_csv_env("NSS_CLASSIFICATION_BENCHMARK_API_MODELS", tuple(API_CLASSIFICATION_BENCHMARK_MODELS)), + "local_hf": _parse_csv_env( + "NSS_CLASSIFICATION_BENCHMARK_LOCAL_MODELS", tuple(LOCAL_HF_CLASSIFICATION_BENCHMARK_MODELS) + ), + } + + unknown_backends = set(backends) - set(CLASSIFICATION_BENCHMARK_BACKENDS) + if unknown_backends: + pytest.fail(f"Unknown classification benchmark backends: {sorted(unknown_backends)}") + + unknown_api_models = set(model_names_by_backend["api"]) - set(API_CLASSIFICATION_BENCHMARK_MODELS) + if unknown_api_models: + pytest.fail(f"Unknown API classification benchmark models: {sorted(unknown_api_models)}") + + unknown_local_models = set(model_names_by_backend["local_hf"]) - set(LOCAL_HF_CLASSIFICATION_BENCHMARK_MODELS) + if unknown_local_models: + pytest.fail(f"Unknown local classification benchmark models: {sorted(unknown_local_models)}") + + return [ + (backend, model_name, _models_for_backend(backend)[model_name]) + for backend in backends + for model_name in model_names_by_backend[backend] + ] + + +def _models_for_backend(backend: str) -> dict[str, str]: + if backend == "api": + return API_CLASSIFICATION_BENCHMARK_MODELS + return LOCAL_HF_CLASSIFICATION_BENCHMARK_MODELS + + +def _all_benchmark_params() -> list[Any]: + return [ + pytest.param(backend, model_name, model, id=f"{backend}-{model_name}") + for backend in CLASSIFICATION_BENCHMARK_BACKENDS + for model_name, model in _models_for_backend(backend).items() + ] + + +def _benchmark_report_path(backend: str, model_name: str) -> Path: + configured = os.environ.get("NSS_CLASSIFICATION_BENCHMARK_OUTPUT") + if configured is None: + return Path(f"column-classification-benchmark-{backend}-{model_name}.json") + + output_path = Path(configured) + return output_path.with_name(f"{output_path.stem}-{backend}-{model_name}{output_path.suffix}") + + +def _benchmark_summary_paths() -> tuple[Path, Path]: + configured = os.environ.get("NSS_CLASSIFICATION_BENCHMARK_OUTPUT") + if configured is None: + json_path = Path("column-classification-benchmark-summary.json") + else: + output_path = Path(configured) + json_path = output_path.with_name(f"{output_path.stem}-summary.json") + return json_path, json_path.with_suffix(".md") + + +def _sort_result_key(result: dict[str, Any]) -> tuple[int, int]: + backend = result["backend"] + model_name = result["model_name"] + backend_idx = CLASSIFICATION_BENCHMARK_BACKENDS.index(backend) + model_names = tuple(_models_for_backend(backend)) + model_idx = model_names.index(model_name) + return backend_idx, model_idx + + +def _upsert_result(results: list[dict[str, Any]] | Any, result: dict[str, Any]) -> list[dict[str, Any]]: + filtered = [ + existing + for existing in results + if not (existing["backend"] == result["backend"] and existing["model_name"] == result["model_name"]) + ] + filtered.append(result) + return sorted(filtered, key=_sort_result_key) + + +def _render_benchmark_markdown(report: dict[str, Any]) -> str: + header = ["Backend", "Model", "Macro Accuracy", "Macro Positive Recall", *CLASSIFICATION_BENCHMARK_CASE_IDS] + lines = [ + "# Column Classification Benchmark", + "", + "| " + " | ".join(header) + " |", + "| " + " | ".join(["---"] * len(header)) + " |", + ] + for result in report["results"]: + case_scores = {case["case_id"]: case["accuracy"] for case in result["cases"]} + row = [ + result["backend"], + result["model_name"], + f"{result['macro_accuracy']:.3f}", + f"{result['macro_positive_recall']:.3f}", + *( + f"{case_scores[case_id]:.3f}" if case_id in case_scores else "n/a" + for case_id in CLASSIFICATION_BENCHMARK_CASE_IDS + ), + ] + lines.append("| " + " | ".join(row) + " |") + lines.append("") + return "\n".join(lines) + + +def _write_benchmark_reports(result: dict[str, Any]) -> tuple[Path, Path, Path]: + report = { + "api_models": API_CLASSIFICATION_BENCHMARK_MODELS, + "local_hf_models": LOCAL_HF_CLASSIFICATION_BENCHMARK_MODELS, + "results": [result], + } + output_path = _benchmark_report_path(result["backend"], result["model_name"]) + output_path.write_text(json.dumps(report, indent=2) + "\n") + + summary_json_path, summary_md_path = _benchmark_summary_paths() + if summary_json_path.exists(): + summary = json.loads(summary_json_path.read_text()) + else: + summary = { + "api_models": API_CLASSIFICATION_BENCHMARK_MODELS, + "local_hf_models": LOCAL_HF_CLASSIFICATION_BENCHMARK_MODELS, + "results": [], + } + summary["api_models"] = API_CLASSIFICATION_BENCHMARK_MODELS + summary["local_hf_models"] = LOCAL_HF_CLASSIFICATION_BENCHMARK_MODELS + summary["results"] = _upsert_result(summary["results"], result) + summary_json_path.write_text(json.dumps(summary, indent=2) + "\n") + summary_md_path.write_text(_render_benchmark_markdown(summary)) + + return output_path, summary_json_path, summary_md_path + + +@contextmanager +def _classification_benchmark_lock(): + """Serialize live benchmark cases even when pytest-xdist is enabled.""" + try: + import fcntl + except ImportError: + pytest.skip("Column classification live benchmark locking requires fcntl.") + + CLASSIFICATION_BENCHMARK_LOCK_PATH.parent.mkdir(parents=True, exist_ok=True) + with CLASSIFICATION_BENCHMARK_LOCK_PATH.open("w") as lock_file: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + + +def _evaluate_classifier( + backend: Literal["api", "local_hf"], model_name: str, model: str, cases: list[dict[str, Any]] +) -> dict[str, Any]: + config = PiiReplacerConfig.get_default_config() + config.globals.classify.backend = backend + config.globals.classify.model = model + config.globals.classify.entities = sorted( + classify_config_from_params(config).valid_entities | _expected_positive_labels(cases) + ) + classify_config = classify_config_from_params(config) + model_env = {"NSS_INFERENCE_MODEL": model} if backend == "api" else {} + + case_results = [] + with patch.dict("os.environ", model_env, clear=False): + classifier = None + try: + classifier = get_column_classifier(config) + for case in cases: + df = pd.DataFrame(case["rows"]) + predicted = classifier.detect_types(df, classify_config.valid_entities) + accuracy = _column_accuracy(predicted, case["expected_entities"]) + positive_recall = _positive_recall(predicted, case["expected_entities"]) + case_results.append( + { + "case_id": case["id"], + "accuracy": accuracy, + "positive_recall": positive_recall, + "expected": case["expected_entities"], + "predicted": {column: _normalize_label(label) for column, label in predicted.items()}, + } + ) + finally: + if classifier is not None: + classifier.close() + + macro_accuracy = sum(case_result["accuracy"] for case_result in case_results) / len(case_results) + macro_positive_recall = sum(case_result["positive_recall"] for case_result in case_results) / len(case_results) + return { + "backend": backend, + "model_name": model_name, + "model": model, + "status": "ok", + "macro_accuracy": macro_accuracy, + "macro_positive_recall": macro_positive_recall, + "cases": case_results, + } + + +def test_column_classification_gold_fixture_schema(pii_test_data_dir): + cases = _load_gold(pii_test_data_dir / "column_classification_gold.json") + valid_entities = classify_config_from_params(PiiReplacerConfig.get_default_config()).valid_entities + + assert {case["id"] for case in cases} == { + "pii_dataset_structured", + "credit_card_transactions", + "adult_negative_control", + } + for case in cases: + df = pd.DataFrame(case["rows"]) + expected = case["expected_entities"] + assert case["source"].startswith("cleaned/") + assert set(expected) == set(df.columns) + assert all(isinstance(label, str) and label for label in expected.values()) + assert set(expected.values()) <= valid_entities | {UNKNOWN_ENTITY} + + +def test_column_accuracy_metric_normalizes_none_labels(): + expected = {"name": "name", "height": UNKNOWN_ENTITY} + predicted = {"name": "name", "height": None} + + assert _column_accuracy(predicted, expected) == 1.0 + + +def test_column_accuracy_metric_rejects_missing_predictions(): + expected = {"name": "name", "height": UNKNOWN_ENTITY} + predicted = {"height": None} + + with pytest.raises(AssertionError, match="name"): + _column_accuracy(predicted, expected) + + +def test_positive_recall_ignores_negative_columns(): + expected = { + "name": "name", + "email": "email", + "height": UNKNOWN_ENTITY, + } + predicted = { + "name": "name", + "email": UNKNOWN_ENTITY, + "height": UNKNOWN_ENTITY, + } + + assert _positive_recall(predicted, expected) == 0.5 + + +def test_classification_benchmark_matrix_defaults(monkeypatch): + monkeypatch.delenv("NSS_CLASSIFICATION_BENCHMARK_BACKENDS", raising=False) + monkeypatch.delenv("NSS_CLASSIFICATION_BENCHMARK_API_MODELS", raising=False) + monkeypatch.delenv("NSS_CLASSIFICATION_BENCHMARK_LOCAL_MODELS", raising=False) + + assert _benchmark_matrix() == [ + ("api", "default", DefaultLLMConfig.DEFAULT_CONFIG_ID), + ("local_hf", "smollm3", "HuggingFaceTB/SmolLM3-3B"), + ("local_hf", "mistral", "mistralai/Mistral-7B-Instruct-v0.3"), + ("local_hf", "tinyllama", "TinyLlama/TinyLlama-1.1B-Chat-v1.0"), + ] + + +def test_benchmark_summary_report_upserts_and_renders_table(tmp_path, monkeypatch): + monkeypatch.setenv("NSS_CLASSIFICATION_BENCHMARK_OUTPUT", str(tmp_path / "classification-benchmark.json")) + result = { + "backend": "local_hf", + "model_name": "smollm3", + "model": LOCAL_HF_CLASSIFICATION_BENCHMARK_MODELS["smollm3"], + "status": "ok", + "macro_accuracy": 0.5, + "macro_positive_recall": 0.25, + "cases": [ + {"case_id": "pii_dataset_structured", "accuracy": 1.0}, + {"case_id": "credit_card_transactions", "accuracy": 0.5}, + {"case_id": "adult_negative_control", "accuracy": 0.0}, + ], + } + + _write_benchmark_reports(result) + result["macro_accuracy"] = 0.75 + _write_benchmark_reports(result) + + summary_json = json.loads((tmp_path / "classification-benchmark-summary.json").read_text()) + assert len(summary_json["results"]) == 1 + assert summary_json["results"][0]["macro_accuracy"] == 0.75 + summary_md = (tmp_path / "classification-benchmark-summary.md").read_text() + assert ( + "| Backend | Model | Macro Accuracy | Macro Positive Recall | pii_dataset_structured | credit_card_transactions | adult_negative_control |" + in summary_md + ) + assert "| local_hf | smollm3 | 0.750 | 0.250 | 1.000 | 0.500 | 0.000 |" in summary_md + + +@pytest.mark.slow +@pytest.mark.parametrize(("backend", "model_name", "model"), _all_benchmark_params()) +def test_live_column_classification_model_comparison(pii_test_data_dir, backend, model_name, model): + """Optional benchmark for one API/local classifier and model-family pair. + + Set ``NSS_RUN_CLASSIFICATION_BENCHMARK=1`` to run. By default this evaluates + the production API classifier baseline plus local HF SmolLM3, Mistral, and + TinyLlama as separate pytest cases. Narrow the + run with comma-separated ``NSS_CLASSIFICATION_BENCHMARK_BACKENDS``, + ``NSS_CLASSIFICATION_BENCHMARK_API_MODELS``, and + ``NSS_CLASSIFICATION_BENCHMARK_LOCAL_MODELS``. Each case writes a JSON report + named ``column-classification-benchmark-{backend}-{model}.json`` by default. + The test takes a file lock so benchmark cases run serially under xdist. + """ + if os.environ.get("NSS_RUN_CLASSIFICATION_BENCHMARK") != "1": + pytest.skip("Set NSS_RUN_CLASSIFICATION_BENCHMARK=1 to run live column classification benchmark.") + if (backend, model_name, model) not in _benchmark_matrix(): + pytest.skip("Benchmark case excluded by NSS_CLASSIFICATION_BENCHMARK_* filters.") + + with _classification_benchmark_lock(): + cases = _load_gold(pii_test_data_dir / "column_classification_gold.json") + result = _evaluate_classifier(backend, model_name, model, cases) + + output_path, summary_json_path, summary_md_path = _write_benchmark_reports(result) + + print(f"Column classification benchmark report written to {output_path}") + print(f"Column classification benchmark summary written to {summary_json_path}") + print(f"Column classification benchmark table written to {summary_md_path}") + assert len(result["cases"]) == len(cases) diff --git a/tests/pii_replacer/test_detect.py b/tests/pii_replacer/test_detect.py index 18bbbc091..89c8e90ae 100644 --- a/tests/pii_replacer/test_detect.py +++ b/tests/pii_replacer/test_detect.py @@ -1,11 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402 import pytest # Skip all tests in this module if torch is not available -pytest.importorskip("torch", reason="torch is required for these tests (install with: uv sync --extra cpu)") +torch = pytest.importorskip("torch", reason="torch is required for these tests (install with: uv sync --extra cpu)") +from typing import TYPE_CHECKING, cast from unittest.mock import MagicMock, patch import numpy as np @@ -17,6 +19,7 @@ DEFAULT_ENTITIES, UNKNOWN_ENTITY, ClassifyConfig, + ColumnClassifierHF, ColumnClassifierLLM, DefaultLLMConfig, EntityExtractorGliner, @@ -28,6 +31,9 @@ from nemo_safe_synthesizer.pii_replacer.data_editor.environment import redact_entities_fn from nemo_safe_synthesizer.pii_replacer.ner.ner import NERPrediction +if TYPE_CHECKING: + from transformers import PreTrainedTokenizerBase + class TestDefaultLLMConfigId: def test_uses_env_override(self, monkeypatch): @@ -378,6 +384,117 @@ def attach_mock_response(content: str): assert "LLM failed" in exc.value.args[0] +def test_detect_types_local_hf_uses_shared_prompt_and_parser(): + classifier = ColumnClassifierHF(model_name_or_path="local/smollm", num_samples=1) + df = pd.DataFrame( + { + "email": ["alice@example.com", "bob@example.com"], + "height": [170, 180], + } + ) + + class FakeTokenizer: + chat_template = "{{ messages }}" + pad_token_id = 0 + eos_token_id = 1 + + def apply_chat_template(self, messages, *, add_generation_prompt, return_tensors): + assert add_generation_prompt is True + assert return_tensors == "pt" + assert messages[0]["content"] == DefaultLLMConfig.SYSTEM_PROMPT + assert "email:" in messages[1]["content"] + return { + "input_ids": torch.tensor([[10, 11]]), + "attention_mask": torch.tensor([[1, 1]]), + } + + def decode(self, generated_ids, *, skip_special_tokens): + assert skip_special_tokens is True + return '{"email": "email", "height": "none"}' + + class FakeModel: + device = torch.device("cpu") + + def generate(self, *, input_ids, attention_mask, max_new_tokens, do_sample, pad_token_id, eos_token_id): + assert input_ids.tolist() == [[10, 11]] + assert attention_mask.tolist() == [[1, 1]] + assert max_new_tokens == DefaultLLMConfig.MAX_OUTPUT_TOKENS + assert do_sample is False + assert pad_token_id == 0 + assert eos_token_id == 1 + return torch.tensor([[10, 11, 12, 13]]) + + classifier._tokenizer = cast("PreTrainedTokenizerBase", FakeTokenizer()) + classifier._model = FakeModel() + + assert classifier.detect_types(df, DEFAULT_ENTITIES) == { + "email": "email", + "height": "none", + } + + +def test_detect_types_local_hf_falls_back_without_chat_template(): + classifier = ColumnClassifierHF(model_name_or_path="local/smollm", num_samples=1) + df = pd.DataFrame({"email": ["alice@example.com"]}) + + class FakeTokenizer: + chat_template = None + pad_token_id = 0 + eos_token_id = 1 + + def __call__(self, prompt, *, return_tensors): + assert return_tensors == "pt" + assert prompt.startswith(DefaultLLMConfig.SYSTEM_PROMPT) + assert "email:" in prompt + return { + "input_ids": torch.tensor([[20, 21]]), + "attention_mask": torch.tensor([[1, 1]]), + } + + def decode(self, generated_ids, *, skip_special_tokens): + assert skip_special_tokens is True + return '{"email": "email"}' + + class FakeModel: + device = torch.device("cpu") + + def generate(self, *, input_ids, attention_mask, max_new_tokens, do_sample, pad_token_id, eos_token_id): + assert input_ids.tolist() == [[20, 21]] + return torch.tensor([[20, 21, 22]]) + + classifier._tokenizer = cast("PreTrainedTokenizerBase", FakeTokenizer()) + classifier._model = FakeModel() + + assert classifier.detect_types(df, DEFAULT_ENTITIES) == {"email": "email"} + + +def test_local_hf_uses_repo_id_for_incomplete_online_cache(tmp_path, monkeypatch): + snapshot = tmp_path / "models--org--model" / "snapshots" / "abc" + snapshot.mkdir(parents=True) + model_ref = MagicMock( + repo_id="org/model", + local_path=snapshot, + target=MagicMock(return_value=snapshot), + ) + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + + assert ColumnClassifierHF._load_target(model_ref) == "org/model" + + +def test_local_hf_keeps_incomplete_cache_path_when_offline(tmp_path, monkeypatch): + snapshot = tmp_path / "models--org--model" / "snapshots" / "abc" + snapshot.mkdir(parents=True) + model_ref = MagicMock( + repo_id="org/model", + local_path=snapshot, + target=MagicMock(return_value=snapshot), + ) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + + assert ColumnClassifierHF._load_target(model_ref) == snapshot + + def test_redact_from_entities(): entities = [ NERPrediction("", start=3, end=5, label="name", source="test", score=9.0), diff --git a/tests/pii_replacer/test_nemo_pii.py b/tests/pii_replacer/test_nemo_pii.py index ca63ef193..400bf380b 100644 --- a/tests/pii_replacer/test_nemo_pii.py +++ b/tests/pii_replacer/test_nemo_pii.py @@ -10,13 +10,16 @@ import pandas as pd +from nemo_safe_synthesizer.config.replace_pii import PiiReplacerConfig from nemo_safe_synthesizer.defaults import DEFAULT_NSS_INFERENCE_ENDPOINT +from nemo_safe_synthesizer.pii_replacer.data_editor.detect import ColumnClassifierHF, ColumnClassifierLLM from nemo_safe_synthesizer.pii_replacer.data_editor.edit import TransformFnAccounting from nemo_safe_synthesizer.pii_replacer.nemo_pii import ( ColumnClassification, NemoPII, _build_column_statistics, _get_classify_endpoint_url, + get_column_classifier, ) @@ -33,6 +36,44 @@ def test_unset_falls_back_to_default(self, monkeypatch): monkeypatch.delenv("NSS_INFERENCE_ENDPOINT", raising=False) assert _get_classify_endpoint_url() == DEFAULT_NSS_INFERENCE_ENDPOINT + +class TestGetColumnClassifier: + @patch("nemo_safe_synthesizer.pii_replacer.nemo_pii.OpenAI") + def test_api_backend_uses_configured_num_samples(self, mock_openai): + config = PiiReplacerConfig.get_default_config() + config.globals.classify.num_samples = 2 + + classifier = get_column_classifier(config) + + assert isinstance(classifier, ColumnClassifierLLM) + assert classifier._num_samples == 2 + mock_openai.assert_called_once() + + @patch("nemo_safe_synthesizer.pii_replacer.nemo_pii.OpenAI") + def test_api_backend_close_closes_openai_client(self, mock_openai): + config = PiiReplacerConfig.get_default_config() + client = MagicMock() + mock_openai.return_value = client + + classifier = get_column_classifier(config) + classifier.close() + + client.close.assert_called_once() + + @patch("nemo_safe_synthesizer.pii_replacer.nemo_pii.OpenAI") + def test_local_hf_backend_does_not_create_openai_client(self, mock_openai): + config = PiiReplacerConfig.get_default_config() + config.globals.classify.backend = "local_hf" + config.globals.classify.model = "local/smollm" + config.globals.classify.num_samples = 4 + + classifier = get_column_classifier(config) + + assert isinstance(classifier, ColumnClassifierHF) + assert classifier._model_name_or_path == "local/smollm" + assert classifier._num_samples == 4 + mock_openai.assert_not_called() + @pytest.mark.parametrize("blank", ["", " ", "\t"]) def test_blank_falls_back_to_default(self, monkeypatch, blank): # A blank endpoint must resolve to the default, never reach the OpenAI @@ -94,6 +135,7 @@ def test_nemo_pii_classify_df(_build_entity_extractor, fake_people_csv): "date of birth": "date", "notes": "text", } + mock_column_classifier.close.assert_called_once() @patch("nemo_safe_synthesizer.pii_replacer.nemo_pii.build_entity_extractor", return_value=MagicMock()) diff --git a/tests/preflight/test_preflight.py b/tests/preflight/test_preflight.py index 143157c7e..6a26e59f7 100644 --- a/tests/preflight/test_preflight.py +++ b/tests/preflight/test_preflight.py @@ -19,7 +19,7 @@ from nemo_safe_synthesizer.config.parameters import SafeSynthesizerParameters from nemo_safe_synthesizer.config.time_series import TimeSeriesParameters from nemo_safe_synthesizer.config.training import TrainingHyperparams -from nemo_safe_synthesizer.defaults import DEFAULT_MAX_SEQ_LENGTH, PSEUDO_GROUP_COLUMN +from nemo_safe_synthesizer.defaults import DEFAULT_MAX_SEQ_LENGTH, DEFAULT_PII_CLASSIFY_LOCAL_MODEL, PSEUDO_GROUP_COLUMN from nemo_safe_synthesizer.llm.metadata import ModelMetadata from nemo_safe_synthesizer.llm.utils import ModelRef from nemo_safe_synthesizer.preflight import ( @@ -479,6 +479,67 @@ def test_missing_key_takes_priority_over_blank_model(self, default_config): codes = {i.code for i in issues} assert codes == {"inference_key_missing"} + def test_local_hf_backend_skips_api_env_warnings(self, default_config): + default_config.replace_pii.globals.classify.backend = "local_hf" + model_ref = MagicMock( + repo_id="HuggingFaceTB/SmolLM3-3B", + local_path=None, + cache_root=Path("/hf-cache"), + ) + model_ref.partial_cached_snapshot.return_value = None + + with ( + patch("nemo_safe_synthesizer.preflight.checks.environment.ModelRef.parse", return_value=model_ref), + patch("nemo_safe_synthesizer.preflight.checks.environment.hf_offline_enabled", return_value=False), + patch.dict("os.environ", {}, clear=True), + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + + codes = {i.code for i in issues} + assert "inference_key_missing" not in codes + assert "classify_hf_model_not_cached" in codes + + def test_local_hf_backend_uses_default_smollm3_model(self, default_config): + default_config.replace_pii.globals.classify.backend = "local_hf" + model_ref = MagicMock( + repo_id="HuggingFaceTB/SmolLM3-3B", + local_path=None, + cache_root=Path("/hf-cache"), + ) + model_ref.partial_cached_snapshot.return_value = None + + with ( + patch("nemo_safe_synthesizer.preflight.checks.environment.ModelRef.parse", return_value=model_ref) as parse, + patch("nemo_safe_synthesizer.preflight.checks.environment.hf_offline_enabled", return_value=False), + patch.dict("os.environ", {"HF_TOKEN": "hf_xxx"}, clear=True), + ): + InferenceModelCheck().run(make_ctx(config=default_config)) + + parse.assert_called_once_with(DEFAULT_PII_CLASSIFY_LOCAL_MODEL) + + def test_local_hf_backend_missing_local_path_errors(self, default_config): + default_config.replace_pii.globals.classify.backend = "local_hf" + default_config.replace_pii.globals.classify.model = "/missing/local/model" + + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + + assert any(i.code == "classify_local_model_missing" and i.severity == "error" for i in issues) + + def test_local_hf_backend_incomplete_cache_warns(self, default_config, hf_cached_snapshot_factory, monkeypatch): + cache_root, _ = hf_cached_snapshot_factory(files=("config.json", "tokenizer.json")) + default_config.replace_pii.globals.classify.backend = "local_hf" + default_config.replace_pii.globals.classify.model = "nvidia/Nemotron-Mini-4B-Instruct" + monkeypatch.setattr(ModelRef, "_default_hf_cache_root", staticmethod(lambda: cache_root)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + monkeypatch.delenv("HF_TOKEN", raising=False) + monkeypatch.delenv("HUGGING_FACE_HUB_TOKEN", raising=False) + + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + + assert any(i.code == "classify_hf_model_cache_incomplete" and i.severity == "warning" for i in issues) + assert any(i.code == "hf_token_missing" and i.check == "env.inference" for i in issues) + @pytest.mark.unit class TestHFModelAvailabilityCheck: diff --git a/tests/test_data/pii/column_classification_gold.json b/tests/test_data/pii/column_classification_gold.json new file mode 100644 index 000000000..11871ce63 --- /dev/null +++ b/tests/test_data/pii/column_classification_gold.json @@ -0,0 +1,133 @@ +{ + "cases": [ + { + "id": "pii_dataset_structured", + "source": "cleaned/pii_dataset.csv", + "rows": [ + { + "name": "Aaliyah Popova", + "email": "aaliyah.popova4783@aol.edu", + "phone": "(95) 94215-7906", + "job": "jeweler", + "address": "97 Lincoln Street", + "text": "Contact Aaliyah Popova by phone at (95) 94215-7906 or by email at aaliyah.popova4783@aol.edu.", + "hobby": "Podcasting" + }, + { + "name": "Daniel Martinez", + "email": "daniel.martinez@example.com", + "phone": "323-555-6724", + "job": "physician", + "address": "912 Cedar Avenue", + "text": "Daniel Martinez lives at 912 Cedar Avenue and can be reached at 323-555-6724.", + "hobby": "cycling" + } + ], + "expected_entities": { + "name": "name", + "email": "email", + "phone": "phone_number", + "job": "none", + "address": "address", + "text": "none", + "hobby": "none" + } + }, + { + "id": "credit_card_transactions", + "source": "cleaned/user0_credit_card_transactions.csv", + "rows": [ + { + "User": 0, + "Card": 0, + "Year": 2002, + "Month": 9, + "Day": 1, + "Time": "06:21", + "Amount": "$134.09", + "Use Chip": "Swipe Transaction", + "Merchant Name": "3527213246127876953", + "Merchant City": "La Verne", + "Merchant State": "CA", + "Zip": "91750.0", + "MCC": 5300, + "Is Fraud?": "No" + }, + { + "User": 0, + "Card": 0, + "Year": 2002, + "Month": 9, + "Day": 1, + "Time": "06:42", + "Amount": "$38.48", + "Use Chip": "Swipe Transaction", + "Merchant Name": "-727612092139916043", + "Merchant City": "Monterey Park", + "Merchant State": "CA", + "Zip": "91754.0", + "MCC": 5411, + "Is Fraud?": "No" + } + ], + "expected_entities": { + "User": "none", + "Card": "credit_debit_card", + "Year": "none", + "Month": "none", + "Day": "none", + "Time": "none", + "Amount": "none", + "Use Chip": "none", + "Merchant Name": "none", + "Merchant City": "city", + "Merchant State": "state", + "Zip": "postcode", + "MCC": "none", + "Is Fraud?": "none" + } + }, + { + "id": "adult_negative_control", + "source": "cleaned/adult.csv", + "rows": [ + { + "age": 39, + "workclass": "State-gov", + "education": "Bachelors", + "marital.status": "Never-married", + "occupation": "Adm-clerical", + "relationship": "Not-in-family", + "race": "White", + "sex": "Male", + "native.country": "United-States", + "income": "<=50K" + }, + { + "age": 50, + "workclass": "Self-emp-not-inc", + "education": "Bachelors", + "marital.status": "Married-civ-spouse", + "occupation": "Exec-managerial", + "relationship": "Husband", + "race": "White", + "sex": "Male", + "native.country": "United-States", + "income": "<=50K" + } + ], + "expected_entities": { + "age": "none", + "workclass": "none", + "education": "none", + "marital.status": "none", + "occupation": "none", + "relationship": "none", + "race": "none", + "sex": "none", + "native.country": "none", + "income": "none" + } + } + ] +}