Update app.py

app.py

@@ -1,113 +1,369 @@
+"""
+CLOAK PII Demo — ONNX INT8 CPU inference + Gradio UI + REST /detect API.
+Downloads model.int8.onnx from HF Hub at startup.
+Set HF_TOKEN secret in Space settings if the repo is private.
+The /detect endpoint is what the CLOAK API (Next.js) calls via callHfSpace:
+  POST /detect  Authorization: Bearer <API_SECRET>  {"text": "..."}
+  -> {entities, nested, redacted_text, model, processing_time_ms}
+"""
+
+import html as _html
+import json
 import os
+import time
+
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
-from key_store import (
-    generate_api_key, validate_key, increment_usage,
-    revoke_key, delete_key, list_keys
-)
-
-# ── Load your private HF model ────────────────────────────────────────────────
-HF_TOKEN = os.environ.get("HF_TOKEN")
-MODEL_ID  = "Wardline/CLOAK_V2.2.11"
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-model     = AutoModelForTokenClassification.from_pretrained(MODEL_ID, token=HF_TOKEN)
-ner_pipe  = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
-
-# ── Stitcher: merge split sub-word tokens ─────────────────────────────────────
-def stitch_entities(raw_entities: list) -> list:
-    merged = []
-    for r in raw_entities:
-        if (merged
-                and merged[-1]["entity_group"] == r["entity_group"]
-                and merged[-1]["end"] == r["start"]):
-            merged[-1]["word"] += r["word"].replace(" ", "")
-            merged[-1]["end"]   = r["end"]
-        else:
-            merged.append(dict(r))  # copy so we don't mutate the original
-    return merged
-
-# ── PII detection ─────────────────────────────────────────────────────────────
-def detect_pii(text: str, api_key: str) -> dict:
-    entry = validate_key(api_key)
-    if not entry:
-        return {"error": "Invalid or revoked API key.", "entities": []}
-    increment_usage(api_key)
-
-    raw      = ner_pipe(text)
-    entities = stitch_entities(raw)
+import numpy as np
+import onnxruntime as ort
+import pandas as pd
+from fastapi import FastAPI, HTTPException, Request
+from huggingface_hub import hf_hub_download, login
+from pydantic import BaseModel, Field
+from transformers import AutoTokenizer
+
+# ── Config ────────────────────────────────────────────────────────────────────
+
+MODEL_REPO = os.getenv("MODEL_REPO", "Wardline/CLOAK_V3.6")
+HF_TOKEN   = os.getenv("HF_TOKEN")
+# Bearer token the CLOAK API must send. callHfSpace sends `Bearer ${HF_TOKEN}`,
+# so default API_SECRET to HF_TOKEN — the same secret works for both unless set.
+API_SECRET = os.getenv("API_SECRET", HF_TOKEN)
+
+MAX_LEN   = 512
+STRIDE    = 128
+THRESHOLD = 0.0
+
+# ── Thread count: read cgroup CPU quota, not sched_getaffinity ────────────────
+# sched_getaffinity returns the host's full CPU set inside Docker containers.
+# The real allocated vCPU count lives in the cgroup CFS quota.
+
+def _container_cpu_count() -> int:
+    # cgroup v2
+    try:
+        with open("/sys/fs/cgroup/cpu.max") as f:
+            quota, period = f.read().split()
+        if quota != "max":
+            return max(1, round(int(quota) / int(period)))
+    except Exception:
+        pass
+    # cgroup v1
+    try:
+        with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as f:
+            quota = int(f.read())
+        with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as f:
+            period = int(f.read())
+        if quota > 0:
+            return max(1, quota // period)
+    except Exception:
+        pass
+    return os.cpu_count() or 2
+
+_n_threads = _container_cpu_count()
+print(f"Container CPU count: {_n_threads}  (os.cpu_count={os.cpu_count()})")
+
+# ── Load at startup ───────────────────────────────────────────────────────────
+
+if HF_TOKEN:
+    login(token=HF_TOKEN)
+
+_config_path = hf_hub_download(MODEL_REPO, "config.json", token=HF_TOKEN)
+_int8_path   = hf_hub_download(MODEL_REPO, "model.int8.onnx", token=HF_TOKEN)
+
+with open(_config_path, encoding="utf-8") as fh:
+    _config = json.load(fh)
+
+_id2label  = {int(k): v for k, v in _config["id2label"].items()}
+_tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, subfolder="tokenizer", token=HF_TOKEN)
+
+_opts = ort.SessionOptions()
+_opts.intra_op_num_threads = _n_threads
+_opts.inter_op_num_threads = 1
+_opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+_sess = ort.InferenceSession(_int8_path, sess_options=_opts, providers=["CPUExecutionProvider"])
+
+print(f"ONNX INT8 session ready — {_n_threads} threads, {len(_id2label)} classes")
+
+
+# ── Inference ─────────────────────────────────────────────────────────────────
+
+def _suppress_same_label(entities: list) -> list:
+    kept = []
+    for ent in sorted(entities, key=lambda x: -x["score"]):
+        if any(k["entity_group"] == ent["entity_group"]
+               and ent["start"] < k["end"] and k["start"] < ent["end"] for k in kept):
+            continue
+        kept.append(ent)
+    return kept
+
+
+def detect(text: str, threshold: float = THRESHOLD):
+    enc = _tokenizer(
+        text, max_length=MAX_LEN, truncation=True, stride=STRIDE, padding=True,
+        return_overflowing_tokens=True, return_offsets_mapping=True, return_tensors="np",
+    )
+    all_offsets = enc.pop("offset_mapping")
+    enc.pop("overflow_to_sample_mapping", None)
+
+    best: dict = {}
+    for w in range(all_offsets.shape[0]):
+        feed = {
+            "input_ids":      enc["input_ids"][w:w + 1].astype(np.int64),
+            "attention_mask": enc["attention_mask"][w:w + 1].astype(np.int64),
+        }
+        logits  = _sess.run(["logits"], feed)[0][0]   # [classes, L, L]
+        offsets = all_offsets[w]
+
+        cidx, sidx, eidx = np.where(logits > threshold)
+        if len(cidx) == 0:
+            continue
+
+        scores = 1.0 / (1.0 + np.exp(-logits[cidx, sidx, eidx]))
+        cs = offsets[sidx, 0]
+        ce = offsets[eidx, 1]
+        keep = cs < ce
+        for cid, c0, c1, sc in zip(cidx[keep], cs[keep], ce[keep], scores[keep]):
+            key = (int(cid), int(c0), int(c1))
+            if key not in best or sc > best[key]["score"]:
+                best[key] = {
+                    "entity_group": _id2label[int(cid)],
+                    "start": int(c0), "end": int(c1),
+                    "score": float(sc),
+                }
+
+    doc_len = max(1, len(text))
+    cands = [e for e in best.values() if (e["end"] - e["start"]) <= 0.5 * doc_len]
+    cands = _suppress_same_label(cands)
+
+    flat = []
+    for ent in sorted(cands, key=lambda x: (-(x["end"] - x["start"]), -x["score"])):
+        if any(ent["start"] < k["end"] and k["start"] < ent["end"] for k in flat):
+            continue
+        flat.append(ent)
+    flat = sorted(flat, key=lambda x: x["start"])
+    return flat, cands
+
+
+# ── UI helpers ────────────────────────────────────────────────────────────────
+
+COLOR_MAP = {
+    "PERSON_NAME": "#FF6B6B", "AADHAAR": "#6a0dad", "PASSPORT_NUMBER": "#1e90ff",
+    "PAN": "#ff8c00", "SSN": "#dc143c", "DRIVER_LICENSE": "#ff7f50", "VOTER_ID": "#ff69b4",
+    "PHONE_NUMBER": "#20b2aa", "EMAIL": "#2e8b57", "USERNAME": "#3cb371", "PASSWORD": "#8b0000",
+    "BANK_ACCOUNT": "#8b4513", "CREDIT_CARD": "#ff1493", "CVV": "#c71585",
+    "CREDIT_CARD_EXPIRATION": "#db7093", "UPI_ID": "#4682b4", "MONEY": "#2f4f4f",
+    "DATE_OF_BIRTH": "#228b22", "DATE": "#6b8e23", "ADDRESS": "#8b008b",
+    "LOCATION": "#008b8b", "ORGANIZATION": "#cd853f", "EMAIL_ADDRESS": "#2e8b57",
+    "IP_ADDRESS": "#483d8b", "URL": "#4169e1", "API_KEY": "#800000",
+}
+_FALLBACK = ["#45B7D1", "#96CEB4", "#FFEAA7", "#DDA0DD", "#F7DC6F", "#85C1E9"]
+_extra: dict = {}
+
+def _label_color(label: str) -> str:
+    if label in COLOR_MAP:
+        return COLOR_MAP[label]
+    if label not in _extra:
+        _extra[label] = _FALLBACK[len(_extra) % len(_FALLBACK)]
+    return _extra[label]
+
+
+def _highlight_html(text: str, cands: list) -> str:
+    if not cands:
+        return f"<div style='line-height:2.4;font-size:15px;font-family:sans-serif'>{_html.escape(text)}</div>"
+
+    ents = sorted(cands, key=lambda e: (e["start"], -(e["end"] - e["start"])))
+    for e in ents:
+        e["_children"] = []
+
+    roots = []
+    for e in ents:
+        parent = None
+        for c in ents:
+            if (c is not e
+                    and c["start"] <= e["start"] and e["end"] <= c["end"]
+                    and (c["end"] - c["start"]) > (e["end"] - e["start"])):
+                if parent is None or (c["end"] - c["start"]) < (parent["end"] - parent["start"]):
+                    parent = c
+        (parent["_children"] if parent else roots).append(e)
+
+    def render(lo, hi, children):
+        children = sorted(children, key=lambda c: c["start"])
+        out, cur = "", lo
+        for ch in children:
+            if ch["start"] < cur:
+                continue
+            out  += _html.escape(text[cur:ch["start"]])
+            color = _label_color(ch["entity_group"])
+            label = ch["entity_group"]
+            pct   = round(ch["score"] * 100, 1)
+            inner = render(ch["start"], ch["end"], ch["_children"])
+            badge = f'<small style="opacity:.8;font-size:10px"> {label}</small>'
+            title = f"{label} {pct}%"
+            out  += (f'<span style="background:{color};color:white;padding:2px 5px;'
+                     f'border-radius:4px;font-weight:600;" title="{title}">'
+                     f'{inner}{badge}</span>')
+            cur   = ch["end"]
+        out += _html.escape(text[cur:hi])
+        return out
 
+    body = render(0, len(text), roots)
+    return f"<div style='line-height:2.6;font-size:15px;font-family:sans-serif'>{body}</div>"
+
+
+def run_analyze(text: str, threshold: float):
+    if not text or not text.strip():
+        return "<div style='color:#888'>Paste text above and click Analyze.</div>", pd.DataFrame(), "", ""
+
+    t0 = time.time()
+    flat, cands = detect(text, threshold)
+    elapsed = round(time.time() - t0, 2)
+
+    html_out = _highlight_html(text, cands)
+
+    redacted = text
+    for e in sorted(flat, key=lambda x: x["start"], reverse=True):
+        redacted = redacted[:e["start"]] + f'[{e["entity_group"]}]' + redacted[e["end"]:]
+
+    rows = [{
+        "Type":       e["entity_group"],
+        "Text":       text[e["start"]:e["end"]].strip(),
+        "Confidence": f'{round(e["score"] * 100, 1)}%',
+        "Start":      e["start"],
+        "End":        e["end"],
+    } for e in flat]
+    df = pd.DataFrame(rows)
+
+    summary = f"{elapsed}s  |  {len(flat)} entities  |  {len(text)} chars  |  {_n_threads} CPU threads"
+    return html_out, df, redacted, summary
+
+
+# ── Gradio UI ─────────────────────────────────────────────────────────────────
+
+EXAMPLES = [
+    ["My name is John Smith, email john@example.com, SSN 123-45-6789, card 4532 8812 7731 9823.", 0.0],
+    ["Patient Priya Sharma, DOB 12/03/1991, Aadhaar 2345 6789 0123, prescribed metformin 500mg.", 0.0],
+    ["Invoice to Acme Corp, 400 Market St, San Francisco CA 94105. Contact: +1 415 555 0100.", 0.0],
+]
+
+with gr.Blocks(title="CLOAK PII Detector", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        "## CLOAK V3.6 — PII Detector\n"
+        "Detects personally identifiable information using nested NER (ONNX INT8 CPU). "
+        "Nested highlighting shows overlapping spans."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            txt_input = gr.Textbox(
+                label="Input text",
+                placeholder="Paste text here…",
+                lines=8, max_lines=30,
+            )
+            threshold_slider = gr.Slider(
+                minimum=-2.0, maximum=6.0, value=0.0, step=0.5,
+                label="Detection threshold  (higher = fewer, more confident detections)",
+            )
+            with gr.Row():
+                btn       = gr.Button("Analyze", variant="primary", scale=2)
+                clear_btn = gr.Button("Clear", scale=1)
+            summary_box = gr.Textbox(label="", lines=1, interactive=False, show_label=False)
+
+    gr.Examples(examples=EXAMPLES, inputs=[txt_input, threshold_slider], label="Try an example")
+
+    highlighted = gr.HTML(label="Highlighted output")
+
+    with gr.Row():
+        with gr.Column():
+            entity_table = gr.Dataframe(
+                headers=["Type", "Text", "Confidence", "Start", "End"],
+                label="Detected entities",
+                interactive=False, wrap=True,
+            )
+        with gr.Column():
+            redacted_out = gr.Textbox(label="Redacted text", lines=8, interactive=False)
+
+    btn.click(
+        fn=run_analyze,
+        inputs=[txt_input, threshold_slider],
+        outputs=[highlighted, entity_table, redacted_out, summary_box],
+    )
+    txt_input.submit(
+        fn=run_analyze,
+        inputs=[txt_input, threshold_slider],
+        outputs=[highlighted, entity_table, redacted_out, summary_box],
+    )
+    clear_btn.click(
+        fn=lambda: ("", pd.DataFrame(), "", ""),
+        outputs=[txt_input, entity_table, redacted_out, summary_box],
+    )
+
+# ── REST API for the CLOAK backend (callHfSpace contract) ─────────────────────
+
+api = FastAPI(title="CLOAK PII API", version="1.0.0")
+
+
+class DetectRequest(BaseModel):
+    text: str = Field(..., min_length=1, max_length=20_000)
+
+
+def _fmt(text: str, e: dict) -> dict:
+    """Shape one entity to the contract the CLOAK API expects."""
     return {
-        "text": text,
-        "entities": [
-            {
-                "entity": e["entity_group"],
-                "word":   e["word"],
-                "score":  round(e["score"], 4),
-                "start":  e["start"],
-                "end":    e["end"],
-            }
-            for e in entities
-        ],
+        "type":       e["entity_group"],
+        "text":       text[e["start"]:e["end"]].strip(),
+        "start":      e["start"],
+        "end":        e["end"],
+        "confidence": round(e["score"], 4),
     }
 
-# ── Key management helpers ─────────────────────────────────────────────────���──
-def create_key(label):
-    if not label.strip():
-        return "Please enter a label.", get_key_table()
-    key = generate_api_key(label.strip())
-    return f"✅ New key created (save it now):\n\n{key}", get_key_table()
 
-def get_key_table():
-    rows = list_keys()
-    if not rows:
-        return []
-    return [[r["label"], r["key_preview"], r["status"], r["usage"], r["created_at"]] for r in rows]
+def _verify(request: Request):
+    auth = request.headers.get("Authorization", "")
+    if not auth.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Missing Authorization header")
+    token = auth.removeprefix("Bearer ").strip()
+    if not API_SECRET or token != API_SECRET:
+        raise HTTPException(status_code=401, detail="Invalid token")
 
-def do_revoke(full_key):
-    revoke_key(full_key.strip())
-    return "Key revoked.", get_key_table()
 
-def do_delete(full_key):
-    delete_key(full_key.strip())
-    return "Key deleted.", get_key_table()
+@api.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_REPO, "engine": "onnx-int8-cpu",
+            "threads": _n_threads, "classes": len(_id2label)}
 
-# ── Gradio UI ─────────────────────────────────────────────────────────────────
-with gr.Blocks(title="PII Detection API") as demo:
-    gr.Markdown("# PII Detection API\nPowered by your private Hugging Face model.")
-
-    with gr.Tab("Detect PII"):
-        with gr.Row():
-            api_key_in = gr.Textbox(label="API key", placeholder="pii_sk_...")
-            text_in    = gr.Textbox(label="Input text", lines=4,
-                                    placeholder="Enter text to scan for PII...")
-        detect_btn  = gr.Button("Detect PII", variant="primary")
-        output_json = gr.JSON(label="Results")
-        detect_btn.click(fn=detect_pii, inputs=[text_in, api_key_in], outputs=output_json)
-
-    with gr.Tab("API Key Management"):
-        gr.Markdown("### Create a new key")
-        with gr.Row():
-            label_in   = gr.Textbox(label="Key label", placeholder="e.g. production")
-            create_btn = gr.Button("Generate key")
-        create_msg = gr.Textbox(label="New key (copy now!)", interactive=False, lines=2)
-
-        gr.Markdown("### All keys")
-        key_df = gr.DataFrame(
-            value=get_key_table(),
-            headers=["Label", "Key preview", "Status", "Usage", "Created"],
-            interactive=False,
-        )
-
-        gr.Markdown("### Revoke or delete a key")
-        with gr.Row():
-            manage_key_in = gr.Textbox(label="Full API key", placeholder="pii_sk_...")
-            revoke_btn    = gr.Button("Revoke")
-            delete_btn    = gr.Button("Delete", variant="stop")
-        manage_msg = gr.Textbox(label="Result", interactive=False)
-
-        create_btn.click(create_key, inputs=label_in,      outputs=[create_msg, key_df])
-        revoke_btn.click(do_revoke,  inputs=manage_key_in, outputs=[manage_msg, key_df])
-        delete_btn.click(do_delete,  inputs=manage_key_in, outputs=[manage_msg, key_df])
-
-demo.launch()
+
+@api.post("/detect")
+def detect_api(body: DetectRequest, request: Request):
+    _verify(request)
+    t0 = time.time()
+    try:
+        flat, cands = detect(body.text)
+    except Exception as exc:
+        raise HTTPException(status_code=503, detail=f"Inference failed: {exc}")
+
+    entities = [_fmt(body.text, e) for e in flat
+                if body.text[e["start"]:e["end"]].strip()]
+    nested = [_fmt(body.text, e)
+              for e in sorted(cands, key=lambda x: (x["start"], -(x["end"] - x["start"])))
+              if body.text[e["start"]:e["end"]].strip()]
+
+    redacted = body.text
+    for e in sorted(flat, key=lambda x: x["start"], reverse=True):
+        redacted = redacted[:e["start"]] + f'[{e["entity_group"]}]' + redacted[e["end"]:]
+
+    return {
+        "entities":           entities,
+        "nested":             nested,
+        "redacted_text":      redacted,
+        "model":              MODEL_REPO,
+        "processing_time_ms": int((time.time() - t0) * 1000),
+    }
+
+
+# Mount the Gradio UI at "/" on the same FastAPI app, so one Space serves both
+# the demo UI and the /detect + /health API.
+app = gr.mount_gradio_app(api, demo, path="/")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))