Upload train_ner_pii.py with huggingface_hub

train_ner_pii.py

@@ -0,0 +1,266 @@
+# /// 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}")