feat: add option to use local classification#615
Conversation
Signed-off-by: mkornfield <mkornfield@nvidia.com>
github-actions
Bot
added
area:config
area:data-processing
area:pii
area:tests
area:docs
labels
Greptile SummaryThis PR adds a
Confidence Score: 4/5Safe to merge after addressing the two open issues from prior review rounds (wrong kwarg dtype vs torch_dtype in AutoModelForCausalLM and the unreachable entity labels in the gold fixture); the new code paths are well-tested and the resource lifecycle is correct. The architecture is solid — the template-method refactor is clean, the close() lifecycle is correctly wired into the finally block, and the preflight extension follows the existing pattern faithfully. The two previously-identified issues (wrong torch_dtype kwarg causing bfloat16 to be silently ignored on CUDA, and credit_debit_card/state fixture labels that are not in DEFAULT_ENTITIES) remain unresolved and will cause real problems: GPU memory waste for larger models and a failing schema assertion when torch is available. src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py (ColumnClassifierHF.initialize — dtype kwarg) and tests/test_data/pii/column_classification_gold.json (credit_debit_card / state labels vs DEFAULT_ENTITIES). Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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.", | ||
| ) |
There was a problem hiding this comment.
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.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
kendrickb-nvidia
left a comment
kendrickb-nvidia
left a comment
There was a problem hiding this comment.
I'm not opposed to adding this now to existing PII replacement, but we're also going to be redoing PII replacement in July, so could be easier to just support local mode then? But this is already written.
CC @nina-xu on the column_classification_gold.json (not sure where @mckornfield sourced this from) and using this or similar in evaluation of new pii replacement.
|
will DM about it, mostly was just trying to pick up old, interesting issues. This one seemed like a decent one, but I know the PII replace is going to go elsewhere than this current state haha |
binaryaaron
left a comment
binaryaaron
left a comment
There was a problem hiding this comment.
agent-assisted: Review of the local Hugging Face classification path.
| return | ||
| collector.warning( | ||
| "classify_hf_model_cache_incomplete", | ||
| f"{message} Model loading will contact Hugging Face unless the full model snapshot is pre-downloaded.", |
There was a problem hiding this comment.
agent-assisted: This says the loader can fetch the missing files, but _load() may pass the incomplete cache directory instead of the Hugging Face model ID. Transformers treats that directory as local, so it cannot fill in missing config, tokenizer, or shard files. Could we use the model ID when online, or make an incomplete cache an error?
| tokenizer = self._tokenizer | ||
| model = self._model | ||
|
|
||
| if hasattr(tokenizer, "apply_chat_template"): |
There was a problem hiding this comment.
agent-assisted: hasattr() is true for standard Transformers tokenizers even when they do not have a chat template. apply_chat_template() then raises, so the fallback below never runs. Could we check for an actual template or catch that specific error and use the plain prompt?
|
|
||
| num_samples: int | None = Field(description="Number of column values to sample for classification.", default=3) | ||
|
|
||
| backend: Literal["api", "local_hf"] = Field( |
There was a problem hiding this comment.
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.
| def _check_column_classifier_hf_model(model_name: str, collector: IssueCollector) -> None: | ||
| """Validate the Hugging Face model used by local PII column classification.""" | ||
| if not model_name: | ||
| collector.error("classify_model_ref_empty", "`replace_pii.globals.classify.model` must not be empty.") |
There was a problem hiding this comment.
agent-assisted: These new classify_* codes are missing from the troubleshooting table. This path also emits shared HF codes under env.inference, while the docs list them under env.hf_model_availability. Could we update the table so the documented check names match what users actually receive?
| def _column_accuracy(predicted: dict[str, str | None], expected: dict[str, str]) -> float: | ||
| if not expected: | ||
| return 1.0 | ||
| correct = sum(_normalize_label(predicted.get(column)) == label for column, label in expected.items()) |
There was a problem hiding this comment.
agent-assisted: A missing prediction is counted as none here. With the current gold data, an empty result gets 0.714 accuracy even though it finds none of the PII columns. Could we reject missing columns and add a positive-recall check to the benchmark?
| ) | ||
|
|
||
|
|
||
| def _check_column_classifier_hf_model(model_name: str, collector: IssueCollector) -> None: |
There was a problem hiding this comment.
agent-assisted: This repeats most of the model availability logic from HFModelAvailabilityCheck below and calls some of its private helpers. Could we pull the shared cache, offline, token, and remote-code checks into one helper and use it from both places?
binaryaaron
left a comment
binaryaaron
left a comment
There was a problem hiding this comment.
agent (
review-pr): Requesting changes to the local classification backend
Thanks for putting this together. The local option is useful, but I think the backend boundary needs another pass before this merges. We are working through a similar local/remote split in #575: the base owns the shared flow, while each backend owns its typed setup, model call, and idempotent cleanup.
Using the same shape here would address the resource ownership and typing concerns below without making this PR depend on #575. The existing inline comments also identify correctness gaps in cache loading, chat-template fallback, and benchmark scoring that should be resolved before merge.
| or self.pii_replacer_config.globals.classify.backend == "local_hf", | ||
| ) | ||
| finally: | ||
| column_classifier.close() |
There was a problem hiding this comment.
agent (
review-pr): Close the API client during cleanup
One cleanup issue here: get_column_classifier() creates an OpenAI client, but ColumnClassifierLLM inherits the no-op close() method. This finally block therefore releases the local HF model but leaves the API client's HTTP transport open. Can we have the API backend close its client too? This is the same lifecycle we are using for RemoteBackend in #575.
| ) | ||
|
|
||
|
|
||
| class ColumnClassifierHF(ColumnClassifier): |
There was a problem hiding this comment.
agent (
review-pr): Follow the backend split from #575
This looks similar to the local and remote backend split we are working through in #575. There, the base class owns the shared generation flow, while each backend handles its own model call, setup, and cleanup.
I think the same split would work well here. ColumnClassifier.detect_types() could handle sampling, prompt construction, parsing, and filtering. The API and Hugging Face implementations would only need to provide the model call and their resource lifecycle. That would also give _model and _tokenizer concrete types instead of Any.
This does not need to depend on #575. Keeping the two designs close would just make them easier to bring together later.
|
Following up on the ownership discussion here: I rebased #575 onto I think that gives this PR a cleaner shape for |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📜 Recent review details⏰ Context from checks skipped due to timeout. (4)
WalkthroughAdds a ChangesLocal HF classification backend
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
| if torch.cuda.is_available(): | ||
| load_kwargs["device_map"] = "auto" | ||
| load_kwargs["dtype"] = torch.bfloat16 |
There was a problem hiding this comment.
AutoModelForCausalLM.from_pretrained accepts torch_dtype, not dtype. Passing dtype=torch.bfloat16 lands in **kwargs, is forwarded as an unrecognized argument, and is silently ignored, so the model loads in float32 on CUDA hardware. This doubles GPU memory consumption and will cause OOM for larger models.
| if torch.cuda.is_available(): | |
| load_kwargs["device_map"] = "auto" | |
| load_kwargs["dtype"] = torch.bfloat16 | |
| if torch.cuda.is_available(): | |
| load_kwargs["device_map"] = "auto" | |
| load_kwargs["torch_dtype"] = torch.bfloat16 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py (1)
445-459: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winDuplicate init-guard raises two different exception types for the same condition.
detect_types(line 447-448, pre-existing) raises bareExceptionwhenself._llm is None; the new_classify_prompt(line 453-454) raisesRuntimeErrorfor the identical precondition. Sincedetect_typescallssuper().detect_types()which callsself._classify_prompt(), the check is now duplicated with inconsistent exception types, and neither uses the project's error hierarchy (InternalError/UserError) required for user-facing failures. Callers catching a specific exception type will behave differently depending on which code path raises first.🐛 Proposed fix
def detect_types(self, df: pd.DataFrame, entities: set[str]) -> dict[str, Optional[str]]: """Sample column data and call the inference API to classify columns into entity types.""" - if self._llm is None: - raise Exception("InferenceAPI classifier not initialized. Use get_classifier() method.") - 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.") + raise InternalError("InferenceAPI classifier not initialized. Use get_classifier() method.") return _classify_prompt_with_openai(As per path instructions, "user-facing failures should raise project errors from
nemo_safe_synthesizer.errorswith precise messages."Source: Path instructions
🧹 Nitpick comments (4)
src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py (3)
479-647: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing docstrings on nontrivial
ColumnClassifierHFmethods.
initialize()(cache/offline/device handling),_encode_prompt()(chat-template vs. plain-prompt fallback), andteardown()have non-obvious logic but no docstrings, and the class docstring omitsArgs:for the constructor parameters (model_name_or_path,num_samples). Per STYLE_GUIDE.md, a docstring is mandatory for functions with non-obvious logic, and constructorArgs:belong in the class docstring so IDEs surface them on hover.Source: Path instructions
608-611: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
_loadalias adds an unnecessary second entry point.This is net-new code in this PR (no external
ColumnClassifierHFconsumers predate it), so the "backward-compatible alias for older tests and callers" is really just carrying an earlier draft's method name forward. Consider updating the referencing tests to callinitialize()directly and dropping this alias to avoid two names for the same operation.
591-593: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winHard-coded
torch.bfloat16may not suit all GPUs the local backend targets.
device_map="auto"+dtype=torch.bfloat16is applied whenever any CUDA device is available, regardless of GPU architecture. Pre-Ampere GPUs (compute capability < 8.0) don't have native bf16 support and can hit degraded performance or backend errors depending on the op. Since this backend is meant to lower the adoption barrier (per issue#121) it will likely run on a wider variety of hardware than the training path. Consider gating ontorch.cuda.is_bf16_supported()and falling back totorch.float16/"auto"otherwise.tests/pii_replacer/test_detect.py (1)
467-492: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLGTM on core coverage; consider adding a "complete cache" branch test.
Both
_load_targettests exercise the incomplete-cache branch (online → repo id, offline → snapshot path). The complementary branch — a complete cache returningmodel_ref.target()— isn't covered here, though it may be covered elsewhere.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ae0c23be-f2f2-4a89-83de-2dc73cb3f5b7
📒 Files selected for processing (19)
README.mddocs/tutorials/safe-synthesizer-101.ipynbdocs/user-guide/configuration.mddocs/user-guide/docker.mddocs/user-guide/environment.mddocs/user-guide/evaluating-data.mddocs/user-guide/running.mddocs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/config/replace_pii.pysrc/nemo_safe_synthesizer/defaults.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/config/test_nss_config.pytests/pii_replacer/test_column_classification_eval.pytests/pii_replacer/test_detect.pytests/pii_replacer/test_nemo_pii.pytests/preflight/test_preflight.pytests/test_data/pii/column_classification_gold.json
📜 Review details
⏰ Context from checks skipped due to timeout. (4)
- GitHub Check: Unit Tests (3.13)
- GitHub Check: Unit Tests (3.11)
- GitHub Check: Unit Tests (3.12)
- GitHub Check: Greptile Review
🧰 Additional context used
📓 Path-based instructions (18)
**/*.{md,markdown,py}
📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)
**/*.{md,markdown,py}: Avoid decorative bold (**text**) in list items, body text, and docstrings; use structural cues (headers, list markers, colons, backticks) for emphasis instead
Use backticks for code identifiers, paths, and CLI commands in markdown and docstrings
Files:
src/nemo_safe_synthesizer/defaults.pytests/config/test_nss_config.pytests/preflight/test_preflight.pydocs/user-guide/troubleshooting.mddocs/user-guide/environment.mdtests/pii_replacer/test_nemo_pii.pyREADME.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mdsrc/nemo_safe_synthesizer/config/replace_pii.pytests/pii_replacer/test_detect.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/pii_replacer/test_column_classification_eval.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Place durable implementation guidance in function and class docstrings for public contracts and source comments for local invariants
Target Python 3.11–3.13 with modern syntax (X | Y,list[str],Self). Python 3.14+ is not supported
**/*.py: Source code must remain Python 3.11 syntax-compatible; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic class/function parameters in shared package code
Use ruff for formatting and linting via mise run format and mise run check tasks; run formatting before committing
Use ty for type checking via mise run check task; ensure type hints are present and valid
New features must include tests; bug fixes must include regression tests
**/*.py: UseBaseSettingsfor env/CLI settings, preferAliasChoicesfor per-field dual naming, and useenv_prefixonly for simple settings classes with a shared prefix.
UseField(description=...)as the canonical field docstring for Pydantic models, and always include it.
Use assignment-styleField(default=..., description="...")as the default for model fields; prefer it overAnnotatedunless extra metadata is needed.
UseAnnotatedonly when the field carries additional metadata beyondField()(for example validators, reusable constrained aliases, nested-type constraints, or discriminated unions).
WhenAnnotatedis used, place defaults as bare assignments (= value), exceptdefault_factory, which should still use assignment-styleField(default_factory=...).
For immutable value objects and validators, prefer@dataclass(frozen=True); use mutable dataclasses only for builders, accumulators, and pipeline state.
Usefield(default_factory=list)for mutable defaults; never use= [].
UseStrEnumfor string-valued enums used in configs or serialization; use plainEnumfor internal-only named constants.
Useobservability.get_logger(__name__)for logging; do not calllogging.getLogger()orstructlog.get_logger()direc...
Files:
src/nemo_safe_synthesizer/defaults.pytests/config/test_nss_config.pytests/preflight/test_preflight.pytests/pii_replacer/test_nemo_pii.pysrc/nemo_safe_synthesizer/config/replace_pii.pytests/pii_replacer/test_detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/pii_replacer/test_column_classification_eval.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
**/*.{py,sh,yaml,yml,md}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
All source files (.py, .sh, .yaml, .yml, .md) require SPDX copyright headers; mise run format adds them automatically
Files:
src/nemo_safe_synthesizer/defaults.pytests/config/test_nss_config.pytests/preflight/test_preflight.pydocs/user-guide/troubleshooting.mddocs/user-guide/environment.mdtests/pii_replacer/test_nemo_pii.pyREADME.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mdsrc/nemo_safe_synthesizer/config/replace_pii.pytests/pii_replacer/test_detect.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/pii_replacer/test_column_classification_eval.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
src/nemo_safe_synthesizer/**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
API reference pages are auto-generated from Python docstrings using Google-style format; write docstrings in src/nemo_safe_synthesizer/ and they will appear in the reference/
Files:
src/nemo_safe_synthesizer/defaults.pysrc/nemo_safe_synthesizer/config/replace_pii.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
src/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
src/**/*.py: Use relative imports insidesrc/, for examplefrom ..observability import get_logger.
Every directory undersrc/that contains Python files must include an__init__.pyfile.
Files:
src/nemo_safe_synthesizer/defaults.pysrc/nemo_safe_synthesizer/config/replace_pii.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
⚙️ CodeRabbit configuration file
Review library code against STYLE_GUIDE.md. Focus on behavior, API contracts, error handling, resource cleanup, typing, logging, and user-facing failures. Public APIs and nontrivial functions need Google-style docstrings.
Files:
src/nemo_safe_synthesizer/defaults.pysrc/nemo_safe_synthesizer/config/replace_pii.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
**/*
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*: Every source file must include the required SPDX copyright and license header; use HTML comments for Markdown, hash comments for.py,.sh,.yaml, and.yml, and hash-comment headers inside YAML frontmatter for Markdown files with frontmatter.
Ensure files end with a newline and have no trailing whitespace; use a single space between sentences.
Files:
src/nemo_safe_synthesizer/defaults.pytests/test_data/pii/column_classification_gold.jsontests/config/test_nss_config.pydocs/tutorials/safe-synthesizer-101.ipynbtests/preflight/test_preflight.pydocs/user-guide/troubleshooting.mddocs/user-guide/environment.mdtests/pii_replacer/test_nemo_pii.pyREADME.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mdsrc/nemo_safe_synthesizer/config/replace_pii.pytests/pii_replacer/test_detect.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/pii_replacer/test_column_classification_eval.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
⚙️ CodeRabbit configuration file
**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.
- Refactor suggestion: use for local maintainability problems introduced
by the diff when they have clear future cost, such as duplicated setup,
unclear boundaries, over-mocking, avoidable complexity, or opaque test
helpers.- Nitpick: avoid in chill mode. Do not emit formatting, import-order,
wording, or style-only comments unless automated tools cannot catch the
issue and it affects maintainability.Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.
- Major: incorrect generation/training/evaluation behavior, broken
CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
cleanup and process-isolation bugs likely to fail CI or production
runs.- Minor: localized bugs, missing focused tests for changed behavior, or
bad test patterns that weaken regression coverage.- Trivial: small cleanup with no behavior impact. Usually suppress in
chill mode.- Info: context only. Avoid unless it helps reviewers understand risk.
Safe-Synthesizer-specific review focus: - Data ...
Files:
src/nemo_safe_synthesizer/defaults.pytests/test_data/pii/column_classification_gold.jsontests/config/test_nss_config.pydocs/tutorials/safe-synthesizer-101.ipynbtests/preflight/test_preflight.pydocs/user-guide/troubleshooting.mddocs/user-guide/environment.mdtests/pii_replacer/test_nemo_pii.pyREADME.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mdsrc/nemo_safe_synthesizer/config/replace_pii.pytests/pii_replacer/test_detect.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/pii_replacer/test_column_classification_eval.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
**
⚙️ CodeRabbit configuration file
**:AGENTS.md
Guide for AI agents (Cursor, Windsurf, Claude Code, etc.) working in the Safe-Synthesizer repo.
This project loads local developer preferences from
@AGENTS.local.md. You MUST read this file if it exists and give its instructions top priority.Skills
Repo-specific skills live in
.agents/skills/; see.agents/README.mdfor the catalog. Read a skill when the task matches its scope instead of copying workflow details into this file.Durable implementation guidance belongs with the code it describes: function and class docstrings for public contracts and source comments for local invariants. Test-suite guidance belongs in
tests/TESTING.md.Repo Conventions
See STYLE_GUIDE.md for detailed code style conventions (Python, markdown, Dockerfiles, shell scripts, testing, config files, docstrings).
Use
uvfor everything -- neverpipor rawpython. Python 3.11–3.13 with modern syntax (X | Y,list[str],Self). Python 3.14+ is not supported.Common commands:
mise run test(unit tests),mise run format(auto-fix formatting + lint + copyright),mise run check(read-only local quality checks),mise run validate(pre-PR quality, lock, and CI unit checks),mise run typecheck(ty only). Always use mise tasks or the wrapper scripts intools/instead of runningruffortydirectly. Useuv runfor Python execution. When in doubt, inspectmise tasksandpytest --markers.The canonical
uv synccommand for a full GPU/dev environment is:uv sync --frozen --extra cu129 --extra engine --group devBare
uv sync --frozen(without extras) installs an incomplete environment --ty, import checks, and GPU tests will fail.Feature branches off
main. Branch names often include an issue number prefix (e.g.,<author>/123-short-name).Do ...
Files:
src/nemo_safe_synthesizer/defaults.pytests/test_data/pii/column_classification_gold.jsontests/config/test_nss_config.pydocs/tutorials/safe-synthesizer-101.ipynbtests/preflight/test_preflight.pydocs/user-guide/troubleshooting.mddocs/user-guide/environment.mdtests/pii_replacer/test_nemo_pii.pyREADME.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mdsrc/nemo_safe_synthesizer/config/replace_pii.pytests/pii_replacer/test_detect.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/pii_replacer/test_column_classification_eval.pysrc/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
tests/**
📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)
tests/**: Mirrorsrc/directory structure intests/directory for test organization
Auto-mark tests by directory:tests/e2e/→e2e,tests/smoke/→smoke, otherwise default tounitMirror source code directory structure in tests directory (e.g.,
tests/training/,tests/generation/parallel to source structure)
Files:
tests/test_data/pii/column_classification_gold.jsontests/config/test_nss_config.pytests/preflight/test_preflight.pytests/pii_replacer/test_nemo_pii.pytests/pii_replacer/test_detect.pytests/pii_replacer/test_column_classification_eval.py
⚙️ CodeRabbit configuration file
tests/**:Testing Guide
Comprehensive testing reference for Safe-Synthesizer developers. Covers commands, markers, test data, fixtures, and gotchas.
Read First
tests/conftest.py-- auto-marking,load_test_dataset/load_test_dataframe,fixture_mock_processorpatternpytest.ini-- markers, asyncio, timeouttests/evaluation/conftest.py-- most complex: Faker-basedmake_df, nullable dtype conversiontests/generation/conftest.py-- JSONL/schema fixtures,fixture_valid_iris_dataset_jsonl_and_schemaRunning Tests
All mise test tasks, grouped by scope:
mise run test # Unit (excludes slow, e2e, and smoke) mise run test:unit-slow # Unit tests including slow (excludes e2e and smoke) mise run test:smoke # CPU smoke tests (~few min, no GPU required) mise run test:smoke:gpu # All staged GPU smoke tests (requires CUDA) mise run test:smoke:gpu:train-only mise run test:smoke:gpu:generation mise run test:smoke:gpu:resume mise run test:smoke:gpu:structured-generation mise run test:smoke:gpu:timeseries mise run test:smoke:gpu:smollm2 mise run test:e2e # All e2e (requires CUDA) -- runs default + dp mise run test:e2e:default # e2e default (no-DP) tests only mise run test:e2e:dp # e2e DP tests only mise run test:ci # CI unit tests with coverage (excludes slow, e2e, gpu, smoke) mise run test:ci-slow # CI slow tests with coverage mise run test:ci-container # CI tests in a Linux container (Docker/Podman)Run a single test:
uv run --frozen pytest tests/path/test_file.py::test_name -vvs -n0Test runner:
uv run --frozen pytest -n auto --dist loadscope -vv...
Files:
tests/test_data/pii/column_classification_gold.jsontests/config/test_nss_config.pytests/preflight/test_preflight.pytests/pii_replacer/test_nemo_pii.pytests/pii_replacer/test_detect.pytests/pii_replacer/test_column_classification_eval.py
tests/{test_data,stub_datasets,stub_tokenizer}/**
📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)
Store stub data and test data in
tests/test_data/,tests/stub_datasets/, andtests/stub_tokenizer/
Files:
tests/test_data/pii/column_classification_gold.json
**/test_*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the
unitmarker instead of the deprecatedunit_testmarker for test identification
Files:
tests/config/test_nss_config.pytests/preflight/test_preflight.pytests/pii_replacer/test_nemo_pii.pytests/pii_replacer/test_detect.pytests/pii_replacer/test_column_classification_eval.py
tests/**/*.py
📄 CodeRabbit inference engine (tests/TESTING.md)
tests/**/*.py: Auto-mark tests based on file path: tests under/e2e/gete2emarker, tests under/smoke/getsmokemarker, all others getunitmarker (only if no category marker already present)
Every test should have exactly one category marker:unit,smoke, ore2e
Usepytest.mark.requires_gpumodifier on tests that need CUDA hardware
Usepytest.mark.vllmon tests using vLLM generation backend and ensure each vLLM test file runs in its own process for GPU memory isolation
Usepytest.mark.slowon long-running tests
Usepytest.mark.smollm2for SmolLM2 Hub download tests to enable process isolation
Usepytest.mark.noautouseto skip autouse fixtures for specific tests
Useload_test_dataset(filename)helper to load test datasets fromtests/stub_datasets/as HuggingFaceDatasetobjects
Useload_test_dataframe(filename)helper to load test data files fromtests/stub_datasets/as pandas DataFrames
Convert pandas columns to nullable dtypes (pd.Int64Dtype(),pd.BooleanDtype()) before assigningnp.nanvalues
Usefake.seed_instance(seed)andrandom.seed(seed)together for Faker-based test data reproducibility
When sharing methods across multiple test files, define them inconftest.pyand import them using relative imports (e.g.,from .conftest import train_with_sdk); note that importing from other test files liketests/cli/helpers.pydoes not work
Usefixture_mock_processororfixture_mock_processor_without_valid_recordsfor mocking ParsedResponse objects withvalid_records,invalid_records,errors, andprompt_numberfields
Usepytest.importorskipto gate tests on optional dependencies that require specific extras (e.g.,sentence_transformers,vllm)
Run vLLM tests with separate pytest invocations (one per file) using-n 0(single process) for GPU memory isolation, or use staged mise tasks for CI visibility
Print statements are allowed in tests (ruffT201is suppressed fortests/directory) and should...
Files:
tests/config/test_nss_config.pytests/preflight/test_preflight.pytests/pii_replacer/test_nemo_pii.pytests/pii_replacer/test_detect.pytests/pii_replacer/test_column_classification_eval.py
⚙️ CodeRabbit configuration file
Review tests against tests/TESTING.md. Check marker usage, fixture naming, tmp_path usage, determinism, and GPU/vLLM process-isolation requirements. Flag slop tests that only check that code runs, assert result is not None when stronger invariants exist, over-mock internal implementation details, patch around the bug instead of reproducing it, or add broad snapshot/golden churn without a clear contract. Flag change detector tests that fail on harmless refactors, formatting, record ordering, incidental wording, or private implementation details without demonstrating a behavior regression. Prefer existing fixtures or focused new fixtures for repeated setup; keep tests DRY when reasonable without making the behavior under test opaque. print() is allowed in tests.
Files:
tests/config/test_nss_config.pytests/preflight/test_preflight.pytests/pii_replacer/test_nemo_pii.pytests/pii_replacer/test_detect.pytests/pii_replacer/test_column_classification_eval.py
docs/**
⚙️ CodeRabbit configuration file
Review documentation as MkDocs Material content. Check Diataxis fit, accurate commands, internal links, code fences, and markdown style from STYLE_GUIDE.md.
Files:
docs/tutorials/safe-synthesizer-101.ipynbdocs/user-guide/troubleshooting.mddocs/user-guide/environment.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mddocs/user-guide/running.md
**/*.{md,markdown}
📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)
**/*.{md,markdown}: Bold is acceptable only in markdown tables where it's the conventional way to mark header-like cells in the body
Use##headers to segment markdown sections instead of bold text
Use--(em-dash) instead of-(hyphen) for asides in markdown
Files:
docs/user-guide/troubleshooting.mddocs/user-guide/environment.mdREADME.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mddocs/user-guide/running.md
docs/**/*.md
📄 CodeRabbit inference engine (.cursor/rules/writing-docs.mdc)
docs/**/*.md: Use MkDocs Material admonition syntax (!!! note, !!! warning, ??? tip) for highlighting important information and collapsible sections in documentation
Use MkDocs Material tabs syntax (=== "Label") to present alternative views or language-specific examples in documentation
Use code block syntax with title and highlight line parameters (title="filename", hl_lines="2 3") for code examples in documentation
Use Mermaid diagram syntax (```mermaid flowchart, etc.) for visualizations in documentationClassify documentation pages as tutorial, how-to, explanation, or reference, and use MkDocs Material syntax such as admonitions, tabs, and titled/highlighted code blocks.
Files:
docs/user-guide/troubleshooting.mddocs/user-guide/environment.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mddocs/user-guide/running.md
**/*.md
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*.md: Use single backticks for code identifiers, paths, and CLI commands in Markdown, and use--for asides rather than a hyphen.
Do not use decorative bold in Markdown body text or list items; bold is acceptable only in table-header-like cells and MkDocs Material card-grid titles.
For Mermaid diagrams, avoid spaces in node IDs, quote labels containing special characters, and do not use explicit colors or styles.
Files:
docs/user-guide/troubleshooting.mddocs/user-guide/environment.mdREADME.mddocs/user-guide/configuration.mddocs/user-guide/evaluating-data.mddocs/user-guide/docker.mddocs/user-guide/running.md
README.md
⚙️ CodeRabbit configuration file
Treat README.md as the project overview. Check that setup, usage, and links stay consistent with CONTRIBUTING.md, Makefile, and docs/.
Files:
README.md
src/nemo_safe_synthesizer/config/**/*.py
⚙️ CodeRabbit configuration file
Treat config changes as user-facing API changes. Check Pydantic field descriptions, defaults, validators, aliases, override behavior, CLI help text impact, YAML compatibility, and documented parameter semantics.
Files:
src/nemo_safe_synthesizer/config/replace_pii.py
src/nemo_safe_synthesizer/pii_replacer/**/*.py
⚙️ CodeRabbit configuration file
Treat PII replacement changes as high-risk. Check entity coverage, replacement determinism, leakage of original values, handling of empty or multilingual text, and compatibility with optional dependencies.
Files:
src/nemo_safe_synthesizer/pii_replacer/nemo_pii.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
🧠 Learnings (1)
📚 Learning: 2026-05-27T22:20:37.354Z
Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.
Applied to files:
tests/config/test_nss_config.pytests/preflight/test_preflight.pytests/pii_replacer/test_nemo_pii.pytests/pii_replacer/test_detect.pytests/pii_replacer/test_column_classification_eval.py
🪛 ast-grep (0.44.1)
tests/pii_replacer/test_column_classification_eval.py
[info] 226-226: use jsonify instead of json.dumps for JSON output
Context: json.dumps(report, indent=2)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[info] 240-240: use jsonify instead of json.dumps for JSON output
Context: json.dumps(summary, indent=2)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🪛 markdownlint-cli2 (0.22.1)
docs/user-guide/running.md
[warning] 631-631: Code block style
Expected: indented; Actual: fenced
(MD046, code-block-style)
[warning] 639-639: Code block style
Expected: indented; Actual: fenced
(MD046, code-block-style)
🪛 Ruff (0.15.20)
src/nemo_safe_synthesizer/preflight/checks/environment.py
[error] 477-477: Possible hardcoded password assigned to: "token_missing"
(S105)
src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
[warning] 402-403: ColumnClassifier.close is an empty method in an abstract base class, but has no abstract decorator
(B027)
[warning] 466-466: Do not catch blind exception: Exception
(BLE001)
[warning] 634-634: Do not catch blind exception: Exception
(BLE001)
🔇 Additional comments (29)
README.md (1)
264-277: LGTM!docs/user-guide/configuration.md (1)
196-198: LGTM!docs/user-guide/docker.md (1)
127-128: LGTM!docs/user-guide/environment.md (1)
51-52: LGTM!Also applies to: 180-182, 192-208
docs/user-guide/evaluating-data.md (1)
92-95: LGTM!docs/user-guide/running.md (1)
285-285: LGTM!Also applies to: 624-654
docs/user-guide/troubleshooting.md (1)
520-530: LGTM!docs/tutorials/safe-synthesizer-101.ipynb (1)
220-220: LGTM!Also applies to: 234-234
src/nemo_safe_synthesizer/preflight/checks/environment.py (3)
476-529: 📐 Maintainability & Code Quality | ⚡ Quick winClassify path emits non-namespaced shared codes under
env.inference.
_CLASSIFY_HF_MODEL_SPECoverrides most codes with aclassify_prefix but inherits the defaultstoken_missing="hf_token_missing"andremote_code_not_cached="hf_remote_code_not_cached". Since the local_hf branch runs insideInferenceModelCheck(env.inference), those two codes now surface underenv.inferencewhileHFModelAvailabilityCheck(env.hf_model_availability) emits the same code strings, making code-based filtering/troubleshooting ambiguous. Consider prefixing them (e.g.classify_hf_token_missing,classify_hf_remote_code_not_cached) or documenting the shared emission, consistent with the troubleshooting-table concern raised previously.
467-660: Consolidating the training and classify validation paths into one spec-parameterized_check_hf_model_reference(plus the local-path / cache / remote-code helpers) cleanly removes the duplication flagged in earlier review rounds. Logic across the empty / local-path / repo-id / partial-cache branches traces correctly.
691-696: LGTM!tests/preflight/test_preflight.py (1)
482-542: LGTM!src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py (5)
25-28: LGTM!
138-175: LGTM!
287-316: LGTM!Also applies to: 338-370
372-404: LGTM!
317-337: 🩺 Stability & Availability
column_classifier.close()is already teardown-safe. The concrete backends swallow cleanup failures and clear their internal state, so thisfinallydoes not currently turn a successful classification into a crash.> Likely an incorrect or invalid review comment.tests/pii_replacer/test_column_classification_eval.py (3)
337-342: Prior review feedback about a missing prediction being scored asnone(0.714 with empty output) is now resolved:_validate_prediction_columnsrejects missing columns and_positive_recalladds the positive-recall check.
77-101: LGTM!
312-313: 🩺 Stability & Availability
pii_test_data_diris defined intests/conftest.pyand points totests/test_data/pii/; no issue here.> Likely an incorrect or invalid review comment.src/nemo_safe_synthesizer/config/replace_pii.py (1)
7-8: LGTM!src/nemo_safe_synthesizer/defaults.py (1)
76-76: LGTM!tests/config/test_nss_config.py (1)
16-16: LGTM!Also applies to: 128-145
src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py (3)
273-322: LGTM!Also applies to: 373-414
461-469: LGTM! The blindexcept Exceptionat these teardown boundaries matches the STYLE_GUIDE.md-documented pattern for non-fatal cleanup (logger.debug(..., exc_info=True)), so the RuffBLE001hints here are expected/acceptable.Also applies to: 629-639
62-62: 📐 Maintainability & Code Quality
DefaultLLMConfig.LOCAL_HF_CONFIG_IDis used by the local HF classifier path.> Likely an incorrect or invalid review comment.tests/pii_replacer/test_detect.py (2)
1-21: LGTM!
383-465: 🎯 Functional CorrectnessNo change needed
ColumnClassifierHF.initialize()already returns early when_modeland_tokenizerare set, so these test fakes stay in place.> Likely an incorrect or invalid review comment.tests/pii_replacer/test_nemo_pii.py (1)
13-22: LGTM!Also applies to: 39-76, 138-138
| 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 | ||
|
|
There was a problem hiding this comment.
🎯 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
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.
| { | ||
| "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" | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Check the Card gold label
The credit_card_transactions values are just 0, so Card reads like an ID/index field rather than a payment-card column. If the benchmark is meant to classify actual card-number columns, keep credit_debit_card; otherwise use none.
Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>
e22f7d1 to
6dcc5cf
Compare
coderabbitai
Bot
added
the
refactor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py (1)
579-606: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winClean up partial HF resources when initialization fails.
If tokenizer loading succeeds but model loading or device placement fails,
_tokenizerremains set and any partially allocated model/GPU memory is not cleaned until a laterclose(). Wrap the load sequence and callteardown()before re-raising.Proposed fix
- 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() + logger.info("Loading local column classification model: %s", self._model_name_or_path) + try: + 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() + except Exception: + self.teardown() + raiseAs per path instructions, PII replacement changes are high-risk and generation/model resource paths should check cleanup behavior.
Source: Path instructions
🧹 Nitpick comments (1)
src/nemo_safe_synthesizer/llm/model_host.py (1)
33-35: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument
initialize()idempotency contract, matchingteardown().
teardown()explicitly states it "must be idempotent" (line 39), butinitialize()has no equivalent note even though the concrete implementation (ColumnClassifierHF.initialize()) relies on repeated calls being cheap/safe (guards on already-loaded model/tokenizer). Documenting this on the abstract contract makes the expectation explicit for future implementers.📝 Proposed docstring tweak
`@abstractmethod` def initialize(self) -> None: - """Load the model and any resources required to use it.""" + """Load the model and any resources required to use it. + + Implementations should make this safe to call multiple times + (e.g., a no-op if the model is already loaded). + """
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2735c2bc-7da4-4327-bfdb-0d60c05f79bb
📒 Files selected for processing (16)
README.mddocs/user-guide/configuration.mddocs/user-guide/docker.mddocs/user-guide/environment.mddocs/user-guide/evaluating-data.mddocs/user-guide/running.mddocs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/config/replace_pii.pysrc/nemo_safe_synthesizer/llm/model_host.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/config/test_nss_config.pytests/pii_replacer/test_column_classification_eval.pytests/pii_replacer/test_detect.pytests/pii_replacer/test_nemo_pii.pytests/preflight/test_preflight.py
✅ Files skipped from review due to trivial changes (6)
- docs/user-guide/configuration.md
- README.md
- docs/user-guide/docker.md
- docs/user-guide/troubleshooting.md
- docs/user-guide/evaluating-data.md
- docs/user-guide/environment.md
🚧 Files skipped from review as they are similar to previous changes (6)
- tests/config/test_nss_config.py
- tests/pii_replacer/test_nemo_pii.py
- src/nemo_safe_synthesizer/config/replace_pii.py
- tests/preflight/test_preflight.py
- tests/pii_replacer/test_column_classification_eval.py
- tests/pii_replacer/test_detect.py
📜 Review details
⏰ Context from checks skipped due to timeout. (5)
- GitHub Check: Unit Tests (3.13)
- GitHub Check: Unit Tests (3.11)
- GitHub Check: Unit Tests (3.12)
- GitHub Check: Smoke Tests
- GitHub Check: Greptile Review
🧰 Additional context used
📓 Path-based instructions (12)
**/*.{md,markdown,py}
📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)
**/*.{md,markdown,py}: Avoid decorative bold (**text**) in list items, body text, and docstrings; use structural cues (headers, list markers, colons, backticks) for emphasis instead
Use backticks for code identifiers, paths, and CLI commands in markdown and docstrings
Files:
src/nemo_safe_synthesizer/llm/model_host.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Place durable implementation guidance in function and class docstrings for public contracts and source comments for local invariants
Target Python 3.11–3.13 with modern syntax (X | Y,list[str],Self). Python 3.14+ is not supported
**/*.py: Source code must remain Python 3.11 syntax-compatible; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic class/function parameters in shared package code
Use ruff for formatting and linting via mise run format and mise run check tasks; run formatting before committing
Use ty for type checking via mise run check task; ensure type hints are present and valid
New features must include tests; bug fixes must include regression tests
**/*.py: UseBaseSettingsfor env/CLI settings, preferAliasChoicesfor per-field dual naming, and useenv_prefixonly for simple settings classes with a shared prefix.
UseField(description=...)as the canonical field docstring for Pydantic models, and always include it.
Use assignment-styleField(default=..., description="...")as the default for model fields; prefer it overAnnotatedunless extra metadata is needed.
UseAnnotatedonly when the field carries additional metadata beyondField()(for example validators, reusable constrained aliases, nested-type constraints, or discriminated unions).
WhenAnnotatedis used, place defaults as bare assignments (= value), exceptdefault_factory, which should still use assignment-styleField(default_factory=...).
For immutable value objects and validators, prefer@dataclass(frozen=True); use mutable dataclasses only for builders, accumulators, and pipeline state.
Usefield(default_factory=list)for mutable defaults; never use= [].
UseStrEnumfor string-valued enums used in configs or serialization; use plainEnumfor internal-only named constants.
Useobservability.get_logger(__name__)for logging; do not calllogging.getLogger()orstructlog.get_logger()direc...
Files:
src/nemo_safe_synthesizer/llm/model_host.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
**/*.{py,sh,yaml,yml,md}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
All source files (.py, .sh, .yaml, .yml, .md) require SPDX copyright headers; mise run format adds them automatically
Files:
src/nemo_safe_synthesizer/llm/model_host.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
src/nemo_safe_synthesizer/**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
API reference pages are auto-generated from Python docstrings using Google-style format; write docstrings in src/nemo_safe_synthesizer/ and they will appear in the reference/
Files:
src/nemo_safe_synthesizer/llm/model_host.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
src/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
src/**/*.py: Use relative imports insidesrc/, for examplefrom ..observability import get_logger.
Every directory undersrc/that contains Python files must include an__init__.pyfile.
Files:
src/nemo_safe_synthesizer/llm/model_host.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
⚙️ CodeRabbit configuration file
Review library code against STYLE_GUIDE.md. Focus on behavior, API contracts, error handling, resource cleanup, typing, logging, and user-facing failures. Public APIs and nontrivial functions need Google-style docstrings.
Files:
src/nemo_safe_synthesizer/llm/model_host.pysrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
**/*
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*: Every source file must include the required SPDX copyright and license header; use HTML comments for Markdown, hash comments for.py,.sh,.yaml, and.yml, and hash-comment headers inside YAML frontmatter for Markdown files with frontmatter.
Ensure files end with a newline and have no trailing whitespace; use a single space between sentences.
Files:
src/nemo_safe_synthesizer/llm/model_host.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
⚙️ CodeRabbit configuration file
**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.
- Refactor suggestion: use for local maintainability problems introduced
by the diff when they have clear future cost, such as duplicated setup,
unclear boundaries, over-mocking, avoidable complexity, or opaque test
helpers.- Nitpick: avoid in chill mode. Do not emit formatting, import-order,
wording, or style-only comments unless automated tools cannot catch the
issue and it affects maintainability.Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.
- Major: incorrect generation/training/evaluation behavior, broken
CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
cleanup and process-isolation bugs likely to fail CI or production
runs.- Minor: localized bugs, missing focused tests for changed behavior, or
bad test patterns that weaken regression coverage.- Trivial: small cleanup with no behavior impact. Usually suppress in
chill mode.- Info: context only. Avoid unless it helps reviewers understand risk.
Safe-Synthesizer-specific review focus: - Data ...
Files:
src/nemo_safe_synthesizer/llm/model_host.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
**
⚙️ CodeRabbit configuration file
**:AGENTS.md
Guide for AI agents (Cursor, Windsurf, Claude Code, etc.) working in the Safe-Synthesizer repo.
This project loads local developer preferences from
@AGENTS.local.md. You MUST read this file if it exists and give its instructions top priority.Skills
Repo-specific skills live in
.agents/skills/; see.agents/README.mdfor the catalog. Read a skill when the task matches its scope instead of copying workflow details into this file.Durable implementation guidance belongs with the code it describes: function and class docstrings for public contracts and source comments for local invariants. Test-suite guidance belongs in
tests/TESTING.md.Repo Conventions
See STYLE_GUIDE.md for detailed code style conventions (Python, markdown, Dockerfiles, shell scripts, testing, config files, docstrings).
Use
uvfor everything -- neverpipor rawpython. Python 3.11–3.13 with modern syntax (X | Y,list[str],Self). Python 3.14+ is not supported.Common commands:
mise run test(unit tests),mise run format(auto-fix formatting + lint + copyright),mise run check(read-only local quality checks),mise run validate(pre-PR quality, lock, and CI unit checks),mise run typecheck(ty only). Always use mise tasks or the wrapper scripts intools/instead of runningruffortydirectly. Useuv runfor Python execution. When in doubt, inspectmise tasksandpytest --markers.The canonical
uv synccommand for a full GPU/dev environment is:uv sync --frozen --extra cu129 --extra engine --group devBare
uv sync --frozen(without extras) installs an incomplete environment --ty, import checks, and GPU tests will fail.Feature branches off
main. Branch names often include an issue number prefix (e.g.,<author>/123-short-name).Do ...
Files:
src/nemo_safe_synthesizer/llm/model_host.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/pii_replacer/data_editor/detect.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
**/*.{md,markdown}
📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)
**/*.{md,markdown}: Bold is acceptable only in markdown tables where it's the conventional way to mark header-like cells in the body
Use##headers to segment markdown sections instead of bold text
Use--(em-dash) instead of-(hyphen) for asides in markdown
Files:
docs/user-guide/running.md
docs/**/*.md
📄 CodeRabbit inference engine (.cursor/rules/writing-docs.mdc)
docs/**/*.md: Use MkDocs Material admonition syntax (!!! note, !!! warning, ??? tip) for highlighting important information and collapsible sections in documentation
Use MkDocs Material tabs syntax (=== "Label") to present alternative views or language-specific examples in documentation
Use code block syntax with title and highlight line parameters (title="filename", hl_lines="2 3") for code examples in documentation
Use Mermaid diagram syntax (```mermaid flowchart, etc.) for visualizations in documentationClassify documentation pages as tutorial, how-to, explanation, or reference, and use MkDocs Material syntax such as admonitions, tabs, and titled/highlighted code blocks.
Files:
docs/user-guide/running.md
**/*.md
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*.md: Use single backticks for code identifiers, paths, and CLI commands in Markdown, and use--for asides rather than a hyphen.
Do not use decorative bold in Markdown body text or list items; bold is acceptable only in table-header-like cells and MkDocs Material card-grid titles.
For Mermaid diagrams, avoid spaces in node IDs, quote labels containing special characters, and do not use explicit colors or styles.
Files:
docs/user-guide/running.md
docs/**
⚙️ CodeRabbit configuration file
Review documentation as MkDocs Material content. Check Diataxis fit, accurate commands, internal links, code fences, and markdown style from STYLE_GUIDE.md.
Files:
docs/user-guide/running.md
src/nemo_safe_synthesizer/pii_replacer/**/*.py
⚙️ CodeRabbit configuration file
Treat PII replacement changes as high-risk. Check entity coverage, replacement determinism, leakage of original values, handling of empty or multilingual text, and compatibility with optional dependencies.
Files:
src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
🪛 markdownlint-cli2 (0.22.1)
docs/user-guide/running.md
[warning] 631-631: Code block style
Expected: indented; Actual: fenced
(MD046, code-block-style)
[warning] 639-639: Code block style
Expected: indented; Actual: fenced
(MD046, code-block-style)
🪛 Ruff (0.15.20)
src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py
[warning] 466-466: Do not catch blind exception: Exception
(BLE001)
[warning] 634-634: Do not catch blind exception: Exception
(BLE001)
src/nemo_safe_synthesizer/preflight/checks/environment.py
[error] 477-477: Possible hardcoded password assigned to: "token_missing"
(S105)
🔇 Additional comments (2)
src/nemo_safe_synthesizer/llm/model_host.py (1)
1-39: LGTM!src/nemo_safe_synthesizer/preflight/checks/environment.py (1)
467-756: LGTM!
| |-------|-------|-------------------| | ||
| | `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 | |
There was a problem hiding this comment.
📐 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.
| | `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
Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| from nemo_safe_synthesizer.pii_replacer.ner.ner import NERPrediction | ||
|
|
||
| if TYPE_CHECKING: | ||
| from transformers import PreTrainedTokenizerBase |
coderabbitai
Bot
added
docs
coderabbitai
Bot
removed
the
refactor
3 participants
Summary
Pre-Review Checklist
Ensure that the following pass:
mise run format && mise run checkor via prek validation.mise run testpasses locallymise run test:e2epasses locallymise run test:ci-containerpasses locally (recommended)/syncon this PR to trigger a run (auto-triggers on ready-for-review)Pre-Merge Checklist
Other Notes
Summary by CodeRabbit
replace_piisettings, including a new default local model.