byt5-pii-masking-ocr

byt5-pii-masking-ocr

A 582M parameter, byte-level (google/byt5-base) model fully fine-tuned for two tasks:

1. PII masking — replaces real personal information in free text with [LABEL] tags, robust to OCR noise (character-level errors from scanned documents).
2. PII column identification — given a table's column headers + sample values, identifies which columns contain PII and their type.

Byte-level modeling was chosen deliberately: OCR errors are character-level corruptions (0↔O, 1↔l, dropped/duplicated characters), and a byte-level model sees and can correct these directly, unlike subword-tokenized models.

Usage

Task 1 — mask PII in text

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tok = AutoTokenizer.from_pretrained("rachitpandey26/byt5-pii-masking-ocr")
model = AutoModelForSeq2SeqLM.from_pretrained("rachitpandey26/byt5-pii-masking-ocr")

text = "Dear Omer, call 0151 555 0192 or email [email protected]."
inputs = tok(text, return_tensors="pt")
out = model.generate(**inputs, max_length=1024, num_beams=1)
print(tok.decode(out[0], skip_special_tokens=True))
# -> "Dear [FIRSTNAME], call [PHONENUMBER] or email [EMAIL]."

Task 2 — identify PII columns in a table

text = (
    "[TABLE]\n"
    "[HDR] name | order_id | email | city\n"
    "[ROW] Omer | 4471 | [email protected] | Berlin\n"
    "[ROW] Sara | 9982 | [email protected] | Lyon\n"
    "[TASK] identify PII columns\n"
    "[ANSWER]"
)
inputs = tok(text, return_tensors="pt")
out = model.generate(**inputs, max_length=256, num_beams=1)
print(tok.decode(out[0], skip_special_tokens=True))
# -> "name: FIRSTNAME\nemail: EMAIL"   (non-PII columns omitted; "[NONE]" if none found)

Evaluation (held-out test set)

Benchmark numbers not available at push time.

Training data

• Task 1 (masking): English subset of ai4privacy/pii-masking-300k, filtered to a curated set of genuine PII labels (names, contact info, financial/government IDs, digital identifiers). Non-PII labels (job title, city, company, etc.) are deliberately not masked, to avoid over-masking ordinary text. ~50% of inputs have synthetic OCR noise injected (character confusables, dropped/duplicated characters) with clean targets, so the model learns to correct errors, not just copy clean text.
• Task 2 (column ID): synthesized fake tables mixing real harvested PII values, real non-PII values from the same dataset (job title, city, company, etc.), and adversarial look-alike columns (formatted codes that resemble PII but aren't) — teaching the model to discriminate on real signal rather than superficial formatting.

Intended use & limitations

• Intended for augmenting PII-scrubbing pipelines (e.g. OCR post-processing for scanned documents, table/column PII discovery), not as a sole compliance control.
• Trained on English, synthetic/semi-synthetic data. Real-world PII formats, languages, or document types outside this distribution may see reduced accuracy — validate on your own data before production use.
• Not guaranteed to catch 100% of PII; always combine with human review for compliance-critical applications.

License

MIT