feat: tolerate small-model drift on detection + replace schemas#174
feat: tolerate small-model drift on detection + replace schemas#174lipikaramaswamy wants to merge 7 commits into
Conversation
Loosen the LLM-facing detection wire schemas so small models (gemma4-e2b/e4b, nemotron-3-nano:4b, qwen3.5:4b) that drift from the strict contract no longer drop records during validation/augmentation/latent tagging: - RawValidationDecisionSchema / ValidationDecisionSchema: coerce free-form decision prose to keep/reclass/drop, preserve None as "no answer", drop value/label from the wire (re-filled downstream from trusted candidates). - LatentEntitySchema: loosen category/confidence/evidence/rationale with before-validators; truncate verbose rationale instead of failing. - LatentEntitiesSchema + detection_workflow: pad empty latent cells with a sentinel so PyArrow can write the parquet column. - AugmentedEntitySchema: drop value/label min_length; a single empty entry no longer fails the whole augmentation batch (empties are skipped downstream). Adds test_small_model_drift.py (detection drift classes) and updates test_chunked_validation for the slimmed wire shape.
EntityReplacementSchema is the LLM output_format for Substitute mode, so a single drifted entry (small models occasionally emit an empty synthetic, or a blank original/label) used to fail-validate the whole replacement map and drop the record. Drop min_length on original/label/synthetic: - empty original/label cannot match a requested entity and are filtered out by _filter_replacement_map_to_input_entities (valid siblings are preserved); - an empty synthetic keys on (original, label), survives the filter, and is applied as a deletion at rewrite time -- privacy-safe (no PII leak), even if utility-poor. Adds regression tests for both drift modes.
Greptile SummaryThis PR loosens Pydantic wire schemas for detection (
Confidence Score: 5/5Safe to merge; all coercion paths are backed by regression tests pinned to observed small-model drift cases. Each loosening is accompanied by a downstream filter or coercion that preserves the existing semantic contract: None-preserving vs. keep-defaulting decision normalizers are carefully separated, empty replacement entries are filtered before apply, and the parquet sentinel is invisible to readers. The class docstring in src/anonymizer/engine/schemas/detection.py (LatentEntitySchema) has a minor inversion but no functional concern. Important Files Changed
Reviews (5): Last reviewed commit: "fix(detection): wrap bare-string latent ..." | Re-trigger Greptile |
lipikaramaswamy
changed the title
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def _fix(cell): | ||
| if isinstance(cell, dict): | ||
| if not cell.get("latent_entities"): | ||
| return {**cell, "latent_entities": sentinel} | ||
| return cell | ||
| if isinstance(cell, list) and not cell: | ||
| return sentinel | ||
| return cell |
There was a problem hiding this comment.
None/NaN cells returned unchanged will still fail PyArrow
The _fix helper covers dict with empty latent_entities and empty list, but silently passes through any other value — including None or float('nan') that pandas places in a column during partial-failure fallback. PyArrow will still raise a type error on those cells. An early None/NaN guard returning sentinel closes this gap.
| def _fix(cell): | |
| if isinstance(cell, dict): | |
| if not cell.get("latent_entities"): | |
| return {**cell, "latent_entities": sentinel} | |
| return cell | |
| if isinstance(cell, list) and not cell: | |
| return sentinel | |
| return cell | |
| def _fix(cell): | |
| if cell is None or (isinstance(cell, float) and pd.isna(cell)): | |
| return sentinel | |
| if isinstance(cell, dict): | |
| if not cell.get("latent_entities"): | |
| return {**cell, "latent_entities": sentinel} | |
| return cell | |
| if isinstance(cell, list) and not cell: | |
| return sentinel | |
| return cell |
| proposed_label: str | None = Field( | ||
| default="", | ||
| description="Correct label when decision is 'reclass', otherwise empty", | ||
| ) |
There was a problem hiding this comment.
str | None annotation never resolves to None at the Python level
_coerce_proposed_label always converts None → "", so the runtime type of proposed_label is always str. The str | None annotation is intentional for JSON Schema emission (so null passes DD's pre-check), but it should be documented inline — otherwise readers will reasonably wonder whether downstream code needs to guard against None.
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!
_normalize_category referenced LatentCategory.latent_sensitive_attribute, which is not a member of the enum (it only defines latent_identifier; sensitive attributes were folded into quasi_identifier on the rewrite side). Any latent category string containing "sensitive" — exactly the kind of drift small models emit — would raise AttributeError at runtime. Non-canonical category drift now normalizes to latent_identifier, and the field description no longer advertises the nonexistent value. Also clarifies _pad_empty_latent_column's shape-preservation intent (the downstream reader tolerates both dict and bare-list cells, so no mixing occurs). Adds regression tests for the sensitive/unknown category drift paths.
_pad_empty_latent_column previously passed None/NaN cells through unchanged. A row absent from a partial-failure fallback merge (reintroduced as NaN by a pandas reindex) could leave the latent column without an inferable struct schema. None/NaN now normalize to the canonical sentinel struct, consistent with the existing empty-cell handling. Adds TestPadEmptyLatentColumn covering empty-struct, populated, bare-list, None/NaN, and missing-column cases.
Greptile P2: the str | None annotation on ValidationDecisionSchema.proposed_label is intentional for jsonschema null-tolerance; _coerce_proposed_label always normalizes None to "" so the runtime value is a str. Document inline so readers do not add unnecessary None guards.
Live small-model validationThese changes make detection + replacement tolerate small-model output drift without dropping records. Validated live against real small models on Setup
Resilience results — the core goal
Zero record loss at every model size, even though the 2B drifted heavily. Pre-change, a single drifted field could fail-validate an entire batch. Drifted output was always privacy-safe — original PII was replaced, never leaked back. Quality (LLM-as-judge, graded by gpt-oss-120b)
Failures split into two classes:
Replacement-prompt clarification (included here)
|
…al values
Small models echoed the "(e.g. X, Y, Z)" wrapper verbatim as the synthetic
value (e.g. an age replaced with the literal string "(e.g. 52, 38, 45, 31)"),
and "one of: ..." made them copy a canned example value instead. Render the
per-label hints as "such as {examples}" and instruct the model to generate a
NEW value of that kind without reusing the examples or copying the reference
text. Validated live on gemma-3n-e2b: type_fidelity and attribute_fidelity
both clean (4/4), no record drops.
_clamp_evidence returned [] for a bare-string evidence value, silently dropping the quote when a small model emits a single evidence string instead of a one-element list. Treat a bare string as a single-item list, consistent with the drift-tolerance strategy. Addresses Greptile review note on #174. Adds test_bare_string_evidence_wrapped_not_dropped.
|
Addressed the remaining Greptile note (non-blocking) from the review summary in c15896c: The three inline comments (P0 |
1 participant
Summary
First of two stacked PRs reworking the LLM-facing schemas so small models (gemma4-e2b/e4b, nemotron-3-nano:4b, qwen3.5:4b) no longer drop records when their structured output drifts from the strict contract. This PR covers the detection schemas plus the replace (Substitute) wire schema; the rewrite schemas + server-side disposition follow in stacked PR #175.
This supersedes the schema work in #130, rebased onto current
mainand split for reviewability.Detection
RawValidationDecisionSchema/ValidationDecisionSchema: coerce free-formdecisionprose intokeep/reclass/drop, preserveNoneas "no answer" (required by the chunked-merge logic), and dropvalue/labelfrom the wire (re-filled downstream from the trusted candidate lookup byenrich_validation_decisions).LatentEntitySchema: loosencategory/confidence/evidence/rationalewith before-validators; verbose rationales truncate instead of failing.LatentEntitiesSchema+detection_workflow: pad empty latent cells with a sentinel so PyArrow can write the parquet column. (Latent entities are produced here but only consumed in rewrite mode — they land as inert-but-correct columns until feat(rewrite): small-model drift tolerance + server-side disposition #175 reads them.)AugmentedEntitySchema: dropvalue/labelmin_length— a single empty entry no longer fails the whole augmentation batch (empties are already skipped byapply_augmented_entities).Replace
EntityReplacementSchema: dropmin_lengthonoriginal/label/synthetic. This is the LLMoutput_formatfor Substitute mode, so a single drifted entry used to fail-validate the whole replacement map and drop the record. Emptyoriginal/labelare filtered downstream (cannot match a requested entity); an emptysyntheticis applied as a deletion — privacy-safe (no leak), if utility-poor.Test plan
make testgreen (837 passing on this branch)make format-checkcleantest_small_model_drift.py(detection drift classes);test_chunked_validationupdated for the slimmed wire shape; new replace drift tests intest_llm_replace_workflow