-
Notifications
You must be signed in to change notification settings - Fork 8
feat: add option to use local classification #615
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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( | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
| 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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When |
||
|
|
||
| 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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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/configuratorRepository: 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
|
||
|
|
||
| class Globals(NSSBaseModel): | ||
| """Global settings for the PII replacer including locales, seed, NER, and classification.""" | ||
|
|
||
| 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.""" |
There was a problem hiding this comment.
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
--validatefor network/gated-model availability that it has not proven.Suggested wording
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
Source: Path instructions