feat: add option to use local classification

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,

src/nemo_safe_synthesizer/config/replace_pii.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import os
-from typing import Annotated, Any, Self
+from typing import Annotated, Any, Literal, Self
 
 from faker.config import AVAILABLE_LOCALES
 from pydantic import Field, field_validator, model_validator
@@ -183,6 +183,16 @@ 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. "

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

src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py

@@ -11,7 +11,7 @@
 from itertools import chain, islice
 from time import monotonic
 from timeit import default_timer as timer
-from typing import Optional
+from typing import Any, Optional
 
 import json_repair
 import pandas as pd
@@ -20,6 +20,8 @@
 from openai import OpenAI
 from pydantic import ConfigDict, TypeAdapter, ValidationError
 
+from ...defaults import DEFAULT_PII_CLASSIFY_LOCAL_MODEL
+from ...llm.utils import ModelRef, cleanup_memory
 from ...observability import get_logger
 from ...utils import hf_offline_enabled
 from ..ner import ner_mp
@@ -49,6 +51,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:
@@ -281,6 +284,15 @@ def classify_columns(
         },
     )
 
+    return _filter_entities(entities_str, entities, on_validation_error)
+
+
+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()}
 
@@ -349,6 +361,9 @@ def detect_types(self, df: pd.DataFrame, entities: Optional[set[str]]) -> dict[s
         """
         ...
 
+    def close(self) -> None:
+        """Release backend resources after classification."""
+
 
 class ColumnClassifierNoop(ColumnClassifier):
     """No-op classifier that assigns ``UNKNOWN_ENTITY`` to every column."""
@@ -408,6 +423,128 @@ def _on_validation_error(self) -> None:
         )
 
 
+class ColumnClassifierHF(ColumnClassifier):
+    """Classify column types with an in-process Hugging Face causal LM."""
+
+    _model_name_or_path: str
+    _num_samples: Optional[int]
+    _model: Any | None
+    _tokenizer: Any | 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
+
+    def detect_types(self, df: pd.DataFrame, entities: set[str]) -> dict[str, Optional[str]]:
+        """Sample columns, run local text generation, and parse JSON entity labels."""
+        formatted_prompt = _format_prompt(df, entities, self._num_samples)
+        if not formatted_prompt:
+            return {}
+
+        self._load()
+        if self._model is None or self._tokenizer is None:
+            raise RuntimeError("Local Hugging Face classifier failed to initialize.")
+
+        messages = [
+            {"role": "system", "content": DefaultLLMConfig.SYSTEM_PROMPT},
+            {"role": "user", "content": formatted_prompt},
+        ]
+        tokenizer = self._tokenizer
+        model = self._model
+
+        if hasattr(tokenizer, "apply_chat_template"):
+            encoded = tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+                return_tensors="pt",
+            )
+        else:
+            prompt = f"{DefaultLLMConfig.SYSTEM_PROMPT}\n\n{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)
+
+        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 _filter_entities(entities_str, entities, self._on_validation_error)
+
+    def _load(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_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(
+            model_ref.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(model_ref.target(), **load_kwargs)
+        if not torch.cuda.is_available():
+            self._model = self._model.to("cpu")
+        self._model.eval()
+
+    def close(self) -> None:
+        self._model = None
+        self._tokenizer = None
+        cleanup_memory()
+
+    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)."""

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: