diff --git a/plan/plan.md b/plan/plan.md new file mode 100644 index 00000000..6a8ae9a7 --- /dev/null +++ b/plan/plan.md @@ -0,0 +1,366 @@ +# Entity Coverage — Implementation Plan + +## Problem + +Currently `detection_valid` scores **precision** of the detected entities. Observed scores are predominantly in the lower ranges issue [#176](https://github.com/NVIDIA-NeMo/Anonymizer/issues/176), partly because the judge flags over-detection and mislabeling among other issues. But the actual privacy-sensitive question users care about is the opposite: **did any real PII slip through the anonymizer?** + +Precision complaints are already surfaced in the entity dropdown for manual inspection, making `detection_valid` feel redundant as a headline metric. What's missing is a **recall** signal: a score that answers "of all the PII that existed, how much did we successfully catch?" + +--- + +## Proposed Solution: Entity Coverage + +Replace `detection_valid` with `entity_coverage` (name TBD): + +``` +entity_coverage = n_entities_detected / (n_entities_detected + n_entities_leaked) +``` + +- **`n_entities_detected`** — identifiers detected by Anonymizer in the original text (via the detection pipeline) +- **`n_entities_leaked`** — identifiers missed by Anonymizer in the original text, found by the evaluation LLM +- **`leaked_entities`** — the list of missed entities (shown in a dropdown when coverage < 100%) + +The evaluation LLM runs on the **original text** and detects all PII that should have been caught. The delta between what it finds and what Anonymizer already detected gives the leaked entities. + +**Key shift:** current judge asks "were our detections correct?" (precision). New judge asks "did any PII survive?" (recall). + +**Question TBD:** The evaluation and anonymization currently use the same LLM, and we may need to address potential bias by using a different model for evaluation. Potential models: Nemotron Ultra and Nemotron Super (previous [research log](https://nvidia.atlassian.net/wiki/spaces/SDGR/pages/3597435385/2026-06-09+Entity+Validator+Augmentor+Model+Selection+for+NeMo+Anonymizer) around experimenting these 2 models) + +--- + +## Design Decisions + +### Coverage display differs by mode + +| Mode | Display format | Example | +|---|---|---| +| **Replace** | `Entity Coverage: — n/d (pct%)` | `Entity Coverage: Not Satisfied — 18/21 (86%)` | +| **Rewrite** | `Entity Coverage: — n/d (pct%)` | `Entity Coverage: 0.86 — 18/21 (86%)` | + +Where `n` = `n_entities_detected`, `d` = `n_entities_detected + n_entities_leaked`, and `pct` is derived from the fraction. The fraction gives absolute scale; the percentage makes it immediately readable without mental math. No explanatory subtitle line is shown — entity coverage is self-explanatory. + +A value of `1.0` (or Satisfied) means no PII leaked into the output. Any value below `1.0` (or Not Satisfied) means at least one entity was missed, and `leaked_entities` names them. + +### Passthrough rows + +Rows where Anonymizer detected **zero entities** trivially score `entity_coverage = 1.0` — there was nothing to protect, so nothing could have leaked. The LLM judge is skipped for these rows (same pattern as the current passthrough shortcut in `_BaseJudgeWorkflow`). + +### Prompt starting point + +The entity judge prompt is based on the [leakage eval script](https://gitlab-master.nvidia.com/sdg-research/anonymizer-experiments/-/blob/lramaswamy/latency-benchmark/experiments/latency-benchmark/scripts/leakage_eval.py?ref_type=heads#L26) from the experiments repo. It will be refined as it is tested against real examples. The core task: given the original text, identify all PII — then report any that Anonymizer missed. + +### Reuse `_BaseJudgeWorkflow` + +The new `EntityCoverageWorkflow` extends `_BaseJudgeWorkflow`. The partition-LLM-merge skeleton is identical; only the schema, prompt, and post-process step differ (coverage fraction computed from counts, not a boolean verdict field). `postprocess()` is overridden in the subclass — the base class stays unchanged. + +--- + +## Scope + +- Replaces `detection_valid` and `detection_invalid_entities` in both replace and rewrite `evaluate()` output. +- No new public API symbols beyond new column names (`entity_coverage`, `leaked_entities`) and a renamed model role. +- `COL_DETECTION_VALID` / `COL_DETECTION_INVALID_ENTITIES` constants are retired; new constants added. +- `DetectionJudgeWorkflow` is replaced by `EntityCoverageWorkflow` — the old class is removed, not kept as an alias. + +--- + +## New Column Names + +| Constant | Value | Type | Description | +|---|---|---|---| +| `COL_ENTITY_COVERAGE` | `"entity_coverage"` | `float \| None` | Fraction of entities successfully anonymized; `None` if judge unavailable | +| `COL_LEAKED_ENTITIES` | `"leaked_entities"` | `list` | Each missed entity with `value`, `label`, `reasoning` | +| `COL_ENTITY_COVERAGE_JUDGE` | `"_entity_coverage_judge"` | internal | Raw judge output; internal only | + +--- + +## Files Changed + +| File | Change | +|---|---| +| `src/anonymizer/engine/constants.py` | Retire `COL_DETECTION_JUDGE`, `COL_DETECTION_VALID`, `COL_DETECTION_INVALID_ENTITIES`; add `COL_ENTITY_COVERAGE_JUDGE`, `COL_ENTITY_COVERAGE`, `COL_LEAKED_ENTITIES` | +| `src/anonymizer/engine/evaluation/detection_judge.py` | Replace entirely with `entity_coverage_judge.py` — new schema, new prompt, new workflow class | +| `src/anonymizer/engine/evaluation/judge_base.py` | No changes needed — `postprocess()` is overridden in the subclass | +| `src/anonymizer/engine/rewrite/rewrite_workflow.py` | Swap `DetectionJudgeWorkflow` → `EntityCoverageWorkflow` in `evaluate()`; remove `_detection_valid_fraction()`; update column references | +| `src/anonymizer/engine/replace/replace_runner.py` | Swap `DetectionJudgeWorkflow` → `EntityCoverageWorkflow`; update column references | +| `src/anonymizer/interface/anonymizer.py` | Update `_build_user_dataframe` allowed columns for both modes; update model role references | +| `src/anonymizer/interface/display.py` | Render `entity_coverage` as Satisfied/Not Satisfied (replace) or fraction (rewrite); update leaked_entities dropdown | +| `src/anonymizer/config/models.py` | Rename `detection_validity_judge` → `entity_coverage_judge` in `EvaluateModelSelection` | +| `src/anonymizer/config/default_model_configs/evaluate.yaml` | Rename the model role key | +| `docs/concepts/evaluation.md` | Replace Entity Detection Judge section with Entity Coverage section in both replace and rewrite docs | +| `skills/anonymizer/SKILL.md` | Update evaluation tips and output template column references | +| `tests/engine/evaluation/test_detection_judge.py` | Replace with `test_entity_coverage_judge.py` — new schema, prompt, and workflow tests | +| `tests/engine/test_rewrite_workflow.py` | Update column name references; remove `_detection_valid_fraction` tests | +| `tests/engine/test_evaluate.py` | Update column name assertions | +| `tests/interface/test_display.py` | Update rendering assertions for entity coverage | + +--- + +## Step 1 — New constants (`constants.py`) + +Retire the three detection judge constants and add three replacements: + +```python +# Retired (remove): +# COL_DETECTION_JUDGE = "_detection_judge" +# COL_DETECTION_VALID = "detection_valid" +# COL_DETECTION_INVALID_ENTITIES = "detection_invalid_entities" + +# New: +COL_ENTITY_COVERAGE_JUDGE = "_entity_coverage_judge" # internal raw output +COL_ENTITY_COVERAGE = "entity_coverage" # user-facing float | None +COL_LEAKED_ENTITIES = "leaked_entities" # user-facing list +``` + +--- + +## Step 2 — New workflow (`entity_coverage_judge.py`) + +Replace `detection_judge.py` with a new file. Key differences from the old judge: + +### Output schema + +```python +class LeakedEntity(BaseModel): + value: str + label: str + reasoning: str # one sentence: why this is PII that survived anonymization + +class EntityCoverageSchema(BaseModel): + leaked_entities: list[LeakedEntity] + # No all_valid field — coverage is computed numerically from counts +``` + +### Prompt + +Runs against the **original text** (`COL_TEXT`), with the Anonymizer-detected entity list injected as context. Starting from the leakage eval script prompt and refined through testing. + +Core task phrasing: +- "Given the original text, identify all PII present. Report any identifiers that were not already caught by Anonymizer — these are the missed entities." +- Returns an empty list when Anonymizer caught everything. + +#### Strict entity protection (`Rewrite.strict_entity_protection`) + +Only applies in rewrite mode. When `strict_entity_protection=True`, the rewrite pipeline disallows `leave_as_is`, bans `combined_risk_level: low`, and overrides the MINIMUM NECESSARY CHANGE and quasi-identifier "not automatically protected" principles. The coverage judge must mirror that posture — otherwise a missed quasi-identifier or low-combined-risk entity would be excused by the judge despite the user having demanded full protection. + +The prompt includes a `` block injected when the flag is set (empty string otherwise): + +``` +STRICT PROTECTION MODE IS ENABLED. + +Flag ALL entities as leaked if they were not caught — including quasi-identifiers +and low-risk entities that would normally be given benefit of the doubt. +Do NOT apply MINIMUM NECESSARY CHANGE reasoning to excuse a missed entity. +Do NOT excuse a missed entity because its combined re-identification risk is low. +Any PII span not caught by Anonymizer is a miss in strict mode. +``` + +The block mirrors the equivalent block in `sensitivity_disposition.py` so the two prompts enforce the same standard. In replace mode the call site passes `strict_entity_protection=False` (the field does not exist on `Detect`); the block is omitted and the judge uses its default charitable posture. + +#### Entity type scope (`Detect.entity_labels`) + +When the user configures `Detect(entity_labels=[...])`, detection is intentionally limited to those label types. The coverage judge must be scoped to match — otherwise it would flag missed PII of types the user never asked to detect, artificially deflating the score. + +The prompt includes an `` block populated per-run in `prepare()`: + +- **`entity_labels` is `None`** (default): `"Evaluate for all PII and sensitive entity types."` +- **`entity_labels` is set**: `"Detection was configured to target ONLY these entity types: . Only report missed entities that belong to one of these types. Do NOT flag PII of other types as leaked — those were intentionally excluded from detection."` + +The `` section includes a corresponding rule: respect the entity_type_scope and do not penalize Anonymizer for types outside it. + +Because `entity_labels` is run-level config (not per-row), the workflow constructor accepts `entity_labels: list[str] | None` and `prepare()` broadcasts it as a constant column (`_entity_type_scope`) across all rows. Call sites (`replace_runner.py`, `rewrite_workflow.py`) pass `config.detect.entity_labels`. + +#### Both flags enabled simultaneously + +Both can be active at the same time — there is no config constraint preventing it: + +```python +AnonymizerConfig( + detect=Detect(entity_labels=["first_name", "email_address"]), + rewrite=Rewrite(strict_entity_protection=True), +) +``` + +The two blocks compose cleanly in the judge prompt: the entity type scope narrows *what* the judge looks for, and the strict protection block raises the bar *within* that scope. No contradiction — scope is applied first, strictness within that scope second. The two blocks are independent injections in `prepare()` and do not interfere with each other. + +### Passthrough condition + +Rows with no detected entities trivially score `entity_coverage = 1.0` — nothing to miss, LLM skipped. + +### Coverage fraction (postprocess override) + +```python +def _compute_coverage(self, n_entities_detected: int, n_entities_leaked: int) -> float: + total = n_entities_detected + n_entities_leaked + return 1.0 if total == 0 else n_entities_detected / total +``` + +Passthrough rows get `entity_coverage = 1.0` and `leaked_entities = []` without calling the LLM. + +### Class declaration + +```python +class EntityCoverageWorkflow(_BaseJudgeWorkflow): + RAW_COL = COL_ENTITY_COVERAGE_JUDGE + VALID_COL = COL_ENTITY_COVERAGE + INVALID_COL = COL_LEAKED_ENTITIES + SCHEMA = EntityCoverageSchema + VERDICT_FIELD = "leaked_entities" + DEFAULT_PAYLOAD = {"leaked_entities": []} + MODEL_ROLE = "entity_coverage_judge" + WORKFLOW_NAME = "entity-coverage-judge" +``` + +`postprocess()` is overridden to compute the float fraction from `n_entities_detected` (from `COL_ENTITIES_BY_VALUE`) and `len(leaked_entities)` rather than storing a boolean verdict. + +> **Note:** `COL_ENTITIES_BY_VALUE` (`_entities_by_value`) is derived from `COL_FINAL_ENTITIES` (`final_entities`) during detection — it groups spans by value with a labels list. The leakage eval script references `final_entities`, which is the same underlying data in a different shape. `COL_ENTITIES_BY_VALUE` is the right source here since it's already what the judge's `prepare()` step and passthrough mask consume. + +> **Note on base class / VALID_COL conflict:** `_BaseJudgeWorkflow` does not write to `VALID_COL` before `postprocess()` runs — the base class `evaluate()` calls `postprocess()` as the only step that touches those columns. Overriding `postprocess()` fully in `EntityCoverageWorkflow` is safe. However, `VALID_COL` on the base class carries an implicit `bool | None` semantic (used by other judges). Since `EntityCoverageWorkflow` writes a `float | None` instead, confirm that no shared display or downstream code reads `VALID_COL` generically and assumes boolean before this is merged. + +--- + +## Step 3 — Wire into replace and rewrite evaluate paths + +### Replace (`replace_runner.py`) + +Swap `DetectionJudgeWorkflow` → `EntityCoverageWorkflow`. The judge runs on `COL_TEXT` (original text) — no change to the input column needed. Update column references for the new output columns. + +### Rewrite (`rewrite_workflow.py`) + +Same swap. The judge also runs on `COL_TEXT` (original text), same as replace. Update `evaluate()`: + +- Remove `_detection_valid_fraction()` — no longer needed since `entity_coverage` is already a float from `postprocess()`. +- Update the try/except around the judge call to default `entity_coverage = None` and `leaked_entities = None` on failure. + +--- + +## Step 4 — Model role rename (`models.py`, `evaluate.yaml`) + +```python +class EvaluateModelSelection(BaseModel): + entity_coverage_judge: str # replaces detection_validity_judge + replace_type_fidelity_judge: str + replace_relational_consistency_judge: str + replace_attribute_fidelity_judge: str + rewrite_judge: str +``` + +Update `evaluate.yaml` default to rename the key. Update `model_loader.py` validation for the new role name. + +--- + +## Step 5 — Update `_build_user_dataframe` (`anonymizer.py`) + +Swap old column names for new in the allowed sets for both replace and rewrite mode: + +```python +# Replace (was: COL_DETECTION_VALID, COL_DETECTION_INVALID_ENTITIES) +COL_ENTITY_COVERAGE, +COL_LEAKED_ENTITIES, + +# Rewrite (same swap) +COL_ENTITY_COVERAGE, +COL_LEAKED_ENTITIES, +``` + +--- + +## Step 6 — Display (`display.py`) + +`display_record()` renders the coverage column as a single line. The format puts the verdict first, then the fraction for scale, then the percentage for readability: + +- **Replace:** `Entity Coverage: Satisfied — 21/21 (100%)` / `Entity Coverage: Not Satisfied — 18/21 (86%)` +- **Rewrite:** `Entity Coverage: 0.86 — 18/21 (86%)` (no Satisfied/Not Satisfied badge — rewrite already uses numeric metrics) + +Where the fraction is `n_entities_detected / (n_entities_detected + n_entities_leaked)` and the percentage is derived from it. The fraction gives absolute scale (18/20 and 18/200 look very different); the percentage makes it immediately readable. + +Drop the explanatory subtitle line that the old Detection Judge displayed — entity coverage is self-explanatory and does not need a definition sentence. Drop the `"LLM alignment score:"` label prefix — that framing belonged to the old judge (LLM-vs-LLM alignment), not to a count ratio. + +The `leaked_entities` dropdown replaces `detection_invalid_entities` with the same layout (value, label, reasoning per row). + +--- + +## Step 7 — Docs and skills + +### `docs/concepts/evaluation.md` + +Replace the **Entity Detection Judge** subsection in both Replace and Rewrite sections with an **Entity Coverage** subsection: + +- Explain the recall focus (vs precision of the old judge) +- Document `entity_coverage` type per mode (float in rewrite, Satisfied/Not Satisfied in replace) +- Document `leaked_entities` list shape +- Update the special-values table (passthrough row → `1.0` / Satisfied) +- Update model role name (`entity_coverage_judge`) + +### `skills/anonymizer/SKILL.md` + +- Update the evaluation tips bullet: `entity_coverage_judge` role; `entity_coverage` type per mode; `leaked_entities` dropdown +- Update output template: replace `detection_valid` column references with `entity_coverage` + +--- + +## Step 8 — Tests + +### New test file (`test_entity_coverage_judge.py`) + +``` +test_judge_prompt_runs_on_original_text_column +test_judge_prompt_includes_detected_entity_context +test_entity_coverage_schema_accepts_empty_leaked_list +test_entity_coverage_schema_accepts_leaked_entities +test_coverage_fraction_all_caught # n_entities_leaked=0 → 1.0 +test_coverage_fraction_partial_miss # n_entities_leaked=1, n_entities_detected=3 → 0.75 +test_coverage_fraction_total_miss # n_entities_leaked=n, n_entities_detected=0 → 0.0 +test_evaluate_short_circuits_when_no_entities # passthrough → 1.0, no LLM call +test_evaluate_invokes_adapter_for_rows_with_entities +test_evaluate_merges_entity_and_empty_rows_in_order +test_evaluate_marks_coverage_unavailable_for_malformed_payload +``` + +### Update existing tests + +``` +tests/engine/test_rewrite_workflow.py # rename column refs; remove _detection_valid_fraction tests +tests/engine/test_evaluate.py # column name assertions +tests/interface/test_display.py # entity_coverage rendering +tests/interface/test_anonymizer_interface.py # allowed column set for both modes +``` + +All tests use stub adapters — no real LLM calls. + +--- + +## Implementation Order + +1. Add new constants to `constants.py`; retire old ones +2. Write `entity_coverage_judge.py` — schema, prompt (v1 from leakage_eval script), workflow class, `postprocess()` override +3. Wire into replace evaluate path (`replace_runner.py`) +4. Wire into rewrite evaluate path (`rewrite_workflow.py`); remove `_detection_valid_fraction()` +5. Rename model role in `models.py`, `evaluate.yaml`, `model_loader.py` +6. Update `_build_user_dataframe` allowed columns (`anonymizer.py`) +7. Update `display.py` rendering +8. Update `docs/concepts/evaluation.md` and `skills/anonymizer/SKILL.md` +9. Replace `test_detection_judge.py` with `test_entity_coverage_judge.py`; update all affected tests +10. Run `make format && make typecheck && make test` + +--- + +## Test Datasets + +The following datasets should be used to validate entity coverage results once the feature is shipped. + +| Dataset | Rows | Text column | Ground-truth annotations | Notes | +|---|---|---|---|---| +| **usage_logs_seed** (`test-data/usage_logs_seed.parquet`) | 20 | `conversation_trace` | None | LLM conversation traces with embedded PII (names, emails, phone numbers, addresses) across business document scenarios; good for sanity-checking coverage on realistic, high-density PII text | +| **openPII_ai4_part1** (`openPII_ai4_part1.csv`) | 1000 | `source_text` | `privacy_mask` (character-offset spans with labels: TITLE, GIVENNAME, SURNAME, EMAIL, TELEPHONENUM, DATE, etc.) | Synthetic PII benchmark with ground-truth spans; use `privacy_mask` as the reference to evaluate whether the judge surfaces the same entities Anonymizer missed | +| **PII_NEW_TESTDATA_part1** (`PII_NEW_TESTDATA_part1.csv`) | 1000 | `text_us` / `text_intl` | `spans_us` / `spans_intl` (character-offset spans with Anonymizer-native labels: first_name, last_name, email, date_of_birth, street_address, bank_routing_number, credit_debit_card, etc.) | Multi-domain synthetic dataset (Life, Finance, etc.) with both US and international text variants and pre-tagged versions; spans use Anonymizer's own label vocabulary, making it the most direct test for coverage scoring | + +Run evaluation in both replace and rewrite modes across all three datasets to confirm `entity_coverage` scores are meaningful and `leaked_entities` lists surface real misses rather than noise. `PII_NEW_TESTDATA_part1` is the highest-signal dataset for this feature because its ground-truth spans use the same label set as Anonymizer's detection pipeline. + +--- + +## Open Questions + +- **Prompt v1:** The leakage_eval script is the starting point. The prompt will provide both the original text and the Anonymizer-detected entity list as context — the detected list is what allows the judge to identify the delta efficiently rather than re-detecting everything from scratch. Removing it would mean the judge can't compute a meaningful delta; keeping it is the right call. Once v1 is shipped and results are observed, a follow-up iteration should explore **NER-style prompting** (instructing the judge to perform explicit named-entity recognition passes over the text before computing the delta) as a potential path to improving recall and reducing missed entities. +- **`n_entities_detected` denominator:** Currently planned to count `(value, label)` pairs (flattened, same as current detection judge). If the same value has multiple labels, it counts multiple times. Consider whether to count unique values instead to avoid inflating the denominator. +- **Replace mode threshold:** "Satisfied" is defined as `entity_coverage == 1.0` (strict — even one leaked entity flips to Not Satisfied). Confirm this is the right threshold or whether a small tolerance is acceptable.