Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
4 changes: 2 additions & 2 deletions docs/tutorials/safe-synthesizer-101.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"display_name": ".venv",
"language": "python",
"name": "python3"
},
Expand All @@ -231,7 +231,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.9"
"version": "3.13.12"
}
},
"nbformat": 4,
Expand Down
4 changes: 3 additions & 1 deletion docs/user-guide/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 |

Expand Down
4 changes: 2 additions & 2 deletions docs/user-guide/docker.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`),
Expand Down
20 changes: 15 additions & 5 deletions docs/user-guide/environment.md
Original file line number Diff line number Diff line change
Expand Up @@ -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) |
Expand Down Expand Up @@ -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"
Expand All @@ -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`

Expand Down
6 changes: 3 additions & 3 deletions docs/user-guide/evaluating-data.md
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down
35 changes: 25 additions & 10 deletions docs/user-guide/running.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 |

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Avoid saying preflight proves the model is fetchable.

The local HF check does not contact Hugging Face; it verifies local/cache state and warns when runtime may need an online fetch. Saying the reference is “fetchable” can make users trust --validate for network/gated-model availability that it has not proven.

Suggested wording
-| `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.inference` | config | Inference config for PII classification: API backend env vars are usable, or the `local_hf` classifier model reference resolves locally; warns when runtime may need Hugging Face access |

As per path instructions, “Review documentation as MkDocs Material content. Check Diataxis fit, accurate commands, internal links, and markdown style from STYLE_GUIDE.md.”

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
| `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.inference` | config | Inference config for PII classification: API backend env vars are usable, or the `local_hf` classifier model reference resolves locally; warns when runtime may need Hugging Face access |

Source: Path instructions

| `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 |
Expand Down Expand Up @@ -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.

Expand Down
11 changes: 9 additions & 2 deletions docs/user-guide/troubleshooting.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 |
Expand Down
25 changes: 24 additions & 1 deletion src/nemo_safe_synthesizer/config/replace_pii.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

agent-assisted: Could we document these new settings? The current guides only cover API classification and still say an inference key is required. It would help to include a local_hf example and mention the default 3B model download, offline/cache behavior, and local model paths.

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.",
)
Comment on lines +192 to +195

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 model field is silently ignored for the api backend

When backend="api", get_column_classifier never reads classify_config.model; the API model comes from the NSS_INFERENCE_MODEL environment variable. A user who sets model="meta-llama/Llama-3-70b-instruct" while keeping the default backend="api" will see their value silently discarded. A model_validator that emits a warning (or raises a ParameterError) when model is non-None and backend != "local_hf" would prevent silent misconfiguration.


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

Comment on lines +187 to +214

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check how ClassifyConfig / PiiReplacerConfig overrides are merged (model_copy vs re-validation).
rg -n -B2 -A8 'model_copy\(update' src/nemo_safe_synthesizer/config src/nemo_safe_synthesizer/configurator

Repository: NVIDIA-NeMo/Safe-Synthesizer

Length of output: 5611


🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

# Inspect the relevant config model and its construction/override paths.
sed -n '1,280p' src/nemo_safe_synthesizer/config/replace_pii.py
printf '\n---\n'
sed -n '1,260p' src/nemo_safe_synthesizer/configurator/parameters.py
printf '\n---\n'
rg -n "replace_pii|classify_model_provider|backend.*local_hf|local_hf" src/nemo_safe_synthesizer -g'*.py'

Repository: NVIDIA-NeMo/Safe-Synthesizer

Length of output: 24158


warn_api_model_ignored needs to run after overrides are merged. Parameters.from_yaml() / from_json() apply overrides with model_copy(update=...), which skips after validators, so replace_pii.globals.classify.model can be set on an api config without this warning ever showing. Re-validate after merging or move the warning into the override path.


class Globals(NSSBaseModel):
"""Global settings for the PII replacer including locales, seed, NER, and classification."""
Expand Down
1 change: 1 addition & 0 deletions src/nemo_safe_synthesizer/defaults.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 39 additions & 0 deletions src/nemo_safe_synthesizer/llm/model_host.py
Original file line number Diff line number Diff line change
@@ -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."""
Loading
Loading