PII Mask

Inputs

Outputs

Source

python

inp = __ivy_ctx__["nodes"][__ivy_node_id__]["input"]
import json
import os
import re

# ---- the rules: extension-mlops-nlp/src/data/rules.ts, ported ----
# JavaScript's \b and \d are ASCII and its \s is Unicode white space; ASCII
# mode with WS spelled out reads the same text the same way.
WS = "[\\s\\u00a0\\u1680\\u2000-\\u200a\\u2028\\u2029\\u202f\\u205f\\u3000\\ufeff]"
A = re.ASCII
IBAN_LENGTH = {"ME": 22, "TR": 26, "RS": 22, "BA": 20, "HR": 21, "DE": 22, "AT": 20, "IT": 27, "GB": 22, "FR": 27, "NL": 18, "CH": 21, "SI": 19, "AL": 28, "MK": 19}


def _digits(s):
    return [int(c) for c in s]


def _mod97(numeric):
    r = 0
    for ch in numeric:
        r = (r * 10 + int(ch)) % 97
    return r


def _letters(s):
    return re.sub("[A-Z]", lambda m: str(ord(m.group(0)) - 55), s)


def iban_valid(raw):
    s = re.sub(WS + "+", "", raw).upper()
    if not re.fullmatch(r"[A-Z]{2}\d{2}[A-Z0-9]{10,30}", s, A):
        return False
    want = IBAN_LENGTH.get(s[:2])
    if (len(s) != want) if want else (len(s) < 15 or len(s) > 34):
        return False
    return _mod97(_letters(s[4:] + s[:4])) == 1


def account_valid(raw):
    m = re.fullmatch(r"(\d{3})-(\d{1,13})-(\d{2})", raw.strip(), A)
    if not m:
        return False
    n = m.group(1) + m.group(2).zfill(13) + m.group(3)
    return str(98 - _mod97(n[:16] + "00")).zfill(2) == n[16:]


def jmbg_valid(raw):
    if not re.fullmatch(r"\d{13}", raw, A):
        return False
    day, month = int(raw[0:2]), int(raw[2:4])
    if day < 1 or day > 31 or month < 1 or month > 12:
        return False
    d = _digits(raw[:12])
    s = 7 * (d[0] + d[6]) + 6 * (d[1] + d[7]) + 5 * (d[2] + d[8]) + 4 * (d[3] + d[9]) + 3 * (d[4] + d[10]) + 2 * (d[5] + d[11])
    m = 11 - s % 11
    if m == 10:
        return False
    return (0 if m == 11 else m) == int(raw[12])


def pib_valid(raw):
    if not re.fullmatch(r"\d{8}", raw, A):
        return False
    r = sum(x * (8 - i) for i, x in enumerate(_digits(raw[:7]))) % 11
    return (0 if r <= 1 else 11 - r) == int(raw[7])


def tckn_valid(raw):
    if not re.fullmatch(r"[1-9]\d{10}", raw, A):
        return False
    d = _digits(raw[:9])
    odd = d[0] + d[2] + d[4] + d[6] + d[8]
    even = d[1] + d[3] + d[5] + d[7]
    tenth = (odd * 7 - even) % 10
    eleventh = (sum(d) + tenth) % 10
    return f"{tenth}{eleventh}" == raw[9:]


def vkn_valid(raw):
    if not re.fullmatch(r"\d{10}", raw, A):
        return False
    total = 0
    for i, x in enumerate(_digits(raw[:9])):
        tmp = (x + 9 - i) % 10
        v = (tmp * 2 ** (9 - i)) % 9
        if tmp != 0 and v == 0:
            v = 9
        total += v
    return (10 - total % 10) % 10 == int(raw[9])


def card_valid(raw):
    s = re.sub("(?:" + WS + "|-)", "", raw)
    if not re.fullmatch(r"\d{13,19}", s, A):
        return False
    total = 0
    for i, x in enumerate(reversed(_digits(s[:-1]))):
        v = x * 2 if i % 2 == 0 else x
        total += v - 9 if v > 9 else v
    return (10 - total % 10) % 10 == int(s[-1])


# (name, label, pattern, check, grouped), in the order a tie is settled.
RULES = [
    ("iban", "IBAN", re.compile(r"\b[A-Z]{2}\d{2}[A-Z0-9]{11,30}\b", A), iban_valid, False),
    ("iban", "IBAN", re.compile(r"\b[A-Z]{2}\d{2}(?: [A-Z0-9]{4}){2,7}(?: [A-Z0-9]{1,4})?\b", A), iban_valid, False),
    ("account", "ACCOUNT", re.compile(r"\b\d{3}-\d{1,13}-\d{2}\b", A), account_valid, False),
    ("card", "CARD", re.compile(r"\b(?:\d{4}[ -]){3}\d{1,7}\b|\b\d{15,16}\b", A), card_valid, False),
    ("jmbg", "JMBG", re.compile(r"\b\d{13}\b", A), jmbg_valid, False),
    ("tckn", "TCKN", re.compile(r"\b[1-9]\d{10}\b", A), tckn_valid, False),
    ("vkn", "VKN", re.compile(r"\b(?:VKN|V\.K\.N\.|Vergi (?:Kimlik |)No\.?|Vergi Numarası)" + WS + r"*:?" + WS + r"*(\d{10})\b", A | re.I), vkn_valid, True),
    ("pib", "PIB", re.compile(r"\b(?:PIB|Mati[čc]ni broj)" + WS + r"*:?[\s\S]{0,40}?\b(\d{8})\b", A), pib_valid, True),
    ("pib", "PIB", re.compile(r"\b(\d{8})\b(?=[\s\S]{0,40}?\b(?:PIB|Mati[čc]ni broj))", A), pib_valid, True),
    ("email", "EMAIL", re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b", A), None, False),
    ("phone", "PHONE", re.compile(
        r"(?:\+382|00382)[ /-]?\d{2}[ /-]?\d{3}[ /-]?\d{3,4}\b|\b0[67]\d[ /-]?\d{3}[ /-]?\d{3,4}\b"
        r"|(?:\+90|0090)[ -]?5\d{2}[ -]?\d{3}[ -]?\d{2}[ -]?\d{2}\b|\b05\d{2}[ -]?\d{3}[ -]?\d{2}[ -]?\d{2}\b", A), None, False),
]


def find_by_rules(text):
    """Every value a rule proves in this text, non-overlapping, the longer kept where two meet."""
    found = []
    for name, label, pattern, check, grouped in RULES:
        for m in pattern.finditer(text):
            value = m.group(1) if grouped else m.group(0)
            if value is None:
                continue
            start = m.start() + m.group(0).rfind(value) if grouped else m.start()
            if check and not check(value):
                continue
            found.append({"start": start, "end": start + len(value), "label": label, "origin": "rule:" + name})
    kept = []
    for span in sorted(found, key=lambda s: -(s["end"] - s["start"])):
        if not any(k["start"] < span["end"] and span["start"] < k["end"] for k in kept):
            kept.append(span)
    return sorted(kept, key=lambda s: (s["start"], -s["end"]))


# ---- the model: train_ner.py's windows and decoding, over onnxruntime ----


def load_model(where):
    """A champion.json, a folder holding one, or an exported package folder."""
    if os.path.isdir(where) and os.path.exists(os.path.join(where, "champion.json")):
        where = os.path.join(where, "champion.json")
    if where.endswith(".json"):
        with open(where, encoding="utf-8") as f:
            champion = json.load(f)
        where = os.path.join(os.path.dirname(where), champion["version"])
    with open(os.path.join(where, "ivyx-ner.json"), encoding="utf-8") as f:
        card = json.load(f)
    try:
        import numpy as np
        import onnxruntime as ort
        from tokenizers import Tokenizer
    except ImportError:
        raise RuntimeError("A span model needs onnxruntime and tokenizers; install them with pip install onnxruntime tokenizers.")
    tokenizer = Tokenizer.from_file(os.path.join(where, "tokenizer.json"))
    tokenizer.no_padding()
    tokenizer.enable_truncation(max_length=int(card["maxLength"]), stride=int(card["stride"]))
    session = ort.InferenceSession(os.path.join(where, card.get("model", "model.onnx")), providers=["CPUExecutionProvider"])
    return {"tokenizer": tokenizer, "session": session, "tags": card["tags"], "np": np}


def model_spans(text, model):
    np = model["np"]
    first_enc = model["tokenizer"].encode(text)
    best = {}  # (start, end) -> (distance from a window edge, tag)
    for enc in [first_enc] + list(first_enc.overflowing):
        seq = enc.sequence_ids
        offsets = [tuple(o) if seq[i] is not None else None for i, o in enumerate(enc.offsets)]
        content = [i for i, o in enumerate(offsets) if o is not None and o[0] < o[1]]
        if not content:
            continue
        ids = np.array([enc.ids], dtype=np.int64)
        logits = model["session"].run(["logits"], {"input_ids": ids, "attention_mask": np.ones_like(ids)})[0][0]
        tags = logits.argmax(-1).tolist()
        first, last = content[0], content[-1]
        for i in content:
            edge = min(i - first, last - i)
            seen = best.get(offsets[i])
            if seen is None or edge > seen[0]:
                best[offsets[i]] = (edge, model["tags"][tags[i]])
    spans, current = [], None
    for (s, e), (_, tag) in sorted(best.items()):
        if tag == "O":
            if current:
                spans.append(current)
            current = None
            continue
        kind, label = tag.split("-", 1)
        if current and kind == "I" and current["label"] == label:
            current["end"] = e
        else:
            if current:
                spans.append(current)
            current = {"start": s, "end": e, "label": label, "origin": "model"}
    if current:
        spans.append(current)
    return spans


# ---- the mask ----

text = inp.get("text")
if not isinstance(text, str):
    raise ValueError("text must be a string.")
where = (inp.get("model") or "").strip()
use_rules = inp.get("rules", True) is not False
only = set(inp.get("labels") or [])
if not where and not use_rules:
    raise ValueError("With rules off, name a model; there is nothing else to find spans with.")

spans = model_spans(text, load_model(where)) if where else []
if use_rules:
    ruled = find_by_rules(text)
    spans = [s for s in spans if not any(s["start"] < r["end"] and r["start"] < s["end"] for r in ruled)] + ruled
spans = sorted((s for s in spans if not only or s["label"] in only), key=lambda s: (s["start"], -s["end"]))

placeholders, mapping, counts, kept = {}, {}, {}, []
for s in spans:
    if kept and s["start"] < kept[-1]["end"]:
        continue
    value = text[s["start"]:s["end"]]
    key = (s["label"], value)
    if key not in placeholders:
        counts[s["label"]] = counts.get(s["label"], 0) + 1
        placeholders[key] = "[" + s["label"] + "_" + str(counts[s["label"]]) + "]"
        mapping[placeholders[key]] = value
    kept.append(dict(s, placeholder=placeholders[key]))

pieces, at = [], 0
for s in kept:
    pieces.append(text[at:s["start"]])
    pieces.append(s["placeholder"])
    at = s["end"]
pieces.append(text[at:])

out = __ivy_ctx__["nodes"][__ivy_node_id__]["output"]
out["masked"] = "".join(pieces)
out["mapping"] = mapping
out["spans"] = kept
out["counts"] = counts

Tests

Requires: python:3.9

• statement-line

  A payment line: the IBAN twice, a PIB and an e-mail; the name stays without a model.

• checksum-fails

  A number shaped like an IBAN whose check digits do not hold is left alone.

• only-labels

  Only the labels asked for are masked.

• empty

  An empty text masks nothing.

• nothing-to-find-with

  Rules off and no model is refused with a sentence.