pii_hybrid_decode.py

"""Self-contained hybrid decode for LiquidAI/pii-detect (v7).

The token-classification head locates PII but, like all byte-BPE token classifiers,
fragments the boundaries of format-bound entities (e.g. it tags `1969` inside a date,
or `charite.de` inside an email). This module adds an inference-time regex layer — the
decode the product is meant to use — which roughly DOUBLES exact-match F1 with no loss
of precision/recall on real text:

  AUTH types  : distinctive, validator-gated formats (email, IBAN, JWT, SSN, MAC, crypto,
                api_key, private_key, connection_string, ip, url, credit_card, swift, imei,
                gps). Regex ADDS these and owns their exact boundaries.
  SNAP types  : FP-prone formats (phone, date_of_birth, amount, postal_code). The MODEL
                must fire; regex only EXPANDS its fragment to the full match (no new FPs).

Everything else (names, addresses, conditions, medications, org, special-category,
username, national_id, passport, etc.) is left to the model.

Usage:
    import torch
    from transformers import AutoTokenizer, AutoModelForTokenClassification
    from pii_hybrid_decode import predict
    tok = AutoTokenizer.from_pretrained("LiquidAI/pii-detect", trust_remote_code=True)
    model = AutoModelForTokenClassification.from_pretrained("LiquidAI/pii-detect",
                trust_remote_code=True).eval()
    spans = predict("Email laura@charite.de or call +49 30 4505 1234.", tok, model)
    # -> [{'start':6,'end':22,'type':'contact.email','text':'laura@charite.de'}, ...]
"""
from __future__ import annotations
import re

def _luhn_ok(num: str) -> bool:
    ds = [int(c) for c in num if c.isdigit()]
    if not (12 <= len(ds) <= 19): return False
    tot, par = 0, len(ds) % 2
    for i, d in enumerate(ds):
        if i % 2 == par:
            d *= 2; d = d - 9 if d > 9 else d
        tot += d
    return tot % 10 == 0

def _iban_ok(s: str) -> bool:
    s = s.replace(" ", "").upper()
    if not re.fullmatch(r"[A-Z]{2}\d{2}[A-Z0-9]{11,30}", s): return False
    r = s[4:] + s[:4]
    return int("".join(str(int(c, 36)) for c in r)) % 97 == 1

# (type, pattern, validator) — distinctive formats the regex layer ADDS + owns boundaries
_AUTH = [
    ("contact.email", re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b"), None),
    ("credential.jwt", re.compile(r"\beyJ[A-Za-z0-9_\-]+\.eyJ[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+"), None),
    ("credential.api_key", re.compile(r"\b(?:AKIA[0-9A-Z]{16}|sk-(?:proj-)?[A-Za-z0-9]{20,}|sk-ant-api03-[A-Za-z0-9_\-]{20,}|ghp_[A-Za-z0-9]{36}|AIza[0-9A-Za-z_\-]{35}|xox[baprs]-[A-Za-z0-9\-]{10,}|hf_[A-Za-z0-9]{30,})\b"), None),
    ("credential.private_key", re.compile(r"-----BEGIN (?:RSA |EC |OPENSSH |PGP )?PRIVATE KEY-----"), None),
    ("credential.connection_string", re.compile(r"\b(?:postgres(?:ql)?|mysql|mongodb(?:\+srv)?|redis|amqp)://[^\s:@/]+:[^\s:@/]+@[^\s/]+"), None),
    ("financial.iban", re.compile(r"\b[A-Z]{2}\d{2}(?:[ ]?[A-Z0-9]{4}){2,7}[ ]?[A-Z0-9]{1,3}\b"), _iban_ok),
    ("financial.crypto_wallet", re.compile(r"\b(?:0x[a-fA-F0-9]{40}|bc1[a-z0-9]{25,90}|[13][a-km-zA-HJ-NP-Z1-9]{25,34})\b"), None),
    ("device.mac_address", re.compile(r"\b(?:[0-9A-Fa-f]{2}[:\-]){5}[0-9A-Fa-f]{2}\b"), None),
    ("location.gps_coordinates", re.compile(r"[\-+]?\d{1,3}\.\d{3,}\s*,\s*[\-+]?\d{1,3}\.\d{3,}"), None),
    ("online.url", re.compile(r"\bhttps?://[^\s]+"), None),
    ("identity.ssn", re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), None),
    ("contact.ip_address", re.compile(r"\b(?:(?:25[0-5]|2[0-4]\d|1?\d?\d)\.){3}(?:25[0-5]|2[0-4]\d|1?\d?\d)\b"), None),
    ("financial.credit_card", re.compile(r"\b(?:\d[ \-]?){13,19}\b"), _luhn_ok),
    ("financial.swift_bic", re.compile(r"\b[A-Z]{4}[A-Z]{2}[A-Z0-9]{2}(?:[A-Z0-9]{3})?\b"), None),
    ("device.imei", re.compile(r"\b\d{15}\b"), _luhn_ok),
]
_SNAP = {
    "contact.phone": re.compile(r"(?<!\d)(?:\+?\d{1,3}[ \-.]?)?(?:\(\d{2,4}\)[ \-.]?)?\d{2,4}[ \-.]?\d{3}[ \-.]?\d{3,4}(?!\d)"),
    "identity.date_of_birth": re.compile(r"\b(?:\d{1,2}[\/.\-]\d{1,2}[\/.\-]\d{2,4}|\d{4}[\/.\-]\d{1,2}[\/.\-]\d{1,2}|(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\.?\s+\d{1,2},?\s+\d{4}|\d{1,2}\s+(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\.?\s+\d{4})\b"),
    "financial.amount": re.compile(r"(?:[$€£¥]\s?\d[\d.,]*(?:\s?[KMB])?|\b(?:USD|EUR|GBP|JPY|CHF|CAD|AUD)\s?\d[\d.,]*(?:\s?[KMB])?\b|\b\d[\d.,]*\s?(?:USD|EUR|GBP|dollars|euros)\b)"),
    "contact.postal_code": re.compile(r"\b(?:\d{5}(?:-\d{4})?|[A-Z]{1,2}\d[A-Z\d]?\s?\d[A-Z]{2})\b"),
}
_AUTH_TYPES = {t for t, _, _ in _AUTH}

# boundary-snap: the token classifier drops leading/trailing sub-tokens of Latin-script
# words ('Ibuprofen'->'profen'). Extend a span across contiguous Latin word-chars to
# complete the partial word(s). Latin-only -> CJK/Arabic left untouched (no over-extend).
_LAT = re.compile(r"[0-9A-Za-zÀ-ÖØ-öø-ÿ]")
def _snap_word(text, s, e):
    n = len(text)
    while s > 0 and _LAT.match(text[s - 1]) and _LAT.match(text[s]): s -= 1
    while e < n and _LAT.match(text[e]) and _LAT.match(text[e - 1]): e += 1
    return s, e
_SWIFT_CUE = re.compile(r"(?i)(swift|bic)")
def _swift_ok(text, s, val):  # kill ALL-CAPS-word false BICs (PARTICULARS, CONFIDENTIAL)
    return bool(re.search(r"\d", val)) or bool(_SWIFT_CUE.search(text[max(0, s - 12):s]))

def hybrid_spans(text: str, model_spans: list[dict]) -> list[dict]:
    """model_spans: [{'start','end','type'}...] from the token classifier. Returns the
    hybrid-decoded spans (dicts with start/end/type/text)."""
    # 1. AUTH regex spans — built INDEPENDENTLY of the model (regex is authoritative
    #    for these distinctive formats; model fragments must not block them).
    auth, claimed = [], [False] * len(text)
    for t, pat, val in _AUTH:
        for mm in pat.finditer(text):
            s, e = mm.start(), mm.end()
            if any(claimed[s:e]): continue
            if val and not val(mm.group(0)): continue
            if t == "financial.swift_bic" and not _swift_ok(text, s, mm.group(0)): continue
            for i in range(s, e): claimed[i] = True
            auth.append({"start": s, "end": e, "type": t, "text": mm.group(0)})
    # 2. model spans for non-AUTH types; SNAP types expand to overlapping regex match
    out = []
    for m in model_spans:
        t = m["type"]
        if t in _AUTH_TYPES:
            continue  # regex owns these
        if t in _SNAP:
            snap = None
            for mm in _SNAP[t].finditer(text):
                if min(mm.end(), m["end"]) > max(mm.start(), m["start"]):
                    snap = mm; break
            if snap:
                out.append({"start": snap.start(), "end": snap.end(), "type": t,
                            "text": text[snap.start():snap.end()]}); continue
        ss, ee = _snap_word(text, m["start"], m["end"])  # complete partial Latin words
        out.append({"start": ss, "end": ee, "type": t, "text": text[ss:ee]})
    out.extend(auth)
    # non-Latin scripts (CJK/Hangul/Hiragana/Katakana/Thai/Arabic/Cyrillic/Devanagari/Hebrew):
    # a 2-char span is a full token (e.g. the name 张敏), not Latin sub-token junk -> keep a >=2
    # floor for non-Latin, >=3 for Latin (otherwise short CJK/Hangul names are wrongly dropped).
    import re as _re
    _NONLATIN = _re.compile(r"[Ѐ-ӿ֐-ۿऀ-ॿ฀-๿぀-ヿ㐀-鿿가-힯豈-﫿]")
    seen, uniq = set(), []
    for sp in sorted(out, key=lambda s: (s["start"], s["end"])):
        k = (sp["start"], sp["end"], sp["type"])
        if k in seen: continue
        frag = text[sp["start"]:sp["end"]].strip()
        if len(frag) < (2 if _NONLATIN.search(frag) else 3): continue  # drop fragments
        seen.add(k); uniq.append(sp)
    # CONTEXT tier: cue-gated alphanumeric IDs (Passport No:/Policy #/MRN: ...) — recovers
    # the structure-less IDs the model can't learn. Authoritative for their types; AUTH wins.
    try:
        from context_cued import context_cued_spans, CONTEXT_TYPES
    except Exception:
        return uniq
    cued = context_cued_spans(text)
    if not cued:
        return uniq
    auth_claim = [False] * len(text)
    for sp in uniq:
        if sp["type"] in _AUTH_TYPES:
            for i in range(sp["start"], sp["end"]): auth_claim[i] = True
    kept = [c for c in cued if not any(auth_claim[c["start"]:c["end"]])]
    rng = [(c["start"], c["end"]) for c in kept]
    def _ov(sp): return any(min(e, sp["end"]) > max(s, sp["start"]) for s, e in rng)
    merged = [sp for sp in uniq if not (sp["type"] in CONTEXT_TYPES and _ov(sp))] + kept
    seen, final = set(), []
    for sp in sorted(merged, key=lambda s: (s["start"], s["end"])):
        k = (sp["start"], sp["end"], sp["type"])
        if k in seen: continue
        seen.add(k); final.append(sp)
    # GROUP-B cue tier: amount/date/phone/postal where the model stayed silent but a
    # field-cue+shape is present (authoritative for those 4 types; AUTH wins; address excluded).
    try:
        from context_cued import group_b_cue_spans
    except Exception:
        return final
    gb = [c for c in group_b_cue_spans(text)
          if c["type"] in ("financial.amount","identity.date_of_birth","contact.phone","contact.postal_code")]
    if not gb:
        return final
    aclaim = [False] * len(text)
    for sp in final:
        if sp["type"] in _AUTH_TYPES:
            for i in range(sp["start"], sp["end"]): aclaim[i] = True
    keptb = []
    for c in gb:
        if any(aclaim[c["start"]:c["end"]]): continue
        s, e = c["start"], c["end"]
        if c["type"] == "financial.amount":
            while e > s and text[e-1] in ".,;": e -= 1
        keptb.append({"start": s, "end": e, "type": c["type"], "text": text[s:e]})
    rb = [(c["start"], c["end"]) for c in keptb]
    def _ovb(sp): return any(min(e, sp["end"]) > max(s, sp["start"]) for s, e in rb)
    GB = {"financial.amount","identity.date_of_birth","contact.phone","contact.postal_code"}
    merged2 = [sp for sp in final if not (sp["type"] in GB and _ovb(sp))] + keptb
    seen, out2 = set(), []
    for sp in sorted(merged2, key=lambda s: (s["start"], s["end"])):
        k = (sp["start"], sp["end"], sp["type"])
        if k in seen: continue
        seen.add(k); out2.append(sp)
    return out2

def model_spans(text: str, tok, model):
    import torch
    enc = tok(text, return_offsets_mapping=True, return_tensors="pt", truncation=True, max_length=2048)
    off = enc.pop("offset_mapping")[0].tolist()
    enc = {k: v.to(model.device) for k, v in enc.items()}
    with torch.no_grad():
        ids = model(**enc).logits[0].argmax(-1).tolist()
    id2label = model.config.id2label
    spans, cur = [], None
    for (a, b), i in zip(off, ids):
        lab = id2label[i]
        if b <= a or lab == "O":
            if cur: spans.append(cur); cur = None
            continue
        typ = lab.split("-", 1)[1] if "-" in lab else lab
        if lab[:2] in ("B-", "S-") or cur is None or cur["type"] != typ:
            if cur: spans.append(cur)
            cur = {"start": a, "end": b, "type": typ}
        else:
            cur["end"] = b
    if cur: spans.append(cur)
    # trim leading/trailing whitespace
    for sp in spans:
        while sp["start"] < sp["end"] and text[sp["start"]].isspace(): sp["start"] += 1
        while sp["end"] > sp["start"] and text[sp["end"] - 1].isspace(): sp["end"] -= 1
    return [s for s in spans if s["end"] > s["start"]]

def predict(text: str, tok, model, hybrid: bool = True) -> list[dict]:
    ms = model_spans(text, tok, model)
    if not hybrid:
        return [{"start": s["start"], "end": s["end"], "type": s["type"],
                 "text": text[s["start"]:s["end"]]} for s in ms]
    return hybrid_spans(text, ms)