NVIDIA-NeMo / NVIDIA-NeMo/Anonymizer
feat(detection): domain-aware entity labels via early domain classification
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 122
- Forks
- 17
- Avg merge
- 4d 2h
- Merged PRs (30d)
- 11
Description
Context
Entity detection currently uses a fixed default label list (55 labels) regardless of the input domain. Domain-specific entities like court_name, case_number, prison_detention_facility (legal) or diagnosis, procedure, medication (medical) are missed in the first detection pass. This causes downstream issues:
- The privacy judge penalizes for entities it doesn't recognize as having been handled
- Latent entity detection has to rediscover entities that should have been caught explicitly
- Sensitivity disposition operates on an incomplete entity set
From e2e testing feedback: adding domain-specific entity labels to legal data improved privacy judge scores (6.65 vs 5.67 with defaults).
Proposal
Move domain classification into the entity detection pipeline, before GLiNER runs:
- Classify domain (single LLM call, reuse existing
DomainClassificationWorkflow) - Look up domain-specific labels from a new
DOMAIN_ENTITY_LABELSmapping - Merge with default labels to form the expanded detection label set
- Run GLiNER + validation + augmentation with the expanded labels
The domain classification result is stored on the dataframe and reused by the rewrite pipeline's pre-generation step (no duplicate LLM call).
Design considerations
- Domain classification runs on the raw text (no entity tags needed), so it can precede detection
- For replace mode, domain classification is an extra LLM call that wasn't there before. Consider making it opt-in for replace, always-on for rewrite
- When
entity_labelsis explicitly provided by the user, domain-specific labels should NOT be added (user knows what they want) - The domain column should be carried through to the rewrite pipeline so the evaluate-repair loop can use it without re-classifying
- Also add broadly useful labels to the default list directly:
nationality,court_name,prison_detention_facility
Acceptance criteria
- Domain classification runs as the first step of detection when rewrite mode is enabled
-
DOMAIN_ENTITY_LABELSmapping exists with at least legal, medical, and biography domains - Detection uses merged (default + domain-specific) labels for GLiNER, validation, and augmentation
- Domain column is available to the rewrite pipeline without a second classification call
-
entity_labelsoverride still works (skips domain-specific additions) - Replace mode is unaffected unless opted in
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with DomainClassificationWorkflow and the existing entity detection pipeline, then trace how GLiNER, validation, augmentation, and the rewrite pipeline exchange dataframe columns. Define the DOMAIN_ENTITY_LABELS mapping and verify the acceptance criteria, especially explicit entity_labels overrides and replace-mode behavior; done means domain classification is reused without a second call and merged labels reach every detection stage.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100