filename-classifier-mmbert-small-v4

Multi-label filename classifier: given only a filename (no file content), predicts which of 114 sensitive-data categories are likely present. Fine-tuned from jhu-clsp/mmBERT-small.

Why this backbone

ettin-encoder-68m (used by the -v3/-v4 ettin checkpoints) is English-only pretrained despite handling en/de/it/es reasonably well via cross-lingual transfer. mmBERT-small (same lab/ architecture family as ettin) is genuinely MLM-pretrained across 1,833 languages — 140M total params but only 42M non-embedding, comparable compute cost to ettin-68m's 68M, and measured comparable single-example CPU inference latency (15-18ms, only ~17% slower than ettin despite the 2x nominal parameter count, since the difference is almost entirely the larger multilingual embedding table, not additional transformer compute).

Architecture — and why loading it needs modeling.py

Mean-pooled encoder output feeds a single 114-way multi-label classifier head. This is not a plain AutoModelForSequenceClassification checkpoint (transformers can't reconstruct the mean- pooling forward pass from config.json alone), so weights are split as backbone/ (standard HF AutoModel format) + classifier_head.pt, and the model class itself is bundled in this repo as modeling.py — download it alongside the weights (e.g. via snapshot_download, see Usage below). No other repo or private code is required to load or run this model.

Training data

data_v4, fully cleaned: 62,241 train / 7,675 val / 7,785 test records. Two rounds of data-quality fixes are reflected in this training run (earlier releases of this checkpoint were trained before both landed):

1. PERSON-label gaps — names embedded in a filename (as a run-on string, or only inside an email address's local-part) without the PERSON label being set. 3,898 records patched project-wide, 1,135 in this dataset.
2. Unnatural literal-value records removed — 3,819 records whose filename baked in a literal PII value (a full email address, phone number, etc.) rather than referencing the category, a pattern with no realistic production analogue — confirmed 100% confined to one synthetic generator pool. Pure removal, no relabeling, no reshuffling of survivors.

See CyberDesk/filename-classifier-ettin68m-v4's README for the additive-only rebuild methodology these both build on.

Hyperparameters

A 40-trial Optuna sweep was run for this backbone specifically (previous releases reused ettin-v4's own tuned hyperparameters untouched, never verified as optimal for mmBERT). Best trial (#36, val eval_f1=0.9123 during the sweep, on a step-capped budget): task-lr=2.30e-04, batch-size=32, warmup-ratio=0.158, weight-decay=0.062, scheduler=linear, dropout=0.207. Trained to the full 20-epoch early-stopping budget (unlike sweep trials, which are step-capped) — best checkpoint at epoch 9 (early-stopped: no gain ≥1% relative for 3 evals), val eval_f1=0.9171.

Note: re-trained on the same data with the previous release's reused-ettin hyperparameters scores statistically indistinguishably on this dataset (test micro-F1 0.9116 vs. 0.9114 here, a 0.0002 difference) — the sweep's chosen hyperparameters are not a clear win over the naive reuse, likely because sweep trials were step-capped and may have selected for fast convergence under a truncated budget rather than genuine superiority once trained to full convergence. This release is still the one to use going forward on the strength of its clean training-data provenance (see above), not because the hyperparameters themselves are proven better.

Decision threshold

Calibrated against data_v4/val.json (never test), sweep 0.30–0.55. Best: 0.48.

Eval results (data_v4/test.json, n=7,785, threshold=0.48)

Split	P	R	F1 (micro)	macro-F1	Exact match
Overall	0.936	0.888	0.911	0.908	0.814
Spanish only (n=1,249)	0.909	0.843	0.875	0.847	0.672
en/de/it only (n=6,536)	0.943	0.900	0.921	0.923	0.842

For reference, filename-classifier-ettin68m-v3 (no Spanish training, older/smaller test set) scores 0.928 micro-F1 / 0.857 exact-match on its own 100%-en/de/it test set. This checkpoint closes nearly all of that en/de/it gap (0.921 vs. 0.928) while substantially improving on Spanish (v3: 0.477).

Usage

import json
import sys
import torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

local_dir = snapshot_download("CyberDesk/filename-classifier-mmbert-small-v4")
sys.path.insert(0, local_dir)
from modeling import FilenameClassifier  # bundled in this repo, see modeling.py

categories = json.load(open(f"{local_dir}/categories.json"))
tokenizer = AutoTokenizer.from_pretrained(local_dir)
model = FilenameClassifier.load_for_inference(local_dir, num_labels=len(categories))

encoding = tokenizer("q1_2024_ssn_export_backup.csv", return_tensors="pt")
with torch.no_grad():
    logits = model(**encoding)
probs = torch.sigmoid(logits)
predicted = [categories[i] for i, p in enumerate(probs[0]) if p >= 0.48]  # calibrated threshold
print(predicted)