Skip to content

Commit ceb94e2

Browse files
Sid Mohanclaude
Sid Mohan
and
claude
committed
Rewrite README with results, HuggingFace link, and development log
Surface v1.0-v1.3 training results, link to DataFog/pii-small-en checkpoint, add dated development log, documentation index, key technical findings, and open problems. Gitignore output-*/ dirs and Windows nul artifact. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 3389eb3 commit ceb94e2

2 files changed

Lines changed: 146 additions & 19 deletions

File tree

‎.gitignore‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,14 @@ Thumbs.db
2929
*.safetensors
3030
wandb/
3131
output/
32+
output-*/
3233
outputs/
3334
checkpoints/
3435
runs/
3536

37+
# Windows artifacts
38+
nul
39+
3640
# Data (downloaded datasets)
3741
data/raw/
3842
data/processed/

‎README.md‎

Lines changed: 142 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,74 @@
11
# DataFog Labs
22

3-
Research and development for DataFog's PII detection models.
3+
Open research and development for lightweight PII detection models. This repo contains the full training code, experiment history, and research behind [DataFog](https://datafog.ai)'s PII-NER model family.
44

5-
## Projects
5+
**Latest checkpoint:** [DataFog/pii-small-en](https://huggingface.co/DataFog/pii-small-en) on HuggingFace (v1.3, Apache 2.0)
66

7-
### [PII-NER v1](pii-ner-v1/)
7+
## PII-NER v1
88

9-
A 22.7M parameter PII detection model combining DeBERTa-v3-xsmall with a character CNN encoder, sigmoid gating fusion, and CRF output layer. Trained on 360K+ open-licensed PII examples.
9+
A 22.7M parameter model for detecting 44 types of personally identifiable information in English text. Combines a pretrained DeBERTa-v3-xsmall backbone with a character CNN encoder, adaptive gating fusion, and CRF output layer.
1010

11-
**Architecture:** DeBERTa-v3-xsmall (contextual) + CharCNN (character patterns) → Gating Fusion → CRF (constrained BIO decoding)
11+
```
12+
Input Text
13+
|
14+
[Tokenization + Word-to-Char mapping]
15+
|
16+
DeBERTa-v3-xsmall (22M) + CharCNN (0.3M)
17+
| |
18+
+-------> Gating Fusion <----+
19+
|
20+
CRF Head (0.2M)
21+
|
22+
BIO Tag Predictions (Viterbi decode)
23+
|
24+
Span-level PII Entities
25+
```
26+
27+
The gating fusion dynamically weights character-level features (for structured PII like SSNs and credit cards) against contextual features (for soft PII like names and addresses) on a per-token basis.
28+
29+
### Results
30+
31+
Best results from v1.3 training on H100 (20 hours, 10 epochs):
1232

13-
**Key features:**
14-
- 163 BIO tags covering ~81 PII entity types across 4 sensitivity tiers
15-
- Adaptive gating between character-level patterns (for structured PII like SSNs) and contextual features (for soft PII like names)
16-
- CRF output layer enforcing valid BIO sequences
17-
- Differential learning rates for pretrained backbone vs randomly initialized head
33+
| Metric | V1.0 | V1.1 | V1.2 | V1.3 |
34+
|--------|------|------|------|------|
35+
| **Overall F1** | 0.904 | 0.901 | 0.901 | **0.907** |
36+
| Precision | 0.907 | 0.906 | 0.905 | 0.898 |
37+
| Recall | 0.902 | 0.895 | 0.896 | **0.916** |
38+
| Tier 1 Recall (SSN, Credit Card, ...) | 0.722 | 0.771 | **0.841** | 0.823 |
39+
| Tier 2 Recall (Person, Email, Phone, ...) | 0.934 | 0.933 | 0.936 | **0.945** |
40+
| Tier 3 Recall (Username, Date, Location, ...) | 0.919 | 0.908 | 0.911 | **0.930** |
41+
| Tier 4 Recall (Employee ID, IBAN, ...) | 0.866 | 0.844 | 0.845 | **0.868** |
1842

19-
**Status:** Smoke test complete (F1=0.947 on 100 examples). Full training pending.
43+
V1.3 has the best overall F1 and recall. V1.2 is better for Tier 1-critical deployments (0.841 vs 0.823).
44+
45+
### Top entity F1 scores (v1.3)
46+
47+
| Entity | F1 | Entity | F1 |
48+
|--------|------|--------|------|
49+
| URL | 0.994 | License Plate | 0.952 |
50+
| Biometric | 0.992 | Gender | 0.946 |
51+
| IP Address | 0.988 | Employee ID | 0.940 |
52+
| Date of Birth | 0.981 | IBAN | 0.935 |
53+
| Vehicle ID | 0.976 | Username | 0.930 |
54+
| Email | 0.968 | SSN | 0.930 |
55+
| Phone | 0.966 | Location | 0.929 |
56+
57+
### Quick start
58+
59+
```python
60+
from datafog_pii_ner.inference import PiiPipeline
61+
62+
pipeline = PiiPipeline.from_pretrained("DataFog/pii-small-en")
63+
entities = pipeline("My SSN is 123-45-6789 and email is john@example.com")
64+
# [PiiEntity(text='123-45-6789', label='SSN', start=10, end=21, tier=1),
65+
# PiiEntity(text='john@example.com', label='EMAIL', start=32, end=48, tier=2)]
66+
```
67+
68+
### Setup
2069

2170
```bash
2271
cd pii-ner-v1
23-
24-
# Install
2572
pip install -e ".[dev]"
2673

2774
# Run tests
@@ -32,19 +79,95 @@ ruff check src/ tests/ scripts/
3279

3380
# Smoke test (requires GPU)
3481
python -m scripts.smoke_test
82+
83+
# Full training
84+
python scripts/train_v1.3.py --config configs/h100-v1.3.yaml
3585
```
3686

37-
### [Research](RESEARCH/)
87+
### Evaluation
3888

39-
Architecture survey, training data catalog, evaluation framework, and gap analysis for PII-NER model design. Includes a 29-slide interactive architecture guide.
89+
```bash
90+
python scripts/eval_benchmark.py \
91+
--model datafog \
92+
--model-path DataFog/pii-small-en \
93+
--dataset combined \
94+
--split test
95+
```
96+
97+
See [eval_benchmark.md](pii-ner-v1/docs/eval_benchmark.md) for flags and options.
4098

41-
## Data Sources
99+
## Training data
42100

43101
| Dataset | Size | License |
44102
|---------|------|---------|
45-
| [AI4Privacy](https://huggingface.co/datasets/ai4privacy/pii-masking-200k) | ~200K | Apache 2.0 |
46-
| [NVIDIA Nemotron-PII](https://huggingface.co/datasets/nvidia/Nemotron-PII) | ~100K | CC-BY-4.0 |
47-
| [Gretel Synthetic PII Finance](https://huggingface.co/datasets/gretelai/synthetic_pii_finance_multilingual) | ~56K | Apache 2.0 |
103+
| [AI4Privacy](https://huggingface.co/datasets/ai4privacy/pii-masking-200k) | ~200K examples | Apache 2.0 |
104+
| [NVIDIA Nemotron-PII](https://huggingface.co/datasets/nvidia/Nemotron-PII) | ~100K examples | CC-BY-4.0 |
105+
| [Gretel Synthetic PII Finance](https://huggingface.co/datasets/gretelai/synthetic_pii_finance_multilingual) | ~56K examples | Apache 2.0 |
106+
107+
Combined: ~169K English examples after filtering and dedup. 44 canonical entity types across 4 sensitivity tiers, unified into 89 BIO labels. The dataset has a 323x frequency imbalance (DATE: 170K occurrences vs PASSPORT: 526), which drives many of the training innovations below.
108+
109+
## Documentation
110+
111+
| Document | Description |
112+
|----------|-------------|
113+
| [Training Chronicle](pii-ner-v1/docs/training_chronicle.md) | Full narrative of the ML journey: 4 NaN sources, backbone instability, tier-weighted loss, freezing experiments |
114+
| [Smoke Test Walkthrough](pii-ner-v1/docs/smoke_test_walkthrough.md) | Why differential learning rates are essential for pretrained+CRF architectures |
115+
| [Evaluation Harness](pii-ner-v1/docs/eval_benchmark.md) | Head-to-head model comparison on the same test split |
116+
| [Design Document](docs/plans/2026-02-01-pii-ner-v1-design.md) | Original architecture decisions and project structure |
117+
118+
### Research
119+
120+
The [RESEARCH/](RESEARCH/pii-model-architecture/) directory contains the pre-implementation research: 8 reports surveying 26 architectures, 9 PII datasets, and the competitive landscape. Includes a [29-slide interactive architecture guide](RESEARCH/pii-model-architecture/architecture-guide.html).
121+
122+
Key finding: no published work combines differentiable character-level pattern recognition with contextual transformers specifically for PII detection.
123+
124+
## Development log
125+
126+
| Date | Version | What changed |
127+
|------|---------|-------------|
128+
| 2026-02-07 | **v1.3** | Best F1 (0.907). Early backbone freeze (epoch 3) + progressive tier weight reduction. Discovered training spikes originate in head components, not backbone. |
129+
| 2026-02-05 | v1.2 | Best Tier 1 recall (0.841). Backbone freezing after epoch 4. Epoch 3 identified as consistent sweet spot. |
130+
| 2026-02-04 | v1.1 | Tier-weighted CRF loss (3x for Tier 1), rare entity oversampling, inference pipeline. Tier 1 recall +4.9pts. |
131+
| 2026-02-04 | — | Training chronicle, entity frequency audit (323x imbalance discovered). |
132+
| 2026-02-03 | v1.0 | First full training on A100. F1=0.904 on 360K examples. Model uploaded to HuggingFace. |
133+
| 2026-02-03 | — | NaN gauntlet: 4 distinct NaN sources identified and fixed (CRF overflow, AdamW bias-correction, BF16 mantissa, FP16 gradient scaler). |
134+
| 2026-02-02 | — | Smoke test passed (F1=0.947 on 100 examples). Differential learning rates proven essential. |
135+
| 2026-02-01 | — | Architecture design. Research phase complete (2,800+ lines across 8 reports). |
136+
137+
## Key technical findings
138+
139+
1. **Differential learning rates are non-negotiable.** A flat LR across pretrained backbone + random CRF head produces F1=0.000. A 50x ratio (backbone 2e-5, head 1e-3) is needed.
140+
141+
2. **AdamW eps=1.0 for pretrained backbones.** Standard eps=1e-8 makes effective updates ~±lr regardless of gradient magnitude, causing NaN on DeBERTa with PyTorch 2.9+. Setting eps=1.0 restores gradient-proportional updates.
142+
143+
3. **The training spike is a head problem, not backbone.** V1.3 proved this definitively: the spike occurred at epoch 5 with the backbone already frozen since epoch 3. The CharCNN/GatingFusion/CRF destabilize under continued training.
144+
145+
4. **Epoch 3 is consistently the best checkpoint.** Across v1.2 and v1.3, the model peaks at epoch 3 then destabilizes. Earlier representations generalize better.
146+
147+
5. **Tier-weighted loss works but amplifies instability.** 3x weight + 3x oversampling = ~9x gradient signal for Tier 1, which accelerates learning but accumulates damage.
148+
149+
## Open problems
150+
151+
- **Tier 1 recall gap**: 0.823 vs 0.98 target. Passport number (0.426 F1) has only 526 training examples.
152+
- **16 zero-occurrence entity types**: NATIONALITY, ETHNICITY, RELIGION, etc. exist in the taxonomy but no training data covers them.
153+
- **Head instability**: Root cause of the epoch 3+ training spike is unknown. Gradient clipping, per-component LR decay, or early stopping are candidate fixes.
154+
- **ONNX export**: CRF Viterbi decode doesn't export cleanly; needs pure-PyTorch reimplementation.
155+
156+
## Project structure
157+
158+
```
159+
datafog-labs/
160+
├── pii-ner-v1/
161+
│ ├── src/datafog_pii_ner/ # Model, data pipeline, training, inference
162+
│ ├── scripts/ # Training runners, evaluation, data download
163+
│ ├── configs/ # YAML configs per GPU/version
164+
│ ├── tests/ # Unit + integration tests (6 modules)
165+
│ ├── notebooks/ # Experiment notebooks (Colab/local)
166+
│ └── docs/ # Training chronicle, eval docs
167+
├── RESEARCH/ # Pre-implementation research (8 reports)
168+
├── docs/plans/ # Design documents
169+
└── .github/workflows/ci.yml # Lint + test CI
170+
```
48171

49172
## License
50173

0 commit comments

Comments
 (0)