build-small-hackathon/trace-field-notes

1"""Optional model-based PII redaction using ``openai/privacy-filter``.2 3The deterministic pipeline always runs regex redaction (:mod:`redaction`). On the4Hugging Face Space GPU this module adds a second pass: a token-classification5model (``openai/privacy-filter``) flags personal or sensitive spans that regex6patterns miss — names, phone numbers, postal addresses, and the like — and masks7them with typed placeholders.8 9Heavy imports (``torch``/``transformers``) load lazily so the deterministic10analyzer, the test suite, and local development keep working without GPU11dependencies. If the model cannot be loaded, the caller falls back to regex-only12redaction and records the reason in the privacy notes.13"""14 15from __future__ import annotations16 17import functools18import time19from collections import Counter20from typing import Any, Callable21 22from model_runtime import resolve_device23from profiling import get_logger24from redaction import RedactionResult25 26logger = get_logger()27 28 29PRIVACY_MODEL_ID = "openai/privacy-filter"30 31# Only mask spans the model is reasonably confident about.32PRIVACY_MIN_SCORE = 0.533 34# Model entity group -> (placeholder written into the text, human label for notes).35PII_TYPES: dict[str, tuple[str, str]] = {36    "private_person": ("[REDACTED_NAME]", "personal name"),37    "private_email": ("[REDACTED_EMAIL]", "email address"),38    "private_phone": ("[REDACTED_PHONE]", "phone number"),39    "private_address": ("[REDACTED_ADDRESS]", "postal address"),40    "private_url": ("[REDACTED_URL]", "personal URL"),41    "private_date": ("[REDACTED_DATE]", "personal date"),42    "account_number": ("[REDACTED_ACCOUNT]", "account number"),43    "secret": ("[REDACTED_SECRET]", "secret"),44}45 46# (texts) -> per-text list of {"start", "end", "label"} spans.47DetectFn = Callable[[list[str]], list[list[dict[str, Any]]]]48 49_PIPELINE_CACHE: dict[str, Any] = {}50 51 52def redact_texts(53    texts: list[str],54    *,55    detect: DetectFn | None = None,56    device: str | None = None,57) -> list[RedactionResult]:58    """Detect and mask PII in each text, returning one result per input.59 60    ``detect`` defaults to :func:`_local_detect` (the lazy model); tests inject a61    stand-in so the masking logic runs without ``torch``. ``device`` forces the62    compute device for the default detector (``cuda`` / ``mps`` / ``cpu``).63    """64 65    detector = detect or functools.partial(_local_detect, device=device)66    spans_per_text = detector(texts)67    return [_apply_spans(text, spans) for text, spans in zip(texts, spans_per_text)]68 69 70def _merge_spans(text: str, spans: list[dict[str, Any]]) -> list[dict[str, Any]]:71    """Drop malformed spans and merge same-label runs into clean, disjoint spans.72 73    ``openai/privacy-filter`` uses BIOES tags, which the pipeline's IOB-oriented74    "simple" aggregation can split into adjacent fragments of one entity (and a75    leading separator can leave a one-character gap). Merging same-label spans76    that overlap or sit within one character keeps each entity to a single77    placeholder; a remaining different-label overlap is clipped to stay disjoint.78    """79 80    valid = [81        span82        for span in spans83        if span.get("label") in PII_TYPES84        and 0 <= int(span["start"]) < int(span["end"]) <= len(text)85    ]86    valid.sort(key=lambda span: (int(span["start"]), int(span["end"])))87 88    merged: list[dict[str, Any]] = []89    for span in valid:90        start, end, label = int(span["start"]), int(span["end"]), span["label"]91        if merged:92            prev = merged[-1]93            if label == prev["label"] and start <= prev["end"] + 1:94                prev["end"] = max(prev["end"], end)95                continue96            if start < prev["end"]:  # different-label overlap: keep them disjoint97                start = prev["end"]98                if start >= end:99                    continue100        merged.append({"start": start, "end": end, "label": label})101    return merged102 103 104def _apply_spans(text: str, spans: list[dict[str, Any]]) -> RedactionResult:105    """Replace detected spans with typed placeholders, right-to-left."""106 107    counts: Counter[str] = Counter()108    redacted = text109    for span in sorted(_merge_spans(text, spans), key=lambda span: span["start"], reverse=True):110        placeholder, label = PII_TYPES[span["label"]]111        redacted = redacted[: span["start"]] + placeholder + redacted[span["end"] :]112        counts[label] += 1113 114    notes = [f"{label}: {count}" for label, count in sorted(counts.items())]115    return RedactionResult(text=redacted, notes=notes, count=sum(counts.values()))116 117 118def _local_detect(texts: list[str], device: str | None = None) -> list[list[dict[str, Any]]]:119    """Run ``openai/privacy-filter`` and return confident PII spans per text.120 121    Imported lazily: ``transformers``/``torch`` only need to exist where the122    model actually runs, never for the deterministic path, tests, or light local123    development.124    """125 126    pipe = _load_pipeline(device=device)127    started = time.perf_counter()128    results: list[list[dict[str, Any]]] = []129    for text in texts:130        if not text.strip():131            results.append([])132            continue133        entities = pipe(text)134        spans = [135            {136                "start": int(entity["start"]),137                "end": int(entity["end"]),138                "label": entity["entity_group"],139            }140            for entity in entities141            if entity.get("entity_group") in PII_TYPES142            and entity.get("start") is not None143            and entity.get("end") is not None144            and float(entity.get("score", 1.0)) >= PRIVACY_MIN_SCORE145        ]146        results.append(spans)147    detected = sum(len(spans) for spans in results)148    logger.debug(149        "privacy-filter scanned %d messages, %d raw spans in %.2fs",150        len(texts),151        detected,152        time.perf_counter() - started,153    )154    return results155 156 157def _load_pipeline(device: str | None = None) -> Any:158    """Lazily build and cache the token-classification pipeline per device."""159 160    resolved = resolve_device(device)161    cached = _PIPELINE_CACHE.get(resolved)162    if cached is not None:163        return cached164 165    from transformers import pipeline166 167    # transformers pipeline device: 0 for cuda, "mps"/"cpu" otherwise.168    pipe_device = 0 if resolved == "cuda" else resolved169    started = time.perf_counter()170    pipe = pipeline(171        "token-classification",172        model=PRIVACY_MODEL_ID,173        aggregation_strategy="simple",174        device=pipe_device,175    )176    logger.info(177        "loaded %s on %s in %.1fs", PRIVACY_MODEL_ID, resolved, time.perf_counter() - started178    )179    _PIPELINE_CACHE[resolved] = pipe180    return pipe181