NVIDIA-NeMo / NVIDIA-NeMo/Anonymizer

feat(detection): domain-aware entity labels via early domain classification

Open
#71 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement medium priority
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:

  1. Classify domain (single LLM call, reuse existing DomainClassificationWorkflow)
  2. Look up domain-specific labels from a new DOMAIN_ENTITY_LABELS mapping
  3. Merge with default labels to form the expanded detection label set
  4. 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_labels is 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_LABELS mapping 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_labels override still works (skips domain-specific additions)
  • Replace mode is unaffected unless opted in

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.