v1.4: epoch 3 best checkpoint (F1=0.889, 41 entity types, 241K training examples)

README.md

@@ -1,178 +1,84 @@
 ---
-library_name: transformers
 license: apache-2.0
 language:
 - en
 tags:
-- token-classification
-- ner
 - pii
+- ner
+- token-classification
 - privacy
 - deberta
-- crf
 datasets:
-- ai4privacy/internationalised_pii_dataset
+- ai4privacy/pii-masking-200k
+- nvidia/Nemotron-PII
+- gretelai/synthetic_pii_finance_multilingual
 - gretelai/gretel-pii-masking-en-v1
+metrics:
+- f1
+- precision
+- recall
 pipeline_tag: token-classification
 model-index:
-- name: datafog-pii-small-en
+- name: pii-small-en
   results:
   - task:
       type: token-classification
-      name: Named Entity Recognition
+      name: PII NER
     metrics:
-    - type: f1
-      value: 0.9071
-      name: Overall F1
-    - type: precision
-      value: 0.8981
-      name: Overall Precision
-    - type: recall
-      value: 0.9162
-      name: Overall Recall
+    - name: F1
+      type: f1
+      value: 0.8894
+    - name: Precision
+      type: precision
+      value: 0.8698
+    - name: Recall
+      type: recall
+      value: 0.9098
 ---
 
-# DataFog PII-NER v1.3
+# DataFog PII-Small-EN (v1.4)
 
-A lightweight token classification model for detecting **Personally Identifiable Information (PII)** in English text. Built on DeBERTa-v3-xsmall with character-level CNN features and a CRF decoding head for structured BIO tag prediction.
-
-**v1.3** is the fourth iteration, achieving the best overall F1 (0.9071) across all versions through early backbone freezing and progressive tier weight reduction.
-
-## Model Details
-
-| Property | Value |
-|----------|-------|
-| Architecture | DeBERTa-v3-xsmall + CharCNN + GatingFusion + CRF |
-| Parameters | ~22.7M total |
-| Labels | 89 BIO tags (44 entity types) |
-| Max sequence length | 256 tokens |
-| Training data | ~169K examples from 3 datasets (with Tier 1 oversampling) |
-| Training hardware | NVIDIA H100 PCIe (80GB), BF16 mixed precision |
-| Training time | 20 hours (10 epochs) |
-| Framework | Transformers 4.49, PyTorch 2.7 |
+A compact PII (Personally Identifiable Information) Named Entity Recognition model for English text.
 
 ## Architecture
 
+- **Backbone:** DeBERTa-v3-xsmall (22M params)
+- **Head:** CharCNN + CRF
+- **Total params:** ~45M
+- **Entity types:** 41 PII categories across 4 sensitivity tiers
 
+## Performance (v1.4, Epoch 3)
 
-The **CharCNN** captures structural PII patterns (SSN: XXX-XX-XXXX, credit cards: XXXX-XXXX-XXXX-XXXX) while **DeBERTa** provides contextual understanding. The **gating fusion** dynamically weights character vs. contextual features per token. The **CRF head** enforces valid BIO tag sequences at the sequence level.
-
-## Supported Entity Types (44 types, 4 tiers)
-
-### Tier 1 -- Critical PII (target: 0.98 recall)
-SSN, Credit Card, Bank Account, Passport Number, Drivers License, Tax ID
-
-### Tier 2 -- High Sensitivity (target: 0.95 recall)
-Person, Email, Phone, Date of Birth, Street Address, IP Address
-
-### Tier 3 -- Moderate Sensitivity (target: 0.90 recall)
-Username, Date, Location, Organization, URL, License Plate, Age, Nationality, Gender, Ethnicity, Religion, Marital Status
-
-### Tier 4 -- Domain-Specific (target: 0.85 recall)
-Medical Record, Employee ID, Student ID, Account Number, PIN, Password, Biometric, Vehicle ID, Device ID, Crypto Wallet, IBAN, Swift Code, Insurance Number, Salary, Criminal Record, Political Affiliation, Sexual Orientation, Health Condition, Genetic Data, Trade Union
-
-## Test Set Results
-
-### Overall Metrics
-
-| Metric | V1.3 | V1.2 | V1.1 | V1 |
-|--------|------|------|------|-----|
-| **Overall F1** | **0.9071** | 0.9005 | 0.9005 | 0.904 |
-| Precision | 0.8981 | 0.9050 | 0.9062 | 0.907 |
-| **Recall** | **0.9162** | 0.8960 | 0.8950 | 0.902 |
-
-### Tier Recall
-
-| Tier | V1.3 | V1.2 | Target | Status |
-|------|------|------|--------|--------|
-| Tier 1 (Critical) | 0.823 | 0.841 | 0.98 | FAIL |
-| Tier 2 (High) | **0.945** | 0.936 | 0.95 | FAIL |
-| Tier 3 (Moderate) | **0.930** | 0.911 | 0.90 | PASS |
-| Tier 4 (Domain) | **0.868** | 0.845 | 0.85 | PASS |
-
-### Per-Entity F1 (Top 20)
-
-| Entity Type | F1 |
-|-------------|------|
-| URL | 0.994 |
-| Biometric | 0.992 |
-| IP Address | 0.988 |
-| Date of Birth | 0.981 |
-| Vehicle ID | 0.976 |
-| Email | 0.968 |
-| Phone | 0.966 |
-| License Plate | 0.952 |
-| Gender | 0.946 |
-| Employee ID | 0.940 |
-| IBAN | 0.935 |
-| Username | 0.930 |
-| SSN | 0.930 |
-| Location | 0.929 |
-| Account Number | 0.923 |
-| Organization | 0.902 |
-| Drivers License | 0.881 |
-| Password | 0.880 |
-| Date | 0.877 |
-| Person | 0.875 |
-
-## Training Details
-
-### V1.3 Approach: Early Freeze + Progressive Tier Weights
-
-Two key innovations based on learnings from V1-V1.2:
-
-1. **Backbone freeze after epoch 3**: DeBERTa weights are frozen after epoch 3 to preserve clean representations before training instability occurs.
-
-2. **Progressive tier weight reduction**: CRF loss weights start at 3x/2x/1.5x/1x (Tier 1-4) for epochs 1-2, then reduce to 2x/1.5x/1.25x/1x from epoch 3 onward. This limits gradient amplification buildup while giving a strong initial learning signal.
-
-### Hyperparameters
-
-| Parameter | Value |
-|-----------|-------|
-| Backbone LR | 1e-5 (with AdamW eps=1.0) |
-| Head LR | 1e-3 (100x faster) |
-| LR Schedule | Cosine |
-| Warmup | 500 steps |
-| Epochs | 10 (3 full + 7 head-only) |
-| Effective batch size | 32 (8 x 4 gradient accumulation) |
-| Mixed precision | BF16 |
-| Best checkpoint | Epoch 3 |
-
-### Training Data
-
-~169K examples from three open-licensed datasets:
-- [AI4Privacy PII Dataset](https://huggingface.co/datasets/ai4privacy/internationalised_pii_dataset) (~43K English examples, Apache 2.0)
-- [NVIDIA Nemotron PII](https://huggingface.co/datasets/ai4privacy/pii-masking-400k) (~100K examples, CC-BY-4.0)
-- [Gretel PII Masking](https://huggingface.co/datasets/gretelai/gretel-pii-masking-en-v1) (~26K examples, Apache 2.0)
-
-Tier 1 entity examples are oversampled 3x to address the 323x frequency imbalance between common entities (DATE: 170K) and rare critical entities (PASSPORT: 526).
-
-## Version History
-
-| Version | F1 | Tier 1 Recall | Key Change |
-|---------|------|--------------|------------|
-| V1 | 0.904 | 0.722 | Baseline |
-| V1.1 | 0.9005 | 0.771 | Tier-weighted loss + oversampling |
-| V1.2 | 0.9005 | 0.841 | Backbone freeze after epoch 4 |
-| **V1.3** | **0.907** | 0.823 | Early freeze (epoch 3) + progressive tier weights |
+| Tier | Recall | Target |
+|------|--------|--------|
+| T1 (Critical: SSN, Credit Card, etc.) | 0.814 | >=0.98 |
+| T2 (High: Person, Email, Phone, etc.) | 0.937 | >=0.95 |
+| T3 (Moderate: Username, Date, Location, etc.) | 0.945 | >=0.90 |
+| T4 (Domain-specific: Employee ID, Crypto, etc.) | 0.937 | >=0.85 |
+| **Overall F1** | **0.889** | |
 
-## Limitations
+## Training Data
 
-- Tier 1 recall (0.823) is below the 0.98 target -- critical PII types like Passport Number (only 526 training examples) remain challenging
-- 16 entity types have zero training examples (Nationality, Ethnicity, Religion, etc.) and cannot be detected
-- English-only
-- Max 256 tokens per input (longer documents need chunking)
-- Custom architecture requires the  source code for loading
+- AI4Privacy (~43K examples, English subset)
+- NVIDIA Nemotron-PII (100K examples)
+- Gretel Synthetic PII Finance (26K examples)
+- Gretel PII Masking EN v1 (50K examples)
+- Synthetic data for rare entity types (22K examples)
+- **Total: ~241K examples**
 
-## Links
+## Entity Types (41)
 
-- **Code**: [github.com/DataFog/datafog-labs](https://github.com/DataFog/datafog-labs)
-- **Training Chronicle**: [Full training log](https://github.com/DataFog/datafog-labs/blob/main/pii-ner-v1/docs/training_chronicle.md)
-- **WandB Run**: [V1.3 training metrics](https://wandb.ai/datafog/huggingface/runs/a66aw6sb)
+### Tier 1 — Critical PII
+SSN, CREDIT_CARD, BANK_ACCOUNT, PASSPORT_NUMBER, DRIVERS_LICENSE, TAX_ID
 
-## Citation
+### Tier 2 — High Sensitivity
+PERSON, EMAIL, PHONE, DATE_OF_BIRTH, STREET_ADDRESS, IP_ADDRESS
 
+### Tier 3 — Moderate Sensitivity
+USERNAME, DATE, LOCATION, ORGANIZATION, URL, LICENSE_PLATE, AGE, NATIONALITY, GENDER, RELIGION, MARITAL_STATUS
 
+### Tier 4 — Domain-Specific
+MEDICAL_RECORD, EMPLOYEE_ID, STUDENT_ID, ACCOUNT_NUMBER, PIN, PASSWORD, BIOMETRIC, VEHICLE_ID, DEVICE_ID, CRYPTO_WALLET, IBAN, SWIFT_CODE, INSURANCE_NUMBER, SALARY, CRIMINAL_RECORD, POLITICAL_AFFILIATION, SEXUAL_ORIENTATION, HEALTH_CONDITION
 
 ## License
 

config.json

@@ -17,189 +17,178 @@
   "char_vocab_size": 256,
   "dropout": 0.1,
   "id2label": {
-    "0": "LABEL_0",
-    "1": "LABEL_1",
-    "2": "LABEL_2",
-    "3": "LABEL_3",
-    "4": "LABEL_4",
-    "5": "LABEL_5",
-    "6": "LABEL_6",
-    "7": "LABEL_7",
-    "8": "LABEL_8",
-    "9": "LABEL_9",
-    "10": "LABEL_10",
-    "11": "LABEL_11",
-    "12": "LABEL_12",
-    "13": "LABEL_13",
-    "14": "LABEL_14",
-    "15": "LABEL_15",
-    "16": "LABEL_16",
-    "17": "LABEL_17",
-    "18": "LABEL_18",
-    "19": "LABEL_19",
-    "20": "LABEL_20",
-    "21": "LABEL_21",
-    "22": "LABEL_22",
-    "23": "LABEL_23",
-    "24": "LABEL_24",
-    "25": "LABEL_25",
-    "26": "LABEL_26",
-    "27": "LABEL_27",
-    "28": "LABEL_28",
-    "29": "LABEL_29",
-    "30": "LABEL_30",
-    "31": "LABEL_31",
-    "32": "LABEL_32",
-    "33": "LABEL_33",
-    "34": "LABEL_34",
-    "35": "LABEL_35",
-    "36": "LABEL_36",
-    "37": "LABEL_37",
-    "38": "LABEL_38",
-    "39": "LABEL_39",
-    "40": "LABEL_40",
-    "41": "LABEL_41",
-    "42": "LABEL_42",
-    "43": "LABEL_43",
-    "44": "LABEL_44",
-    "45": "LABEL_45",
-    "46": "LABEL_46",
-    "47": "LABEL_47",
-    "48": "LABEL_48",
-    "49": "LABEL_49",
-    "50": "LABEL_50",
-    "51": "LABEL_51",
-    "52": "LABEL_52",
-    "53": "LABEL_53",
-    "54": "LABEL_54",
-    "55": "LABEL_55",
-    "56": "LABEL_56",
-    "57": "LABEL_57",
-    "58": "LABEL_58",
-    "59": "LABEL_59",
-    "60": "LABEL_60",
-    "61": "LABEL_61",
-    "62": "LABEL_62",
-    "63": "LABEL_63",
-    "64": "LABEL_64",
-    "65": "LABEL_65",
-    "66": "LABEL_66",
-    "67": "LABEL_67",
-    "68": "LABEL_68",
-    "69": "LABEL_69",
-    "70": "LABEL_70",
-    "71": "LABEL_71",
-    "72": "LABEL_72",
-    "73": "LABEL_73",
-    "74": "LABEL_74",
-    "75": "LABEL_75",
-    "76": "LABEL_76",
-    "77": "LABEL_77",
-    "78": "LABEL_78",
-    "79": "LABEL_79",
-    "80": "LABEL_80",
-    "81": "LABEL_81",
-    "82": "LABEL_82",
-    "83": "LABEL_83",
-    "84": "LABEL_84",
-    "85": "LABEL_85",
-    "86": "LABEL_86",
-    "87": "LABEL_87",
-    "88": "LABEL_88"
+    "0": "O",
+    "1": "B-SSN",
+    "2": "I-SSN",
+    "3": "B-CREDIT_CARD",
+    "4": "I-CREDIT_CARD",
+    "5": "B-BANK_ACCOUNT",
+    "6": "I-BANK_ACCOUNT",
+    "7": "B-PASSPORT_NUMBER",
+    "8": "I-PASSPORT_NUMBER",
+    "9": "B-DRIVERS_LICENSE",
+    "10": "I-DRIVERS_LICENSE",
+    "11": "B-TAX_ID",
+    "12": "I-TAX_ID",
+    "13": "B-PERSON",
+    "14": "I-PERSON",
+    "15": "B-EMAIL",
+    "16": "I-EMAIL",
+    "17": "B-PHONE",
+    "18": "I-PHONE",
+    "19": "B-DATE_OF_BIRTH",
+    "20": "I-DATE_OF_BIRTH",
+    "21": "B-STREET_ADDRESS",
+    "22": "I-STREET_ADDRESS",
+    "23": "B-IP_ADDRESS",
+    "24": "I-IP_ADDRESS",
+    "25": "B-USERNAME",
+    "26": "I-USERNAME",
+    "27": "B-DATE",
+    "28": "I-DATE",
+    "29": "B-LOCATION",
+    "30": "I-LOCATION",
+    "31": "B-ORGANIZATION",
+    "32": "I-ORGANIZATION",
+    "33": "B-URL",
+    "34": "I-URL",
+    "35": "B-LICENSE_PLATE",
+    "36": "I-LICENSE_PLATE",
+    "37": "B-AGE",
+    "38": "I-AGE",
+    "39": "B-NATIONALITY",
+    "40": "I-NATIONALITY",
+    "41": "B-GENDER",
+    "42": "I-GENDER",
+    "43": "B-RELIGION",
+    "44": "I-RELIGION",
+    "45": "B-MARITAL_STATUS",
+    "46": "I-MARITAL_STATUS",
+    "47": "B-MEDICAL_RECORD",
+    "48": "I-MEDICAL_RECORD",
+    "49": "B-EMPLOYEE_ID",
+    "50": "I-EMPLOYEE_ID",
+    "51": "B-STUDENT_ID",
+    "52": "I-STUDENT_ID",
+    "53": "B-ACCOUNT_NUMBER",
+    "54": "I-ACCOUNT_NUMBER",
+    "55": "B-PIN",
+    "56": "I-PIN",
+    "57": "B-PASSWORD",
+    "58": "I-PASSWORD",
+    "59": "B-BIOMETRIC",
+    "60": "I-BIOMETRIC",
+    "61": "B-VEHICLE_ID",
+    "62": "I-VEHICLE_ID",
+    "63": "B-DEVICE_ID",
+    "64": "I-DEVICE_ID",
+    "65": "B-CRYPTO_WALLET",
+    "66": "I-CRYPTO_WALLET",
+    "67": "B-IBAN",
+    "68": "I-IBAN",
+    "69": "B-SWIFT_CODE",
+    "70": "I-SWIFT_CODE",
+    "71": "B-INSURANCE_NUMBER",
+    "72": "I-INSURANCE_NUMBER",
+    "73": "B-SALARY",
+    "74": "I-SALARY",
+    "75": "B-CRIMINAL_RECORD",
+    "76": "I-CRIMINAL_RECORD",
+    "77": "B-POLITICAL_AFFILIATION",
+    "78": "I-POLITICAL_AFFILIATION",
+    "79": "B-SEXUAL_ORIENTATION",
+    "80": "I-SEXUAL_ORIENTATION",
+    "81": "B-HEALTH_CONDITION",
+    "82": "I-HEALTH_CONDITION"
   },
   "label2id": {
-    "LABEL_0": 0,
-    "LABEL_1": 1,
-    "LABEL_10": 10,
-    "LABEL_11": 11,
-    "LABEL_12": 12,
-    "LABEL_13": 13,
-    "LABEL_14": 14,
-    "LABEL_15": 15,
-    "LABEL_16": 16,
-    "LABEL_17": 17,
-    "LABEL_18": 18,
-    "LABEL_19": 19,
-    "LABEL_2": 2,
-    "LABEL_20": 20,
-    "LABEL_21": 21,
-    "LABEL_22": 22,
-    "LABEL_23": 23,
-    "LABEL_24": 24,
-    "LABEL_25": 25,
-    "LABEL_26": 26,
-    "LABEL_27": 27,
-    "LABEL_28": 28,
-    "LABEL_29": 29,
-    "LABEL_3": 3,
-    "LABEL_30": 30,
-    "LABEL_31": 31,
-    "LABEL_32": 32,
-    "LABEL_33": 33,
-    "LABEL_34": 34,
-    "LABEL_35": 35,
-    "LABEL_36": 36,
-    "LABEL_37": 37,
-    "LABEL_38": 38,
-    "LABEL_39": 39,
-    "LABEL_4": 4,
-    "LABEL_40": 40,
-    "LABEL_41": 41,
-    "LABEL_42": 42,
-    "LABEL_43": 43,
-    "LABEL_44": 44,
-    "LABEL_45": 45,
-    "LABEL_46": 46,
-    "LABEL_47": 47,
-    "LABEL_48": 48,
-    "LABEL_49": 49,
-    "LABEL_5": 5,
-    "LABEL_50": 50,
-    "LABEL_51": 51,
-    "LABEL_52": 52,
-    "LABEL_53": 53,
-    "LABEL_54": 54,
-    "LABEL_55": 55,
-    "LABEL_56": 56,
-    "LABEL_57": 57,
-    "LABEL_58": 58,
-    "LABEL_59": 59,
-    "LABEL_6": 6,
-    "LABEL_60": 60,
-    "LABEL_61": 61,
-    "LABEL_62": 62,
-    "LABEL_63": 63,
-    "LABEL_64": 64,
-    "LABEL_65": 65,
-    "LABEL_66": 66,
-    "LABEL_67": 67,
-    "LABEL_68": 68,
-    "LABEL_69": 69,
-    "LABEL_7": 7,
-    "LABEL_70": 70,
-    "LABEL_71": 71,
-    "LABEL_72": 72,
-    "LABEL_73": 73,
-    "LABEL_74": 74,
-    "LABEL_75": 75,
-    "LABEL_76": 76,
-    "LABEL_77": 77,
-    "LABEL_78": 78,
-    "LABEL_79": 79,
-    "LABEL_8": 8,
-    "LABEL_80": 80,
-    "LABEL_81": 81,
-    "LABEL_82": 82,
-    "LABEL_83": 83,
-    "LABEL_84": 84,
-    "LABEL_85": 85,
-    "LABEL_86": 86,
-    "LABEL_87": 87,
-    "LABEL_88": 88,
-    "LABEL_9": 9
+    "O": 0,
+    "B-SSN": 1,
+    "I-SSN": 2,
+    "B-CREDIT_CARD": 3,
+    "I-CREDIT_CARD": 4,
+    "B-BANK_ACCOUNT": 5,
+    "I-BANK_ACCOUNT": 6,
+    "B-PASSPORT_NUMBER": 7,
+    "I-PASSPORT_NUMBER": 8,
+    "B-DRIVERS_LICENSE": 9,
+    "I-DRIVERS_LICENSE": 10,
+    "B-TAX_ID": 11,
+    "I-TAX_ID": 12,
+    "B-PERSON": 13,
+    "I-PERSON": 14,
+    "B-EMAIL": 15,
+    "I-EMAIL": 16,
+    "B-PHONE": 17,
+    "I-PHONE": 18,
+    "B-DATE_OF_BIRTH": 19,
+    "I-DATE_OF_BIRTH": 20,
+    "B-STREET_ADDRESS": 21,
+    "I-STREET_ADDRESS": 22,
+    "B-IP_ADDRESS": 23,
+    "I-IP_ADDRESS": 24,
+    "B-USERNAME": 25,
+    "I-USERNAME": 26,
+    "B-DATE": 27,
+    "I-DATE": 28,
+    "B-LOCATION": 29,
+    "I-LOCATION": 30,
+    "B-ORGANIZATION": 31,
+    "I-ORGANIZATION": 32,
+    "B-URL": 33,
+    "I-URL": 34,
+    "B-LICENSE_PLATE": 35,
+    "I-LICENSE_PLATE": 36,
+    "B-AGE": 37,
+    "I-AGE": 38,
+    "B-NATIONALITY": 39,
+    "I-NATIONALITY": 40,
+    "B-GENDER": 41,
+    "I-GENDER": 42,
+    "B-RELIGION": 43,
+    "I-RELIGION": 44,
+    "B-MARITAL_STATUS": 45,
+    "I-MARITAL_STATUS": 46,
+    "B-MEDICAL_RECORD": 47,
+    "I-MEDICAL_RECORD": 48,
+    "B-EMPLOYEE_ID": 49,
+    "I-EMPLOYEE_ID": 50,
+    "B-STUDENT_ID": 51,
+    "I-STUDENT_ID": 52,
+    "B-ACCOUNT_NUMBER": 53,
+    "I-ACCOUNT_NUMBER": 54,
+    "B-PIN": 55,
+    "I-PIN": 56,
+    "B-PASSWORD": 57,
+    "I-PASSWORD": 58,
+    "B-BIOMETRIC": 59,
+    "I-BIOMETRIC": 60,
+    "B-VEHICLE_ID": 61,
+    "I-VEHICLE_ID": 62,
+    "B-DEVICE_ID": 63,
+    "I-DEVICE_ID": 64,
+    "B-CRYPTO_WALLET": 65,
+    "I-CRYPTO_WALLET": 66,
+    "B-IBAN": 67,
+    "I-IBAN": 68,
+    "B-SWIFT_CODE": 69,
+    "I-SWIFT_CODE": 70,
+    "B-INSURANCE_NUMBER": 71,
+    "I-INSURANCE_NUMBER": 72,
+    "B-SALARY": 73,
+    "I-SALARY": 74,
+    "B-CRIMINAL_RECORD": 75,
+    "I-CRIMINAL_RECORD": 76,
+    "B-POLITICAL_AFFILIATION": 77,
+    "I-POLITICAL_AFFILIATION": 78,
+    "B-SEXUAL_ORIENTATION": 79,
+    "I-SEXUAL_ORIENTATION": 80,
+    "B-HEALTH_CONDITION": 81,
+    "I-HEALTH_CONDITION": 82
   },
   "max_char_len": 20,
   "model_type": "pii_ner",
   "torch_dtype": "float32",
-  "transformers_version": "4.49.0"
-}
+  "transformers_version": "4.45.2",
+  "num_labels": 83
+}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:de2fdce39e2098e4936063aa68caa580739a451bfdb7988298d5870df64db5f9
-size 284508924
+oid sha256:8dbbaa83a0b63f307452ad2f790640d8e3d42aee568feb806f7229e250616808
+size 284495484