Skip to content

feat: add entity coverage as the new Nemo Anonymizer feature#195

Open
memadi-nv wants to merge 23 commits into
mainfrom
memadi/feature/add-entity-coverage
Open

feat: add entity coverage as the new Nemo Anonymizer feature#195
memadi-nv wants to merge 23 commits into
mainfrom
memadi/feature/add-entity-coverage

Conversation

@memadi-nv

@memadi-nv memadi-nv commented Jun 26, 2026

Copy link
Copy Markdown
Contributor

Summary

Changes

New: Entity Coverage Judge (engine/evaluation/entity_coverage_judge.py)

  • LLM extraction phase — the judge independently scans the original (pre-anonymization) text and identifies all in-scope sensitive values that appear in the anonymized output. It is scoped by entity_labels (the detection taxonomy the user configured — e.g. ["first_name", "email"]; when None, all PII types are in scope), data_summary (semantic context), and strict_entity_protection (rewrite-only: also flags inferable/indirect values, not just literal identifiers).

  • Deterministic postprocessing phase — before any result is surfaced, a pure-Python filter removes false positives from the LLM output:

    • Non-literal filter (_filter_nonliteral_entities): drops candidates whose value field does not appear verbatim in the original text — guards against the LLM hallucinating entities.
    • Deduplication (_deduplicate_judge_entities): collapses duplicate candidates.
    • Coverage filter (_filter_covered_leaked_entities): removes candidates already accounted for by Anonymizer's final entities via three matching modes:
      • Exact — tokenized candidate equals a final entity's tokens.
      • Subspan — candidate's core tokens (stopwords stripped) appear as a contiguous, in-order subsequence within a single final entity.
      • Composite — candidate tokens are a concatenation of whole final-entity values (e.g. "AliceSmith" covered by ["Alice", "Smith"]).
  • Output columnsentity_coverage (float: n_final / (n_final + n_leaked), 1.0 = nothing leaked, None = judge unavailable) and leaked_entities (list of {value, label, reasoning} dicts).

  • Model roleentity_coverage_judge, defaults to nemotron-super in evaluate.yaml.

EvaluateConfig: detection validity is now opt-in

  • Added compute_detection_validity: bool = False to EvaluateConfig.
  • Detection validity (tag-precision judge) is now off by default — opt in with EvaluateConfig(compute_detection_validity=True). This is an internal model/threshold diagnostic, not a customer-facing metric.
  • Entity coverage always runs in both Replace and Rewrite modes regardless of EvaluateConfig.

Interface (anonymizer.py, results.py, display.py)

  • Anonymizer.evaluate() now accepts config: EvaluateConfig | None and routes compute_detection_validity to the underlying judge set.
  • entity_coverage and leaked_entities are included in AnonymizerResult columns and surfaced in display_record().
  • Logging updated to report coverage alongside other evaluate metrics.

Docs (docs/concepts/evaluation.md, docs/concepts/models.md)

  • New "Entity Coverage" section explaining the judge, its output columns, edge cases (no entities → 1.0, judge failure → None), and the two-phase (LLM + postprocessing) design.
  • Detection Validity section updated to reflect opt-in status.
  • Model roles table updated to include entity_coverage_judge.

Skill (skills/anonymizer/SKILL.md)

  • Updated to reflect entity coverage always running, the new model role, updated verdict-column descriptions, and the EvaluateConfig knob.

Tests

  • tests/engine/test_entity_coverage_judge.py — unit tests for the postprocessing helpers (exact/subspan/composite matching, non-literal filter, deduplication) and the full workflow.
  • tests/interface/test_anonymizer_interface.py — integration coverage for the new evaluate path.
  • tests/interface/test_anonymizer_logging.py — logging assertions for the new coverage columns.

Type of Change

  • Bug fix
  • New feature
  • Breaking change
  • Documentation update
  • Refactoring

Testing

  • make test passes locally
  • make check passes locally (format + lint + typecheck + lock-check)
  • Added/updated tests for changes

Documentation

  • If docs changed: make docs-build passes locally

Related Issues

Closes #193

@memadi-nv memadi-nv changed the title Memadi/feature/add entity coverage feat: add entity coverage as the new Nemo Anonymizer feature Jun 29, 2026
@memadi-nv
memadi-nv force-pushed the memadi/feature/add-entity-coverage branch from f63da45 to 6186d0c Compare June 29, 2026 23:38
memadi-nv added 14 commits July 20, 2026 09:56
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
@memadi-nv
memadi-nv force-pushed the memadi/feature/add-entity-coverage branch from 1520569 to a532a6e Compare July 20, 2026 16:59
Separate exhaustive candidate extraction from deterministic coverage filtering and add regression coverage for structured responses and prompt behavior.

Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: memadi <memadi@nvidia.com>
@memadi-nv
memadi-nv force-pushed the memadi/feature/add-entity-coverage branch from a532a6e to f999832 Compare July 20, 2026 17:02
The refactor that unified the replace/rewrite evaluate() code paths dropped
check_rewrite_judge=True from validate_model_alias_references, causing a
misconfigured rewrite_judge alias to go undetected until LLM calls were made.
Pass check_rewrite_judge=is_rewrite to restore pre-refactor behaviour.

Signed-off-by: memadi <memadi@nvidia.com>
@memadi-nv
memadi-nv marked this pull request as ready for review July 20, 2026 17:10
@memadi-nv
memadi-nv requested a review from a team as a code owner July 20, 2026 17:10
@memadi-nv
memadi-nv requested a review from a team as a code owner July 20, 2026 17:10
@greptile-apps

greptile-apps Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds an Entity Coverage Judge — a two-phase LLM + deterministic-postprocessing judge that reports residual PII leakage after anonymization. It also makes the existing Detection Validity judge opt-in (EvaluateConfig(compute_detection_validity=True)), previously always-on.

  • New EntityCoverageWorkflow independently extracts PII from the original text, then a deterministic filter removes candidates already accounted for by the anonymizer's final entities (exact, subspan, and composite matching). The result is entity_coverage (0–1 float) and leaked_entities (list) added to every evaluate() call in both Replace and Rewrite modes.
  • EvaluateConfig gains compute_detection_validity: bool = False; the evaluate() signature is updated to accept this config; AnonymizerResult and PreviewResult gain entity_labels, strict_entity_protection, and data_summary fields so the coverage judge uses the same context as the original run.
  • Interface, display, logging, docs, and tests are all updated to reflect the new column and the changed defaults.

Confidence Score: 3/5

The core feature logic is well-designed and the deterministic postprocessing filters are thoroughly tested, but the rewrite-path coverage workflow has a shape-mismatch defect that silently discards all scored rows when the LLM adapter drops even one record.

The entity coverage judge is a meaningful new capability with solid test coverage for the postprocessing helpers. The replace-mode path handles partial adapter failures correctly via a LEFT JOIN. However, the rewrite-mode path calls run_non_critical, which assigns result.dataframe[col].values (M rows) directly into a copy of the original N-row dataframe — if M < N due to any dropped row, pandas raises a length mismatch, the except swallows it, and all N records get entity_coverage=None instead of just the failed ones. This can trigger on any LLM timeout or parse error in production, quietly zeroing out all coverage scores for an entire batch.

entity_coverage_judge.py — specifically run_non_critical (rewrite path, row-count mismatch when adapter drops records) and _build_prompt (dead-code stub with misleading defaults). replace_runner.py has a stale docstring but is otherwise correct.

Important Files Changed

Filename Overview
src/anonymizer/engine/evaluation/entity_coverage_judge.py New entity coverage judge: LLM + deterministic postprocessing (non-literal filter, deduplication, subspan/composite coverage filter). Has a P1 shape-mismatch bug in run_non_critical when the adapter drops rows, and a P2 _build_prompt stub issue.
src/anonymizer/interface/anonymizer.py evaluate() refactored to handle both rewrite and replace paths with entity coverage; entity_labels, strict_entity_protection, and data_summary are now forwarded through the result objects. Logic is correct; runs entity coverage separately for rewrite and as part of merged judges for replace.
src/anonymizer/engine/replace/replace_runner.py Entity coverage judge added to _run_merged_judges as always-first; detection validity gated behind compute_detection_validity. Uses LEFT JOIN for row-drop safety. Docstring stale (P2).
src/anonymizer/engine/rewrite/rewrite_workflow.py Detection validity gated behind compute_detection_validity; passthrough defaults conditional; holistic judge unchanged. Clean changes.
src/anonymizer/interface/display.py New _render_entity_coverage_section added for both replace and rewrite modes; n_detected and n_final counting is consistent via _count_detected_entity_label_pairs. Clean.
src/anonymizer/interface/results.py entity_labels, strict_entity_protection, data_summary added to AnonymizerResult and PreviewResult for round-tripping into evaluate(). Straightforward additions.
src/anonymizer/config/anonymizer_config.py EvaluateConfig gains compute_detection_validity: bool = False. Clean, well-documented addition.
tests/engine/test_entity_coverage_judge.py New unit tests for postprocessing helpers (exact/subspan/composite matching, non-literal filter, deduplication) and column_config smoke test. Good coverage of edge cases.
tests/interface/test_anonymizer_interface.py New tests verify strict_entity_protection and data_summary flow through run/preview/evaluate, and that detection validity defaults to False. Thorough integration coverage.
src/anonymizer/config/default_model_configs/models.yaml Adds nemotron-super alias with thinking disabled (reasoning_effort: none, enable_thinking: false). Config is consistent with evaluation use case.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant U as User
    participant A as Anonymizer.evaluate()
    participant RW as RewriteWorkflow / ReplaceWorkflow
    participant ECW as EntityCoverageWorkflow
    participant LLM as LLM (nemotron-super)
    participant PP as Deterministic Postprocess

    U->>A: evaluate(result, config)
    A->>RW: evaluate(internal_df, compute_detection_validity)
    RW-->>A: rewrite/replace result (with judge_evaluation)

    A->>ECW: run_non_critical(result.dataframe)
    ECW->>ECW: prepare() — flatten EntitiesByValueSchema
    ECW->>LLM: _coverage_prompt(entity_labels, strict_entity_protection, data_summary)
    LLM-->>ECW: "raw JSON {leaked_entities: [...]}"

    ECW->>PP: _parse_leaked_entities(raw)
    PP->>PP: _filter_nonliteral_entities
    PP->>PP: _deduplicate_judge_entities
    PP->>PP: _filter_covered_leaked_entities
    PP-->>ECW: verified leaked_entities list

    ECW->>ECW: "coverage = n_final / (n_final + n_leaked)"
    ECW-->>A: (judged_df, failed_records)

    A->>A: _build_user_dataframe
    A-->>U: AnonymizerResult with entity_coverage + leaked_entities
Loading
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant U as User
    participant A as Anonymizer.evaluate()
    participant RW as RewriteWorkflow / ReplaceWorkflow
    participant ECW as EntityCoverageWorkflow
    participant LLM as LLM (nemotron-super)
    participant PP as Deterministic Postprocess

    U->>A: evaluate(result, config)
    A->>RW: evaluate(internal_df, compute_detection_validity)
    RW-->>A: rewrite/replace result (with judge_evaluation)

    A->>ECW: run_non_critical(result.dataframe)
    ECW->>ECW: prepare() — flatten EntitiesByValueSchema
    ECW->>LLM: _coverage_prompt(entity_labels, strict_entity_protection, data_summary)
    LLM-->>ECW: "raw JSON {leaked_entities: [...]}"

    ECW->>PP: _parse_leaked_entities(raw)
    PP->>PP: _filter_nonliteral_entities
    PP->>PP: _deduplicate_judge_entities
    PP->>PP: _filter_covered_leaked_entities
    PP-->>ECW: verified leaked_entities list

    ECW->>ECW: "coverage = n_final / (n_final + n_leaked)"
    ECW-->>A: (judged_df, failed_records)

    A->>A: _build_user_dataframe
    A-->>U: AnonymizerResult with entity_coverage + leaked_entities
Loading

Comments Outside Diff (2)

  1. src/anonymizer/engine/evaluation/entity_coverage_judge.py, line 902-933 (link)

    P1 Shape mismatch when adapter drops rows causes all-row score loss

    run_non_critical copies the original N-row dataframe into out, then assigns coverage values with out[col] = result.dataframe[col].values. If the LLM adapter drops even one row (timeout, parse error, rate-limit), result.dataframe has M < N rows, and the numpy value array has M elements — causing pandas to raise ValueError: Length of values (M) does not match length of index (N). The enclosing except Exception catches this and falls back to entity_coverage=None for all N rows, discarding the M successfully scored records. In contrast, the replace-mode path (_run_merged_judges) uses a LEFT JOIN on RECORD_ID_COLUMN to handle dropped rows per-record. The rewrite-path coverage judge lacks this guard, so a single partial failure silently wipes every score.

  2. src/anonymizer/engine/replace/replace_runner.py, line 123-131 (link)

    P2 Stale docstring: detection is now opt-in and entity coverage is missing

    The docstring still describes the pre-PR behaviour where detection validity always ran and entity coverage didn't exist. After this change: entity coverage always runs for all strategies; detection validity is opt-in via compute_detection_validity; the "4 judges" count and "detection judge only" description are both wrong.

    Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Reviews (1): Last reviewed commit: "fix: restore check_rewrite_judge validat..." | Re-trigger Greptile

Comment on lines +484 to +486
@classmethod
def _build_prompt(cls) -> str:
return _coverage_prompt(entity_labels=None, strict_entity_protection=False)

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 _build_prompt classmethod ignores instance parameters — latent maintenance trap

The base class marks _build_prompt as the canonical prompt source (its column_config calls self._build_prompt()). This implementation hardcodes entity_labels=None, strict_entity_protection=False, so if the base-class column_config or evaluate path is ever reached — via super() in a subclass, a test calling the base method directly, or a refactor that removes the column_config override — the judge would silently scope to all PII types and non-strict mode regardless of the configured values. At minimum the method should document that it is intentionally unused in favour of the column_config override.

Suggested change
@classmethod
def _build_prompt(cls) -> str:
return _coverage_prompt(entity_labels=None, strict_entity_protection=False)
@classmethod
def _build_prompt(cls) -> str:
# NOTE: This stub exists only to satisfy the abstract base class interface.
# EntityCoverageWorkflow overrides column_config() to build the prompt from
# instance attributes (entity_labels, strict_entity_protection, data_summary).
# _build_prompt() is never called in normal execution. If this class is
# subclassed or the column_config override is removed, this fallback would
# silently ignore instance configuration.
return _coverage_prompt(entity_labels=None, strict_entity_protection=False)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Replace detection_validity score with leakage_judge in NA replace

1 participant