train_ner_pii.py

# /// script
# dependencies = [
#   "transformers>=4.48.0",
#   "datasets>=2.20.0",
#   "evaluate>=0.4.0",
#   "seqeval>=1.2.2",
#   "trackio",
#   "numpy<2.0",
#   "accelerate>=0.34.0",
# ]
# ///
"""
ModernBERT PII NER — remapped to 11 company policy labels.

Trains from answerdotai/ModernBERT-base with a new 23-label classification
head.  Fixes the entity-scan alignment bug: instead of reading char_labels
only at real_s (the first non-space position), we now scan the entire token
span [real_s, tok_e) for the first entity character.  This ensures entities
that start after punctuation (e.g. "(Home" or ":John") are correctly labeled
rather than silently dropped.

Run with:  uv run train_ner_pii.py
"""

import json
import numpy as np
import trackio
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    TrainingArguments,
    Trainer,
    DataCollatorForTokenClassification,
    EarlyStoppingCallback,
)
import evaluate

# ── Label mapping: ai4privacy 56 types → 11 policy categories ────────────────
# Labels absent from this dict → treated as O (not redacted)
LABEL_MAP = {
    # PER — person names and titles
    "FIRSTNAME":        "PER",
    "MIDDLENAME":       "PER",
    "LASTNAME":         "PER",
    "PREFIX":           "PER",      # Mr., Dr., Ms.
    # ORG — organisation names
    "COMPANYNAME":      "ORG",
    # EMAIL
    "EMAIL":            "EMAIL",
    # PHONE
    "PHONENUMBER":      "PHONE",
    # ADDRESS — postal / physical address components
    "BUILDINGNUMBER":   "ADDRESS",
    "STREET":           "ADDRESS",
    "SECONDARYADDRESS": "ADDRESS",  # apt / suite
    "CITY":             "ADDRESS",
    "COUNTY":           "ADDRESS",
    "STATE":            "ADDRESS",
    "ZIPCODE":          "ADDRESS",
    # GOV_ID — government-issued identity numbers
    "SSN":              "GOV_ID",
    # FINANCIAL_ID — payment cards, bank & crypto account identifiers
    "CREDITCARDNUMBER": "FINANCIAL_ID",
    "CREDITCARDCVV":    "FINANCIAL_ID",
    "IBAN":             "FINANCIAL_ID",
    "BIC":              "FINANCIAL_ID",
    "BITCOINADDRESS":   "FINANCIAL_ID",
    "ETHEREUMADDRESS":  "FINANCIAL_ID",
    "LITECOINADDRESS":  "FINANCIAL_ID",
    "MASKEDNUMBER":     "FINANCIAL_ID",
    # ACCOUNT_ID — bank / service account identifiers
    "ACCOUNTNAME":      "ACCOUNT_ID",
    "ACCOUNTNUMBER":    "ACCOUNT_ID",
    "USERNAME":         "ACCOUNT_ID",
    # DEVICE_ID — device, vehicle, and network identifiers
    "IP":               "DEVICE_ID",
    "IPV4":             "DEVICE_ID",
    "IPV6":             "DEVICE_ID",
    "MAC":              "DEVICE_ID",
    "PHONEIMEI":        "DEVICE_ID",
    "USERAGENT":        "DEVICE_ID",
    "VEHICLEVIN":       "DEVICE_ID",
    "VEHICLEVRM":       "DEVICE_ID",
    # DATE_OF_BIRTH — birth dates only (generic DATE → O)
    "DOB":              "DATE_OF_BIRTH",
    # CREDENTIALS — authentication secrets
    "PASSWORD":         "CREDENTIALS",
    "PIN":              "CREDENTIALS",
    # Everything else (AGE, DATE, GENDER, JOBAREA, etc.) → O
}

TARGET_LABELS = [
    "PER", "ORG", "EMAIL", "PHONE", "ADDRESS",
    "GOV_ID", "FINANCIAL_ID", "ACCOUNT_ID", "DEVICE_ID",
    "DATE_OF_BIRTH", "CREDENTIALS",
]

label_list = (
    ["O"]
    + sorted(f"B-{l}" for l in TARGET_LABELS)
    + sorted(f"I-{l}" for l in TARGET_LABELS)
)
id2label = {i: l for i, l in enumerate(label_list)}
label2id = {l: i for i, l in id2label.items()}

# ── Config ────────────────────────────────────────────────────────────────────
MODEL_NAME   = "answerdotai/ModernBERT-base"        # train from base
DATASET_NAME = "ai4privacy/pii-masking-200k"
HUB_MODEL_ID = "jefftherover/modernbert-pii-mapped-v2"
OUTPUT_DIR   = "modernbert-pii-mapped-v2"
MAX_LENGTH   = 512

print(f"Labels ({len(label_list)}): {label_list}")

# ── Dataset ───────────────────────────────────────────────────────────────────
print("Loading dataset...")
full = load_dataset(DATASET_NAME, split="train")
en   = full.filter(lambda x: x["language"] == "en")
print(f"English rows: {len(en)}")

splits   = en.train_test_split(test_size=0.1, seed=42)
train_ds = splits["train"]
eval_ds  = splits["test"]
print(f"Train: {len(train_ds)}  Eval: {len(eval_ds)}")

# ── Tokenizer ─────────────────────────────────────────────────────────────────
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# ── Tokenisation + label alignment ────────────────────────────────────────────
# Identical to v6 logic, with LABEL_MAP applied inside make_char_labels.

def make_char_labels(text, raw):
    spans = json.loads(raw) if isinstance(raw, str) else raw
    cl = ["O"] * len(text)
    for span in spans:
        s, e, src_lbl = int(span[0]), int(span[1]), span[2]
        tgt_lbl = LABEL_MAP.get(src_lbl)
        if tgt_lbl is None:
            continue
        for i in range(s, min(e, len(text))):
            cl[i] = f"B-{tgt_lbl}" if i == s else f"I-{tgt_lbl}"
    return cl

def tokenize_and_align(examples):
    enc = tokenizer(
        examples["source_text"],
        truncation=True,
        max_length=MAX_LENGTH,
        return_offsets_mapping=True,
    )
    all_labels = []
    for idx in range(len(examples["source_text"])):
        text    = examples["source_text"][idx]
        cl      = make_char_labels(text, examples["span_labels"][idx])
        offsets = enc["offset_mapping"][idx]
        labels, prev_end = [], None
        for tok_s, tok_e in offsets:
            if tok_s == tok_e:
                labels.append(-100)
                prev_end = None
                continue
            # Space-offset fix: ModernBERT absorbs leading space into token offset
            real_s = tok_s
            while real_s < tok_e and text[real_s] == " ":
                real_s += 1
            is_word_start = prev_end is None or real_s > tok_s
            # Scan the full token span for the first entity character.
            # Fixes the case where an entity begins after punctuation with no
            # preceding space (e.g. "(Home Loan Account") — previously real_s
            # landed on "(" (O) and the entity was silently dropped.
            lbl = "O"
            for c in range(real_s, min(tok_e, len(cl))):
                if cl[c] != "O":
                    lbl = cl[c]
                    break
            if lbl == "O":
                labels.append(label2id["O"] if is_word_start else -100)
            else:
                labels.append(label2id.get(lbl, label2id["O"]))
            prev_end = tok_e
        all_labels.append(labels)
    enc.pop("offset_mapping")
    enc["labels"] = all_labels
    return enc

print("Tokenising datasets...")
cols      = train_ds.column_names
train_tok = train_ds.map(tokenize_and_align, batched=True, remove_columns=cols)
eval_tok  = eval_ds.map(tokenize_and_align,  batched=True, remove_columns=cols)

# ── Metrics ───────────────────────────────────────────────────────────────────
seqeval = evaluate.load("seqeval")

def compute_metrics(p):
    logits, labels = p
    preds = np.argmax(logits, axis=2)
    true_preds  = [[id2label[pp] for pp, ll in zip(pr, la) if ll != -100]
                   for pr, la in zip(preds,  labels)]
    true_labels = [[id2label[ll] for pp, ll in zip(pr, la) if ll != -100]
                   for pr, la in zip(preds,  labels)]
    res = seqeval.compute(predictions=true_preds, references=true_labels)
    return {
        "precision": res["overall_precision"],
        "recall":    res["overall_recall"],
        "f1":        res["overall_f1"],
        "accuracy":  res["overall_accuracy"],
    }

# ── Model ─────────────────────────────────────────────────────────────────────
print("Loading model (ModernBERT-base + new 23-label head)...")
model = AutoModelForTokenClassification.from_pretrained(
    MODEL_NAME,
    num_labels=len(label_list),
    id2label=id2label,
    label2id=label2id,
)

# ── Trackio ───────────────────────────────────────────────────────────────────
trackio.init(project="modernbert-pii-mapped", name="modernbert-pii-mapped-v2")

# ── Training args ─────────────────────────────────────────────────────────────
args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    gradient_accumulation_steps=2,          # effective batch = 32
    learning_rate=5e-5,
    weight_decay=0.01,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine_with_restarts",
    eval_strategy="steps",
    eval_steps=500,
    save_strategy="steps",
    save_steps=500,
    save_total_limit=3,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    greater_is_better=True,
    push_to_hub=True,
    hub_model_id=HUB_MODEL_ID,
    hub_private_repo=False,
    hub_strategy="every_save",
    report_to="trackio",
    run_name="modernbert-pii-mapped-v2",
    fp16=True,
    logging_steps=100,
    dataloader_num_workers=2,
)

# ── Train ─────────────────────────────────────────────────────────────────────
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_tok,
    eval_dataset=eval_tok,
    data_collator=DataCollatorForTokenClassification(tokenizer),
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
)

print("Starting training...")
trainer.train()
trainer.push_to_hub()
print(f"Done! Model pushed to: https://huggingface.co/{HUB_MODEL_ID}")