# /// 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}")