Upload 6 files

README.md

@@ -1,14 +1,109 @@
 ---
 title: PrivacyShield
-emoji: 🌍
-colorFrom: pink
+emoji: 🛡️
+colorFrom: red
 colorTo: gray
 sdk: gradio
-sdk_version: 6.18.0
-python_version: '3.13'
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
+license: mit
+tags:
+  - build-small-hackathon
+  - privacy
+  - pii
+  - security
+  - llm-guardrails
+  - ner
 short_description: Local PII & secret firewall for LLMs
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🛡️ PrivacyShield — a local firewall for LLMs
+
+**Strip PII and leaked secrets out of text *before* it ever reaches an LLM API — then put the
+real values back into the response. Nothing sensitive leaves your machine.**
+
+Every week another company leaks customer data or an API key into an LLM prompt. The answer
+isn't "stop using AI" — it's a guardrail that runs **locally**, in front of the model.
+
+> **Demo it in two clicks:** open the app → **Try the leaked-secret example → Sanitize.** Watch the
+> AWS key, JWT, emails, Aadhaar (checksum-validated), names and address get masked — **"N blocked ·
+> 0 leaked."** Then hit **Simulate the LLM round-trip**: the model only ever sees placeholders, and
+> the real values are restored on your machine.
+
+## Why this matters
+
+- **The privacy requirement makes a small local model the *correct* design, not a compromise.** You
+  literally cannot send PII to a cloud API to have it redacted. PrivacyShield runs entirely on-device.
+- **It catches what regex can't.** Structured data (Aadhaar, PAN, cards, AWS keys, JWTs) is caught by
+  high-precision, **checksum-validated** detectors. Context-dependent data (**names, addresses,
+  orgs**) is caught by a **fine-tuned model** — regex is blind to these.
+- **The round-trip keeps the LLM useful.** Mask → call the LLM with safe text → restore the originals
+  into the answer locally. You get a real answer; the data never left.
+- **Compliance-ready.** Aligns with privacy regimes (India's DPDP Act, GDPR) that require minimizing
+  exposure of personal data to third parties.
+
+## How it works
+
+```
+your text
+  → DETECT : checksum-validated regex (structured PII + secrets)  ∪  fine-tuned NER (names/addresses)
+  → MASK   : each finding → reversible placeholder, e.g. [PERSON_NAME_1], [SECRET_1]; originals kept
+             only in an in-memory vault (never logged, never sent)
+  → CALL   : send the sanitized text to the LLM (the LLM only ever sees placeholders)
+  → RESTORE: swap placeholders back to the real values in the response, locally
+```
+
+## What it detects
+
+| Layer | Examples | How |
+|---|---|---|
+| **Structured PII** | email, phone, **Aadhaar** (Verhoeff checksum), **PAN**, IFSC, **card** (Luhn), UPI, IP | deterministic, high precision |
+| **Secrets** | AWS keys (`AKIA…`), JWTs, GitHub tokens, private-key blocks, high-entropy strings | regex + Shannon entropy |
+| **Contextual PII** | **person names, addresses, organizations** | **fine-tuned XLM-RoBERTa** |
+
+## The model — real data, real evaluation (not vibes)
+
+- **Base:** `FacebookAI/xlm-roberta-base` (~270M params — runs on CPU, no GPU needed at inference).
+- **Fine-tuned on:** [`ai4privacy/pii-masking-200k`](https://huggingface.co/datasets/ai4privacy/pii-masking-200k)
+  (real, span-labeled) **+ synthetic Indian PII** (valid-format Aadhaar/PAN/IFSC/UPI, Indian names &
+  addresses) so it handles Indian documents, which generic tools miss.
+- **Model:** [`perceptron01/privacyshield-ner`](https://huggingface.co/perceptron01/privacyshield-ner)
+- **Param math:** 0.27B ≪ 32B cap ✅
+
+**Evaluation (held-out mix of ai4privacy + synthetic Indian PII):**
+
+| Method | PERSON recall | ADDRESS recall | structured PII / secrets |
+|---|---|---|---|
+| regex-only baseline | ~0.00 | ~0.00 | high (checksum-validated) |
+| **PrivacyShield (regex + fine-tuned model)** | **~0.97** | **~0.97** | high |
+
+**Recall is the metric we optimize** — a missed secret or PII item is a leak, so a false negative is
+far worse than over-masking. Overall fine-tuned F1 ≈ **0.97** (precision 0.97 / recall 0.97).
+
+*Honest limitations:* the synthetic portion of the test set is formulaic and inflates absolute scores;
+the model occasionally labels an organization as ADDRESS (the value is still masked, so nothing leaks);
+free-text address boundaries are imperfect. The structured/secret layer is the high-precision backbone.
+
+## Privacy by design
+
+No database, no auth, no persistence. The detected values live only in an in-memory vault for the
+duration of a request; the downloadable **audit log contains placeholders only — never raw values.**
+
+## Run locally
+
+```bash
+pip install -r requirements.txt
+python app.py
+```
+
+## Tech
+
+Gradio · Hugging Face Transformers · a fine-tuned XLM-RoBERTa token classifier · deterministic
+detectors with Verhoeff (Aadhaar) and Luhn (card) checksum validation + Shannon-entropy secret detection.
+
+## Submission
+
+- 🤗 **Live Space:** ‹SPACE_URL›
+- 🎥 **Demo video:** ‹VIDEO_URL›
+- ���� **Social post:** ‹LINKEDIN_URL›

app.py

@@ -0,0 +1,150 @@
+"""app.py — PrivacyShield Gradio firewall UI + orchestration (§5).
+
+Layout: paste text -> Sanitize -> sanitized output + findings table +
+"N blocked · 0 leaked" banner -> optional simulated LLM round-trip
+(mask -> stub LLM -> restore) -> downloadable audit log (placeholders only).
+
+The vault (placeholder -> original) lives only in gr.State (per-session,
+in-memory) and is never written to the audit log.
+"""
+import json
+import tempfile
+
+import gradio as gr
+
+from firewall import audit_log, round_trip, sanitize
+
+PII_EXAMPLE = (
+    "Hi team, this is Ramesh Kumar from Acme Traders Pvt Ltd. My number is "
+    "9876543210 and email ramesh.kumar1@gmail.com. Please update my KYC — "
+    "Aadhaar 2341 7634 8900, PAN ABCDE1234F. Refund my payment to card "
+    "4532 0111 1222 2233 and ship the replacement to 12, MG Road, Pune, "
+    "Maharashtra - 411001."
+)
+
+SECRET_EXAMPLE = (
+    "# config.py — DO NOT COMMIT\n"
+    'AWS_ACCESS_KEY_ID = "AKIAIOSFODNN7EXAMPLE"\n'
+    'AWS_SECRET_ACCESS_KEY = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"\n'
+    'JWT_TOKEN = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9'
+    '.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ'
+    '.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"\n'
+    'SLACK_WEBHOOK = "xoxb-1234567890-abcdefghijklmnop"\n'
+    'ADMIN_EMAIL = "admin@acmetraders.com"\n'
+    "-----BEGIN RSA PRIVATE KEY-----\n"
+    "MIIBOgIBAAJBAK0z9k1234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJ\n"
+    "-----END RSA PRIVATE KEY-----"
+)
+
+FINDINGS_HEADERS = ["type", "masked value", "source", "confidence"]
+
+
+def _findings_rows(findings: list[dict]) -> list[list]:
+    return [[f["type"], f["masked_value"], f["source"], f["confidence"]] for f in findings]
+
+
+def _banner(result: dict) -> str:
+    note = ""
+    if result["model_status"] == "off":
+        note = " _(rules-only mode — fine-tuned model disabled)_"
+    elif result["model_status"] == "unavailable":
+        note = " _(rules-only mode — fine-tuned model unavailable, detectors still active)_"
+    return f"### 🛡️ {result['blocked']} items blocked · 0 leaked{note}"
+
+
+def do_sanitize(text: str, pii: bool, secrets: bool, use_model: bool):
+    if not text or not text.strip():
+        return "", [], "### 🛡️ 0 items blocked · 0 leaked", None
+    result = sanitize(text, pii=pii, secrets=secrets, use_model=use_model)
+    return result["sanitized_text"], _findings_rows(result["findings"]), _banner(result), result
+
+
+def do_round_trip(text: str, pii: bool, secrets: bool, use_model: bool):
+    if not text or not text.strip():
+        return "", "", "", [], "### 🛡️ 0 items blocked · 0 leaked", None
+    result = round_trip(text, pii=pii, secrets=secrets, use_model=use_model)
+    return (
+        result["sanitized_text"],
+        result["llm_response"],
+        result["restored_response"],
+        _findings_rows(result["findings"]),
+        _banner(result),
+        result,
+    )
+
+
+def do_download(result: dict | None):
+    if not result:
+        return None
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".json", prefix="privacyshield_audit_", delete=False, encoding="utf-8"
+    ) as f:
+        json.dump(audit_log(result), f, indent=2)
+        path = f.name
+    return path
+
+
+with gr.Blocks(title="PrivacyShield") as demo:
+    gr.Markdown("# 🛡️ PrivacyShield — your local firewall for LLMs")
+    gr.Markdown(
+        "Strip PII & secrets before they ever reach an API. "
+        "Nothing sensitive leaves your machine."
+    )
+
+    input_text = gr.Textbox(
+        label="Paste a prompt / code snippet", lines=8, placeholder="Paste text here…"
+    )
+
+    with gr.Row():
+        pii_example_btn = gr.Button("Try: PII example")
+        secret_example_btn = gr.Button("Try: leaked-secret example")
+
+    with gr.Row():
+        pii_cb = gr.Checkbox(value=True, label="PII")
+        secrets_cb = gr.Checkbox(value=True, label="Secrets")
+        model_cb = gr.Checkbox(value=True, label="Use fine-tuned model")
+        sanitize_btn = gr.Button("🛡️ Sanitize", variant="primary")
+
+    banner = gr.Markdown("### 🛡️ 0 items blocked · 0 leaked")
+    sanitized_output = gr.Textbox(
+        label="Sanitized output (placeholders shown)", lines=8, interactive=False
+    )
+    findings = gr.Dataframe(headers=FINDINGS_HEADERS, label="Findings", interactive=False)
+
+    with gr.Accordion("⚙️ Simulate LLM round-trip", open=False):
+        gr.Markdown(
+            "Sends the **sanitized** text to a stub LLM (templated echo, no API key "
+            "needed) and restores the original values into the response."
+        )
+        round_trip_btn = gr.Button("Simulate LLM round-trip")
+        rt_sanitized = gr.Textbox(label="1. Sanitized prompt -> LLM", lines=4, interactive=False)
+        rt_llm = gr.Textbox(label="2. LLM (stub) response", lines=4, interactive=False)
+        rt_restored = gr.Textbox(
+            label="3. Response with originals RESTORED", lines=4, interactive=False
+        )
+
+    download_btn = gr.Button("⬇ Download audit log (JSON)")
+    audit_file = gr.File(label="audit_log.json (placeholders only — never raw values)")
+
+    state = gr.State()
+
+    pii_example_btn.click(lambda: PII_EXAMPLE, outputs=input_text)
+    secret_example_btn.click(lambda: SECRET_EXAMPLE, outputs=input_text)
+
+    sanitize_btn.click(
+        do_sanitize,
+        inputs=[input_text, pii_cb, secrets_cb, model_cb],
+        outputs=[sanitized_output, findings, banner, state],
+    )
+
+    round_trip_btn.click(
+        do_round_trip,
+        inputs=[input_text, pii_cb, secrets_cb, model_cb],
+        outputs=[rt_sanitized, rt_llm, rt_restored, findings, banner, state],
+    )
+
+    download_btn.click(do_download, inputs=[state], outputs=[audit_file])
+
+
+if __name__ == "__main__":
+    demo.launch()

detectors.py

@@ -0,0 +1,226 @@
+"""Layer 1 — deterministic detectors: regex + validators + entropy.
+
+Each detector returns spans as dicts:
+    {"start": int, "end": int, "type": str, "value": str,
+     "source": "regex", "confidence": 0.99}
+
+High precision is the goal: numeric PII (Aadhaar, card) must pass real
+checksums (Verhoeff / Luhn) so random 12/16-digit numbers are NOT flagged.
+"""
+import math
+import re
+import string
+
+# ---------------------------------------------------------------------------
+# Verhoeff checksum (Aadhaar) — identical tables to Appendix A
+# ---------------------------------------------------------------------------
+_D = [
+    [0,1,2,3,4,5,6,7,8,9],[1,2,3,4,0,6,7,8,9,5],[2,3,4,0,1,7,8,9,5,6],
+    [3,4,0,1,2,8,9,5,6,7],[4,0,1,2,3,9,5,6,7,8],[5,9,8,7,6,0,4,3,2,1],
+    [6,5,9,8,7,1,0,4,3,2],[7,6,5,9,8,2,1,0,4,3],[8,7,6,5,9,3,2,1,0,4],
+    [9,8,7,6,5,4,3,2,1,0]]
+_P = [
+    [0,1,2,3,4,5,6,7,8,9],[1,5,7,6,2,8,3,0,9,4],[5,8,0,3,7,9,6,1,4,2],
+    [8,9,1,6,0,4,3,5,2,7],[9,4,5,3,1,2,6,8,7,0],[4,2,8,6,5,7,3,9,0,1],
+    [2,7,9,3,8,0,6,4,1,5],[7,0,4,6,9,1,3,2,5,8]]
+_INV = [0,4,3,2,1,5,6,7,8,9]
+
+
+def verhoeff_validate(number: str) -> bool:
+    """Validate a Verhoeff checksum (last digit is the check digit)."""
+    c = 0
+    for i, item in enumerate(reversed(number)):
+        c = _D[c][_P[i % 8][int(item)]]
+    return c == 0
+
+
+# ---------------------------------------------------------------------------
+# Luhn checksum (cards)
+# ---------------------------------------------------------------------------
+def luhn_validate(card: str) -> bool:
+    s = 0
+    for i, ch in enumerate(reversed(card)):
+        d = int(ch)
+        if i % 2 == 1:
+            d *= 2
+            if d > 9:
+                d -= 9
+        s += d
+    return s % 10 == 0
+
+
+# ---------------------------------------------------------------------------
+# Shannon entropy (secrets)
+# ---------------------------------------------------------------------------
+def shannon_entropy(s: str) -> float:
+    if not s:
+        return 0.0
+    probs = [s.count(c) / len(s) for c in set(s)]
+    return -sum(p * math.log2(p) for p in probs)
+
+
+# A small dictionary-word check so plain English words of length >= 20
+# (rare, but possible in compound text) don't get flagged as secrets.
+_COMMON_WORDS = {
+    "responsibility", "internationalization", "characterization",
+    "telecommunications", "incomprehensibility", "disproportionately",
+}
+
+
+def _is_dictionary_word(token: str) -> bool:
+    return token.lower() in _COMMON_WORDS
+
+
+# ---------------------------------------------------------------------------
+# Regex patterns
+# ---------------------------------------------------------------------------
+_EMAIL_RE = re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b")
+_PHONE_RE = re.compile(r"(?<!\d)(?:\+?91[-\s]?)?[6-9]\d{9}(?!\d)")
+_AADHAAR_RE = re.compile(r"(?<!\d)\d{4}\s?\d{4}\s?\d{4}(?!\d)")
+_PAN_RE = re.compile(r"\b[A-Z]{5}[0-9]{4}[A-Z]\b")
+_IFSC_RE = re.compile(r"\b[A-Z]{4}0[A-Z0-9]{6}\b")
+_CARD_RE = re.compile(r"(?<!\d)\d(?:[ -]?\d){12,18}(?!\d)")
+_IP_RE = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b")
+_DOB_RE = re.compile(r"\b(?:0[1-9]|[12]\d|3[01])[/-](?:0[1-9]|1[0-2])[/-](?:19|20)\d{2}\b")
+_UPI_RE = re.compile(r"\b[\w.\-]{2,256}@(?:[a-zA-Z]{3,64})\b")
+
+# Secrets
+_AWS_KEY_RE = re.compile(r"\b(?:AKIA|ASIA)[A-Z0-9]{16}\b")
+_GITHUB_TOKEN_RE = re.compile(r"\bgh[pousr]_[A-Za-z0-9]{36,255}\b")
+_SLACK_TOKEN_RE = re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,72}\b")
+_JWT_RE = re.compile(r"\beyJ[A-Za-z0-9_-]+\.eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b")
+_PRIVATE_KEY_RE = re.compile(
+    r"-----BEGIN [A-Z0-9 ]*PRIVATE KEY-----.*?-----END [A-Z0-9 ]*PRIVATE KEY-----",
+    re.DOTALL,
+)
+# Generic API-token-like assignments: api_key="...", token: '...', secret=...
+_GENERIC_TOKEN_RE = re.compile(
+    r"(?i)\b(?:api[_-]?key|api[_-]?secret|access[_-]?token|secret[_-]?key|"
+    r"auth[_-]?token|password|token)\s*[:=]\s*[\"']?([A-Za-z0-9/_+\-]{16,})[\"']?"
+)
+
+# Indian UPI handles (specific bank suffixes, narrower than generic email-ish)
+_UPI_SUFFIXES = (
+    "okhdfcbank", "oksbi", "okicici", "okaxis", "ybl", "paytm", "ibl", "axl", "apl",
+)
+
+
+def _find_aadhaar(text):
+    spans = []
+    for m in _AADHAAR_RE.finditer(text):
+        digits = m.group(0).replace(" ", "")
+        if len(digits) == 12 and verhoeff_validate(digits):
+            spans.append({
+                "start": m.start(), "end": m.end(), "type": "AADHAAR",
+                "value": m.group(0), "source": "regex", "confidence": 0.99,
+            })
+    return spans
+
+
+def _find_card(text):
+    spans = []
+    for m in _CARD_RE.finditer(text):
+        digits = re.sub(r"[ -]", "", m.group(0))
+        if len(digits) in (13, 14, 15, 16, 17, 18, 19) and luhn_validate(digits):
+            spans.append({
+                "start": m.start(), "end": m.end(), "type": "CARD",
+                "value": m.group(0), "source": "regex", "confidence": 0.99,
+            })
+    return spans
+
+
+def _find_upi(text):
+    spans = []
+    for m in _UPI_RE.finditer(text):
+        handle = m.group(0)
+        suffix = handle.rsplit("@", 1)[-1].lower()
+        if suffix in _UPI_SUFFIXES:
+            spans.append({
+                "start": m.start(), "end": m.end(), "type": "UPI",
+                "value": handle, "source": "regex", "confidence": 0.97,
+            })
+    return spans
+
+
+def _find_entropy_secrets(text, threshold=4.0, min_len=20):
+    spans = []
+    for m in re.finditer(r"[A-Za-z0-9+/=_\-]{%d,}" % min_len, text):
+        token = m.group(0)
+        ent = shannon_entropy(token)
+        if ent > threshold and not _is_dictionary_word(token) and not token.isdigit():
+            spans.append({
+                "start": m.start(), "end": m.end(), "type": "SECRET",
+                "value": token, "source": "entropy",
+                "confidence": round(min(0.5 + (ent - threshold) * 0.15, 0.95), 2),
+            })
+    return spans
+
+
+def _find_generic_tokens(text):
+    spans = []
+    for m in _GENERIC_TOKEN_RE.finditer(text):
+        value = m.group(1)
+        if len(value) < 8:
+            continue
+        start, end = m.span(1)
+        spans.append({
+            "start": start, "end": end, "type": "SECRET",
+            "value": value, "source": "regex", "confidence": 0.9,
+        })
+    return spans
+
+
+def _simple_finditer(pattern, type_name, text, confidence=0.99, source="regex"):
+    return [
+        {
+            "start": m.start(), "end": m.end(), "type": type_name,
+            "value": m.group(0), "source": source, "confidence": confidence,
+        }
+        for m in pattern.finditer(text)
+    ]
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+def detect_pii(text: str) -> list[dict]:
+    """Detect structured + context-free PII (Layer 1, PII categories)."""
+    spans = []
+    spans += _simple_finditer(_EMAIL_RE, "EMAIL", text)
+    spans += _find_aadhaar(text)
+    spans += _simple_finditer(_PAN_RE, "PAN", text)
+    spans += _simple_finditer(_IFSC_RE, "IFSC", text)
+    spans += _find_card(text)
+    spans += _simple_finditer(_IP_RE, "IP", text, confidence=0.95)
+    spans += _simple_finditer(_DOB_RE, "DOB", text, confidence=0.9)
+    spans += _find_upi(text)
+    # phone last: avoid double-matching digits already consumed by aadhaar/card
+    spans += _simple_finditer(_PHONE_RE, "PHONE", text, confidence=0.95)
+    return spans
+
+
+def detect_secrets(text: str) -> list[dict]:
+    """Detect credentials / secrets (Layer 1, secret categories)."""
+    spans = []
+    spans += _simple_finditer(_AWS_KEY_RE, "AWS_KEY", text)
+    spans += _simple_finditer(_GITHUB_TOKEN_RE, "GITHUB_TOKEN", text)
+    spans += _simple_finditer(_SLACK_TOKEN_RE, "SLACK_TOKEN", text)
+    spans += _simple_finditer(_JWT_RE, "JWT", text)
+    spans += _simple_finditer(_PRIVATE_KEY_RE, "PRIVATE_KEY", text, confidence=0.99)
+    spans += _find_generic_tokens(text)
+    spans += _find_entropy_secrets(text)
+    return spans
+
+
+def detect_all(text: str, pii: bool = True, secrets: bool = True) -> list[dict]:
+    """Run all enabled detector groups and return raw (possibly overlapping) spans."""
+    spans = []
+    if pii:
+        spans += detect_pii(text)
+    if secrets:
+        spans += detect_secrets(text)
+    return spans
+
+
+# Alias matching the integration doc's naming (`detectors.detect(text)`).
+detect = detect_all

firewall.py

@@ -0,0 +1,149 @@
+"""firewall.py — detect -> mask -> (call) -> restore.
+
+Orchestrates Layer 1 (detectors.py, always on) and Layer 2 (model_ner.py,
+optional/graceful) per ARCHITECTURE in the build spec. Implements the merge
+rule (union spans, resolve overlaps by confidence/length, never overlapping
+masks), a reversible in-memory vault, and a stub LLM round-trip.
+"""
+from __future__ import annotations
+
+from detectors import detect_all
+
+# Types produced by the secret-side detectors / NER -> placeholder category
+_SECRET_TYPES = {
+    "AWS_KEY", "GITHUB_TOKEN", "SLACK_TOKEN", "JWT", "PRIVATE_KEY", "SECRET",
+}
+
+
+def _category(type_name: str) -> str:
+    return "SECRET" if type_name in _SECRET_TYPES else "PII"
+
+
+def merge_spans(spans: list[dict]) -> list[dict]:
+    """Union spans; on overlap keep the higher-confidence / longer span.
+
+    Returns a list of non-overlapping spans sorted by start position.
+    """
+    # Prefer higher confidence, then longer spans, then earlier start —
+    # greedily accept spans in that priority order, skipping anything that
+    # would overlap an already-accepted span.
+    ordered = sorted(
+        spans,
+        key=lambda s: (-s["confidence"], -(s["end"] - s["start"]), s["start"]),
+    )
+    accepted: list[dict] = []
+    for span in ordered:
+        if not any(span["start"] < a["end"] and a["start"] < span["end"] for a in accepted):
+            accepted.append(span)
+    return sorted(accepted, key=lambda s: s["start"])
+
+
+def mask(text: str, spans: list[dict]) -> tuple[str, dict, list[dict]]:
+    """Replace each span with a typed, reversible placeholder.
+
+    Returns (masked_text, vault, findings) where vault maps
+    placeholder -> original substring, and findings carry no raw values.
+    """
+    merged = merge_spans(spans)
+    vault: dict[str, str] = {}
+    findings: list[dict] = []
+    counters: dict[tuple[str, str], int] = {}
+
+    out: list[str] = []
+    last = 0
+    for span in merged:
+        cat = _category(span["type"])
+        key = (cat, span["type"])
+        counters[key] = counters.get(key, 0) + 1
+        placeholder = f"[{cat}_{span['type']}_{counters[key]}]"
+
+        vault[placeholder] = text[span["start"]:span["end"]]
+        out.append(text[last:span["start"]])
+        out.append(placeholder)
+        last = span["end"]
+
+        findings.append({
+            "type": span["type"],
+            "masked_value": placeholder,
+            "source": span["source"],
+            "confidence": span["confidence"],
+        })
+
+    out.append(text[last:])
+    return "".join(out), vault, findings
+
+
+def restore(text: str, vault: dict[str, str]) -> str:
+    """Swap placeholders back to their original values."""
+    for placeholder, original in vault.items():
+        text = text.replace(placeholder, original)
+    return text
+
+
+def sanitize(text: str, pii: bool = True, secrets: bool = True, use_model: bool = True) -> dict:
+    """Run the full DETECT -> MASK pipeline.
+
+    Returns a dict with sanitized_text, vault, findings, blocked count, and
+    model_status ("on" | "off" | "unavailable") describing whether Layer 2
+    (the fine-tuned NER model) contributed.
+    """
+    spans = detect_all(text, pii=pii, secrets=secrets)
+
+    model_status = "off"
+    if use_model:
+        try:
+            from model_ner import detect_entities
+            model_spans = detect_entities(text)
+            model_status = "on" if model_spans else "unavailable"
+            spans += model_spans
+        except Exception:
+            model_status = "unavailable"
+
+    sanitized_text, vault, findings = mask(text, spans)
+    return {
+        "sanitized_text": sanitized_text,
+        "sanitized": sanitized_text,
+        "vault": vault,
+        "findings": findings,
+        "blocked": len(findings),
+        "model_status": model_status,
+    }
+
+
+def call_llm_stub(sanitized_text: str) -> str:
+    """A no-API-key-needed stand-in for an LLM call: echoes the sanitized
+    prompt back inside a templated reply, so the round-trip demo never
+    depends on network access or secrets."""
+    return (
+        "Thanks for the message. Here's a summary of what I received "
+        f"(sanitized):\n\n{sanitized_text}\n\n"
+        "(This is a stub response — no data left your machine.)"
+    )
+
+
+def round_trip(text: str, pii: bool = True, secrets: bool = True, use_model: bool = True, llm_fn=None) -> dict:
+    """sanitize -> (call llm_fn or stub) -> restore originals in the reply."""
+    result = sanitize(text, pii=pii, secrets=secrets, use_model=use_model)
+    llm_fn = llm_fn or call_llm_stub
+    llm_response = llm_fn(result["sanitized_text"])
+    result["llm_response"] = llm_response
+    result["restored_response"] = restore(llm_response, result["vault"])
+    return result
+
+
+def audit_log(result: dict) -> dict:
+    """JSON-safe audit log: placeholders + types + sources only, never the
+    vault or any raw value."""
+    return {
+        "blocked": result["blocked"],
+        "leaked": 0,
+        "findings": [
+            {
+                "type": f["type"],
+                "masked_value": f["masked_value"],
+                "source": f["source"],
+                "confidence": f["confidence"],
+            }
+            for f in result["findings"]
+        ],
+    }

model_ner.py

@@ -0,0 +1,85 @@
+"""model_ner.py — Layer 2: fine-tuned token classifier (PERSON/ADDRESS/ORG).
+
+Loads the fine-tuned NER adapter from the HF Hub (pushed by
+train/train_lora.py) if torch is usable and the model is reachable. On any
+load/inference failure — no torch, no adapter pushed yet, no network — this
+module degrades to returning [] and the firewall keeps working on Layer 1
+alone (build spec §2/§6).
+
+Availability is checked with `transformers.utils.is_torch_available()`,
+never a bare `import torch`.
+"""
+from __future__ import annotations
+
+import os
+
+# The fine-tuned adapter pushed by train/train_lora.py, or override via the
+# PRIVACYSHIELD_NER_MODEL env var.
+MODEL_ID = os.environ.get("PRIVACYSHIELD_NER_MODEL", "perceptron01/privacyshield-ner")
+
+_NER_TYPES = {"PERSON", "ADDRESS", "ORG"}
+
+_pipeline = None
+_unavailable_reason: str | None = None
+
+
+def _load():
+    global _pipeline, _unavailable_reason
+    if _pipeline is not None or _unavailable_reason is not None:
+        return _pipeline
+    try:
+        from transformers.utils import is_torch_available
+        if not is_torch_available():
+            raise RuntimeError("torch is not installed/usable (need torch>=2.4)")
+
+        from transformers import pipeline
+        _pipeline = pipeline(
+            "token-classification",
+            model=MODEL_ID,
+            aggregation_strategy="simple",
+        )
+    except Exception as exc:  # noqa: BLE001 — any failure -> graceful degradation
+        _unavailable_reason = str(exc)
+        _pipeline = None
+    return _pipeline
+
+
+def is_available() -> bool:
+    """True if the fine-tuned model loaded successfully."""
+    return _load() is not None
+
+
+def unavailable_reason() -> str | None:
+    """Human-readable reason the model isn't available, or None if it is."""
+    _load()
+    return _unavailable_reason
+
+
+def detect_entities(text: str) -> list[dict]:
+    """Run Layer 2 NER. Returns [] if the model can't load or run."""
+    pipe = _load()
+    if pipe is None:
+        return []
+    try:
+        out = pipe(text)
+    except Exception:  # noqa: BLE001 — inference failure -> graceful degradation
+        return []
+
+    spans = []
+    for ent in out:
+        label = str(ent.get("entity_group", "")).upper()
+        if label not in _NER_TYPES:
+            continue
+        spans.append({
+            "start": ent["start"],
+            "end": ent["end"],
+            "type": label,
+            "value": ent["word"],
+            "source": "model",
+            "confidence": round(float(ent["score"]), 2),
+        })
+    return spans
+
+
+# Alias matching the integration doc's naming (`model_ner.detect_ner(text)`).
+detect_ner = detect_entities

requirements.txt

@@ -0,0 +1,14 @@
+gradio>=4.0
+regex
+pandas
+pytest
+
+# Layer 2 (optional, graceful) + eval/training — not required for the
+# rules-only firewall or its tests.
+transformers>=4.44
+torch>=2.2
+datasets
+peft
+accelerate
+huggingface_hub
+scikit-learn