| """
|
| Stage 5 - The GUI.
|
|
|
| A Gradio Blocks interface over the four stages. It holds no detection logic: it uploads bytes,
|
| calls pipeline -> detector -> analysis -> embedder, and renders what comes back.
|
|
|
| **The interface comes up first, before anything is downloaded.** The user is asked for a Hugging
|
| Face read token, and only then does the app fetch anything - one part at a time, each ticked off as
|
| it finishes. The scanner appears when loading is done; the result tabs appear when a scan is done.
|
| Nothing empty is ever on screen.
|
|
|
| That ordering is deliberate. Gemma-2-9B in 4-bit is several gigabytes: loading it before `launch()`
|
| means minutes of blank screen and, on a Space, a container the platform may consider dead. It also
|
| turns a gated-repo refusal into a startup stack trace rather than a sentence in the UI.
|
|
|
| Four models, MiMo-7B by default. MiMo is not the most accurate - Gemma-2-9B is, by 0.024 F1 - but
|
| it is the only one of the four that anybody can download, so it is the one the app fetches at
|
| start-up and the one a user with no Hugging Face account can still use. Gemma is gated by Google
|
| and asks for a token; Qwen and Phi are open and load on request. Each entry in the picker states
|
| what it costs you, including Phi-4-mini's, which is that it flags 95% of clean files.
|
| """
|
|
|
| import os
|
| import random
|
| import sys
|
| import threading
|
| import time
|
| from pathlib import Path
|
|
|
| import gradio as gr
|
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
|
| import analysis
|
| import detector
|
| import embedder
|
| import pipeline
|
|
|
|
|
|
|
| __version__ = "5.1-arrows-with-card"
|
|
|
| DEFAULT_MODEL = detector.ACTIVE
|
|
|
|
|
|
|
| READY = {"ok": False, "model": DEFAULT_MODEL}
|
|
|
|
|
|
|
| ON_ZEROGPU = False
|
| try:
|
| import spaces
|
| ON_ZEROGPU = True
|
| except Exception:
|
| class _NoSpaces:
|
| @staticmethod
|
| def GPU(*a, **kw):
|
| def wrap(fn):
|
| return fn
|
| return wrap
|
| spaces = _NoSpaces()
|
|
|
| GPU_SLICE = 8
|
| GPU_DURATION = 120
|
|
|
|
|
| MAX_NEIGHBOURS = 5
|
|
|
| CSS = """
|
| /* Centre everything, at a width that keeps prose readable. */
|
| .gradio-container {max-width: 1000px !important; margin: 0 auto !important;}
|
| .gradio-container .prose {max-width: none !important;}
|
|
|
| #hero {text-align: center; margin: 0.4rem auto 1rem auto;}
|
| #hero h1 {margin-bottom: 0.3rem; font-size: 2rem;}
|
| #hero p {opacity: 0.75; margin: 0;}
|
|
|
| /* The start panel is a narrow card, centred. */
|
| #setup {max-width: 560px; margin: 0 auto 0.5rem auto; padding: 1.1rem 1.3rem;
|
| border-radius: 14px;}
|
| #setup h3 {margin-top: 0;}
|
|
|
| #steps {max-width: 560px; margin: 0.6rem auto; font-size: 0.94rem; line-height: 1.9;}
|
| #steps p {margin: 0.15rem 0;}
|
|
|
| /* The rotating note shown while something loads: a card that sits in front of the page, not a
|
| tint behind it. The colours come from the theme's
|
| own variables so it is opaque in both light and dark mode - the earlier rgba() tint let the
|
| background show through and the text read as if it were behind everything. */
|
| .waiting {max-width: 560px; margin: 1rem auto; padding: 1rem 1.15rem 1.1rem 1.15rem;
|
| border-radius: 14px; position: relative; z-index: 5;
|
| background: var(--block-background-fill, #ffffff);
|
| border: 1px solid var(--border-color-primary, rgba(128,128,128,0.35));
|
| box-shadow: 0 6px 20px rgba(0,0,0,0.22);}
|
| .waiting p {margin: 0.35rem 0; font-size: 0.95rem; line-height: 1.6;}
|
|
|
| /* The first line is the label - "While you're waiting" - and is set apart from the paragraph
|
| under it by a rule, so the eye finds the family name rather than a wall of text. */
|
| .waiting p:first-child {margin: 0; padding-bottom: 0.5rem; margin-bottom: 0.6rem;
|
| border-bottom: 1px solid var(--border-color-primary,
|
| rgba(128,128,128,0.25));
|
| font-size: 0.82rem; letter-spacing: 0.08em; text-transform: uppercase;
|
| opacity: 0.75;}
|
| .waiting p:first-child strong {font-weight: 700;}
|
| .waiting strong {font-weight: 700;}
|
|
|
| /* The one progress bar in this app. A class, not an id, because both phases - loading and
|
| scanning - draw their own, and Gradio's built-in progress is switched off everywhere
|
| (show_progress="hidden") so that this is the only bar a user ever sees. Plain divs rather than
|
| a component: it has to be redrawn from inside a generator, and markup is the only thing a
|
| generator can repaint. */
|
| .pbar {max-width: 560px; margin: 0.7rem auto 0.2rem auto;}
|
| .pbar .track {height: 8px; border-radius: 999px; overflow: hidden;
|
| background: var(--border-color-primary, rgba(128,128,128,0.25));}
|
| .pbar .fill {height: 100%; border-radius: 999px; transition: width 0.35s ease;
|
| background: linear-gradient(90deg, #f97316, #fb923c);}
|
| .pbar .label {display: flex; justify-content: space-between; gap: 1rem; margin-top: 0.45rem;
|
| font-size: 0.88rem; opacity: 0.8;}
|
| .pbar .label .pct {font-variant-numeric: tabular-nums; opacity: 0.7;}
|
|
|
| /* The arrows under a waiting card. Kept small, centred and quiet: they are an offer, not the
|
| thing to look at, and the card auto-advances perfectly well if they are never touched. */
|
| .wnav {max-width: 560px; margin: -0.4rem auto 0.8rem auto !important; justify-content: center;
|
| gap: 0.5rem;}
|
| .wnav button {flex: 0 0 auto !important; border-radius: 999px !important; font-size: 1.1rem;
|
| line-height: 1; padding: 0.25rem 0.9rem !important; opacity: 0.7;}
|
| .wnav button:hover {opacity: 1;}
|
|
|
| /* Tabs as tabs: a resting row of shapes with the active one lifted out of the panel, rather than
|
| three words with one underlined. The panel below joins the active tab, so the two read as one
|
| surface. */
|
| .gradio-container .tab-nav {border-bottom: 1px solid var(--border-color-primary,
|
| rgba(128,128,128,0.3)) !important;
|
| gap: 0.25rem; padding: 0 0.2rem;}
|
| .gradio-container .tab-nav button {border: 1px solid transparent !important;
|
| border-bottom: none !important;
|
| border-radius: 10px 10px 0 0 !important;
|
| padding: 0.55rem 1.05rem !important;
|
| font-size: 0.95rem; font-weight: 500; opacity: 0.72;
|
| margin-bottom: -1px;
|
| transition: background 0.15s ease, opacity 0.15s ease;}
|
| .gradio-container .tab-nav button:hover {opacity: 1;
|
| background: var(--background-fill-secondary,
|
| rgba(128,128,128,0.10)) !important;}
|
| .gradio-container .tab-nav button.selected {opacity: 1; font-weight: 600;
|
| background: var(--block-background-fill, rgba(128,128,128,0.16)) !important;
|
| border-color: var(--border-color-primary, rgba(128,128,128,0.3)) !important;
|
| box-shadow: inset 0 2px 0 0 #f97316;}
|
| .gradio-container .tabitem {border: 1px solid var(--border-color-primary,
|
| rgba(128,128,128,0.3)) !important;
|
| border-top: none !important; border-radius: 0 0 12px 12px !important;
|
| padding: 1rem 1.1rem !important;}
|
|
|
| #verdict {text-align: center; margin: 0.8rem 0 0.2rem 0;}
|
| #verdict h2 {margin: 0.2rem 0;}
|
| #cost {font-size: 0.92rem; opacity: 0.85;}
|
|
|
| /* Result tables: a quiet header row and hairline rules, so a five-column table reads as data
|
| rather than as a grid of boxes. */
|
| #report table th {background: var(--background-fill-secondary, rgba(128,128,128,0.10));
|
| font-weight: 600; font-size: 0.9rem; letter-spacing: 0.02em;}
|
| #report table, #report table td, #report table th {
|
| border-color: var(--border-color-primary, rgba(128,128,128,0.25)) !important;}
|
| #report h2 {font-size: 1.35rem; margin-top: 0.2rem;}
|
| #report h3 {font-size: 1.1rem;}
|
| #report hr {border: none; border-top: 1px solid var(--border-color-primary,
|
| rgba(128,128,128,0.25)); margin: 1.4rem 0;}
|
|
|
| /* Wrap everything: evidence strings and file ids are long and must not force a sideways scroll. */
|
| #report, #report * {overflow-wrap: anywhere; word-break: normal;}
|
| #report table {width: 100%; table-layout: fixed;}
|
| #report td, #report th {white-space: normal !important; overflow-wrap: anywhere;
|
| vertical-align: top; padding: 0.4rem 0.55rem;}
|
| #report code {overflow-wrap: anywhere; white-space: pre-wrap;}
|
| .gradio-container table td, .gradio-container table th {white-space: normal !important;
|
| overflow-wrap: anywhere;}
|
| """
|
|
|
|
|
|
|
| GRADIO_MAJOR = int(gr.__version__.split(".")[0])
|
| _BLOCKS_KW = {} if GRADIO_MAJOR >= 6 else {"css": CSS}
|
| _LAUNCH_KW = {"css": CSS} if GRADIO_MAJOR >= 6 else {}
|
|
|
|
|
|
|
|
|
|
|
|
|
| def _steps(model: str) -> list:
|
| """
|
| The checklist for one model. The first step exists only for a gated one.
|
|
|
| Showing "Check the token" while MiMo loads would be a lie about what the app is doing, and
|
| would suggest a token is needed when the whole point of the default is that none is.
|
| """
|
| spec = detector.DETECTORS[model]
|
| token_step = ["Check the token"] if spec.gated else []
|
| return token_step + ["Download the embedding index",
|
| f"Download {spec.label.split(' - ')[0]}",
|
| "Quantise to 4-bit"]
|
|
|
|
|
| def _checklist(model: str, done: int, current: str = "", failed: str = "") -> str:
|
| """The step list as it stands: finished ticked, current running, the rest waiting."""
|
| lines = []
|
| for i, name in enumerate(_steps(model)):
|
| if failed and i == done:
|
| lines.append(f"✗ **{name}** - {failed}")
|
| elif i < done:
|
| lines.append(f"✓ {name}")
|
| elif i == done and current:
|
| lines.append(f"● **{name}** - {current}")
|
| else:
|
| lines.append(f"<span style='opacity:0.45'>○ {name}</span>")
|
| return "\n\n".join(lines)
|
|
|
|
|
| FACT_SECONDS = 20.0
|
| FACT_POLL = 0.4
|
|
|
|
|
| def new_rotation():
|
| """
|
| A fresh reading order, shuffled, starting at the first entry.
|
|
|
| Re-made at the start of every load and every scan, so two runs in a row do not open on the
|
| same family. A shuffled list rather than an independent random draw each time: drawing
|
| independently repeats families and can leave some never shown at all.
|
| """
|
| order = list(detector.FAMILY_NOTES)
|
| random.shuffle(order)
|
| return {"order": order, "i": 0, "touched": time.monotonic()}
|
|
|
|
|
| def _ensure(rot):
|
| """Rotations live in gr.State, which starts as whatever the first caller finds."""
|
| if not rot:
|
| rot.update(new_rotation())
|
| return rot
|
|
|
|
|
| def rotation_note(rot, advance: bool = True) -> str:
|
| """
|
| The note to show now, advancing only if nobody has pressed an arrow for FACT_SECONDS.
|
|
|
| The clock is reset by the arrow handlers, so pressing one buys another full interval to read
|
| on - an auto-advance two seconds after a deliberate click would be the app arguing with the
|
| person using it.
|
| """
|
| _ensure(rot)
|
| if advance and time.monotonic() - rot["touched"] >= FACT_SECONDS:
|
| rot["i"] = (rot["i"] + 1) % len(rot["order"])
|
| rot["touched"] = time.monotonic()
|
| return _waiting_note(rot)
|
|
|
|
|
| def step_rotation(rot, delta: int):
|
| """An arrow press: move by hand, and restart the idle clock."""
|
| _ensure(rot)
|
| rot["i"] = (rot["i"] + delta) % len(rot["order"])
|
| rot["touched"] = time.monotonic()
|
| return _panel(_waiting_note(rot))
|
|
|
|
|
| def _family_title(family: str) -> str:
|
| """
|
| The family name as a person would write it.
|
|
|
| `.title()` on the registry key gives "Ssrf" and "Llm Prompt Injection" - acronyms are the one
|
| thing it cannot do, and these keys are mostly acronyms, so the exceptions are spelled out.
|
| """
|
| special = {
|
| "ssrf": "SSRF - server-side request forgery",
|
| "llm_prompt_injection": "LLM prompt injection",
|
| "xfa_acroform_injection": "XFA / AcroForm injection",
|
| "dde_template_injection": "DDE template injection",
|
| "uri_redirect_phishing": "URI redirect phishing",
|
| "cross_site_scripting": "Cross-site scripting (XSS)",
|
| }
|
| return special.get(family, family.replace("_", " ").capitalize())
|
|
|
|
|
| def _waiting_note(rot) -> str:
|
| """
|
| One family, as three paragraphs: label, name, explanation.
|
|
|
| Split into paragraphs rather than one block because the CSS styles the first one as the card's
|
| heading; a single paragraph would give a wall of text with a bold phrase buried in it. The
|
| counter tells the reader the arrows lead somewhere and how far round they are.
|
| """
|
| family = rot["order"][rot["i"]]
|
| counter = f"{rot['i'] + 1} of {len(rot['order'])}"
|
| return (f"**While you're waiting** · {counter}\n\n"
|
| f"**{_family_title(family)}**\n\n"
|
| f"{detector.FAMILY_NOTES[family]}")
|
|
|
|
|
| def _nav(text: str):
|
| """
|
| The arrow row that belongs to a waiting card: on screen exactly when the card is.
|
|
|
| Arrows with nothing to page through are a control that does nothing, which is worse than no
|
| control at all - so they are tied to the same string the card is.
|
| """
|
| return gr.update(visible=bool(text))
|
|
|
|
|
| def _panel(text: str):
|
| """
|
| A card that is on screen only when it has something to say.
|
|
|
| An empty Markdown still draws its padding and border, so without this the loading screen
|
| carries an empty box between every note.
|
| """
|
| return gr.update(value=text, visible=bool(text))
|
|
|
|
|
| def _while_working(work, render, rot):
|
| """
|
| Run `work()` on a thread and keep the screen talking while it runs.
|
|
|
| A generator delegated to with `yield from`: it yields whatever `render(note)` builds - a full
|
| set of UI updates - and returns whatever `work()` returned. Downloading six gigabytes is one
|
| blocking call that cannot report from inside itself, so the only way to say anything during it
|
| is to put it on another thread and watch.
|
|
|
| It re-reads the rotation every FACT_POLL and yields only when the text has actually changed,
|
| which covers both ways the note can move: the idle timer running out, and the reader pressing
|
| an arrow - the arrow handler mutates the same rotation object this loop is reading.
|
|
|
| An exception raised on the thread is re-raised here, so the caller's existing `except` still
|
| sees a failed download as a failed download.
|
| """
|
| box = {}
|
|
|
| def runner():
|
| try:
|
| box["value"] = work()
|
| except BaseException as e:
|
| box["error"] = e
|
|
|
| thread = threading.Thread(target=runner, daemon=True)
|
| thread.start()
|
|
|
| shown = rotation_note(rot, advance=False)
|
| yield render(shown)
|
|
|
| while thread.is_alive():
|
| current = rotation_note(rot)
|
| if current != shown:
|
| shown = current
|
| yield render(shown)
|
| time.sleep(FACT_POLL)
|
|
|
| thread.join()
|
| if "error" in box:
|
| raise box["error"]
|
| return box.get("value")
|
|
|
|
|
| def preflight(model: str) -> tuple:
|
| """
|
| Is the token real, and does this account have this model? Returns (status, message).
|
|
|
| Three answers, not two, because a user's next move differs in each: `ok`, `no_access` (the
|
| token is fine, the licence has not been accepted) and `bad_token` (the string is not a token).
|
| Collapsing the middle case into "it failed" is the commonest mistake here - a perfectly good
|
| token on an unaccepted licence then reads as a typo, and the user retypes it forever.
|
|
|
| Two wrong ways to ask, both of which pass when they should not:
|
|
|
| * `model_info` - Gemma is `gated=manual`, which gates the *weights* and leaves the model card
|
| and metadata public, so this succeeds for anybody, garbage token included.
|
| * `hf_hub_download("config.json")` - answers from the local cache once the file has been
|
| fetched before, so it keeps saying yes after access is revoked.
|
|
|
| `auth_check` asks the one question that matters and never answers from cache.
|
| """
|
| spec = detector.DETECTORS[model]
|
| if not spec.gated:
|
| return "ok", "no token needed"
|
|
|
| token = detector.hf_token()
|
| if not token:
|
| return "bad_token", "no token given, and none found in this runtime"
|
|
|
|
|
|
|
| try:
|
| from huggingface_hub import HfApi
|
| HfApi().whoami(token=token)
|
| except Exception as e:
|
| return "bad_token", f"{type(e).__name__}: {e}"
|
|
|
| try:
|
| try:
|
| from huggingface_hub import auth_check
|
| auth_check(spec.repo, token=token)
|
| except ImportError:
|
| from huggingface_hub import get_hf_file_metadata, hf_hub_url
|
| get_hf_file_metadata(hf_hub_url(spec.repo, "config.json"), token=token)
|
| return "ok", "accepted"
|
| except Exception as e:
|
| name = type(e).__name__
|
| if "Gated" in name or "Repository" in name or "401" in str(e) or "403" in str(e):
|
| return "no_access", detector.gate_url(model)
|
| return "error", f"{name}: {e}"
|
|
|
|
|
| def token_hint() -> str:
|
| if detector.hf_token():
|
| return ("A token was found in this runtime. Leave the box empty and press **Send token** "
|
| "to use it.")
|
| return ("Paste a **read** token from huggingface.co/settings/tokens. It is used only to "
|
| "download the model, and is never stored.")
|
|
|
|
|
|
|
|
|
|
|
|
|
| _PREFETCH = {"thread": None}
|
|
|
|
|
| def prefetch_default():
|
| """
|
| Start pulling MiMo into the cache as the app boots, on a background thread.
|
|
|
| On a thread rather than inline: the interface has to be answering requests within seconds or a
|
| Space is considered dead, and a user who intends to pick Gemma anyway should not wait behind a
|
| download they will not use. By the time anyone has read the page and pressed a button the
|
| fetch is usually finished, and `connect()` then ticks the download step off immediately.
|
|
|
| Only the default is prefetched, and only because it is ungated - there is nothing to prefetch
|
| for Gemma until a token exists, and prefetching all four would pull twenty-odd gigabytes for
|
| three models nobody asked for.
|
| """
|
| if _PREFETCH["thread"] is not None:
|
| return _PREFETCH["thread"]
|
|
|
| def work():
|
| try:
|
| from huggingface_hub import snapshot_download
|
| snapshot_download(detector.DETECTORS[DEFAULT_MODEL].repo,
|
| allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"])
|
| except Exception as e:
|
| print(f"[prefetch] {type(e).__name__}: {e}")
|
|
|
| thread = threading.Thread(target=work, daemon=True)
|
| thread.start()
|
| _PREFETCH["thread"] = thread
|
| return thread
|
|
|
|
|
| def connect(model: str, token: str, rot):
|
| """
|
| Everything between "the page loaded" and "you can scan a PDF", one part at a time.
|
|
|
| A generator: each part reports before it starts and ticks off when it finishes, so the bar
|
| advances four times rather than sitting still through a six-gigabyte download.
|
|
|
| Yields (checklist, setup panel, scanner panel, waiting note, token verdict, progress bar).
|
|
|
| No `gr.Progress`: Gradio draws its own bar over every output component of an event, which on
|
| a five-output event meant the same percentage rendered several times down the page. One bar,
|
| drawn by `_bar`, with Gradio's own switched off at the event (`show_progress="hidden"`).
|
| """
|
| steps = _steps(model)
|
| spec = detector.DETECTORS[model]
|
| rot.update(new_rotation())
|
|
|
| def state(done, current="", failed="", waiting="", verdict=""):
|
| finished = done >= len(steps)
|
| return (_checklist(model, done, current, failed),
|
| gr.update(visible=not finished),
|
| gr.update(visible=finished),
|
| _panel(waiting),
|
| _panel(verdict),
|
| "" if finished else _bar(done, len(steps), steps[done]),
|
| _nav(waiting))
|
|
|
| READY["ok"], READY["model"] = False, model
|
|
|
|
|
| step = 0
|
| if spec.gated:
|
| token = (token or "").strip()
|
| if token:
|
| os.environ["HF_TOKEN"] = token
|
| yield state(0, "asking Hugging Face")
|
|
|
| status, detail = preflight(model)
|
| if status == "no_access":
|
| yield state(0, failed="the licence has not been accepted",
|
| verdict=(f"### 🔒 Access not granted\n\nYour token works, but "
|
| f"this account has not been given {spec.label.split(' - ')[0]}. Accept "
|
| f"the licence here, signed in as the same account, then press "
|
| f"**Send token** again:\n\nURL: {detail}"))
|
| return
|
| if status == "bad_token":
|
| yield state(0, failed="token refused",
|
| verdict=("### ❌ The Hugging Face token was not approved\n\n"
|
| "Check it at huggingface.co/settings/tokens and try again, or "
|
| "pick a different model - MiMo, Qwen and Phi need no token."))
|
| return
|
| if status != "ok":
|
| yield state(0, failed=detail,
|
| verdict=f"### ⚠️ Could not reach Hugging Face\n\n`{detail}`")
|
| return
|
|
|
| step = 1
|
| yield state(1, verdict=f"### ✅ Token approved, loading {spec.label.split(' - ')[0]}…")
|
|
|
|
|
| keep = f"### ✅ Token approved, loading {spec.label.split(' - ')[0]}…" if spec.gated else ""
|
| yield state(step, "5 MB", verdict=keep)
|
| try:
|
| idx = yield from _while_working(
|
| embedder.load_index, lambda note: state(step, "5 MB", waiting=note, verdict=keep), rot)
|
| n = idx["matrix"].shape[0]
|
| except Exception as e:
|
| yield state(step, failed=f"{type(e).__name__}: {e}",
|
| verdict=f"### ⚠️ Could not load the embedding index\n\n`{e}`")
|
| return
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| size_note = "about 6 GB, cached after the first run"
|
| yield state(step + 1, size_note, verdict=keep)
|
|
|
| def _fetch_weights():
|
| from huggingface_hub import snapshot_download
|
| return snapshot_download(spec.repo, token=detector.hf_token(),
|
| allow_patterns=["*.json", "*.safetensors", "*.model",
|
| "tokenizer*"])
|
|
|
| try:
|
| yield from _while_working(
|
| _fetch_weights,
|
| lambda note: state(step + 1, size_note, waiting=note, verdict=keep), rot)
|
| except Exception as e:
|
| yield state(step + 1, failed=f"{type(e).__name__}: {e}",
|
| verdict=f"### ⚠️ Download failed\n\n`{e}`")
|
| return
|
|
|
|
|
| if ON_ZEROGPU:
|
|
|
| yield state(step + 2, "deferred to the first scan (ZeroGPU has no device until then)",
|
| verdict=keep)
|
| time.sleep(0.4)
|
| else:
|
| yield state(step + 2, "loading onto the GPU", verdict=keep)
|
| try:
|
| yield from _while_working(
|
| lambda: detector.load(model),
|
| lambda note: state(step + 2, "loading onto the GPU", waiting=note, verdict=keep),
|
| rot)
|
| except Exception as e:
|
|
|
|
|
| yield state(step + 2, failed=f"{type(e).__name__}",
|
| verdict=f"### ⚠️ The model could not be loaded\n\n{e}")
|
| return
|
|
|
| READY["ok"] = True
|
| yield state(len(steps),
|
| verdict=f"### ✅ {spec.label.split(' - ')[0]} is loaded and ready.")
|
|
|
|
|
|
|
|
|
|
|
|
|
| @spaces.GPU(duration=GPU_DURATION)
|
| def _scan_slice(windows):
|
| """One GPU allocation's worth of windows. The only function that touches the GPU."""
|
| return detector.scan(windows, name=READY["model"])
|
|
|
|
|
| def _bar(done: int, total: int, note: str = "") -> str:
|
| """
|
| The progress bar, drawn from the section counts rather than from elapsed time.
|
|
|
| A time-based bar has to guess, and guesses wrongly on the first run of a document twice the
|
| size of the last one. Sections are the unit of work the scan actually consumes, and their
|
| number is known before the first one starts, so the bar is a measurement.
|
| """
|
| pct = 100.0 * done / max(total, 1)
|
| label = note or f"Section {min(done + 1, total)} of {total}"
|
| return (f"<div class='pbar'><div class='track'>"
|
| f"<div class='fill' style='width:{pct:.1f}%'></div></div>"
|
| f"<div class='label'><span>{label}</span>"
|
| f"<span class='pct'>{pct:.0f}%</span></div></div>")
|
|
|
|
|
| def describe(file):
|
| """What the scan will cost, before committing to it."""
|
| if file is None:
|
| return ""
|
| doc = pipeline.pdf_to_text(Path(file).read_bytes())
|
| n = len(pipeline.window_document(doc["masked"]))
|
| total = n * detector.DETECTORS[READY["model"]].seconds_per_window
|
| unit = f"about {total:.0f} seconds" if total < 90 else f"about {total / 60:.0f} minutes"
|
| note = " \n*Large file - trimmed to its first and last sections.*" if doc["was_truncated"] else ""
|
| return (f"**{doc['chars']:,} characters** in **{n} overlapping sections**. "
|
| f"On a GPU that is {unit}.{note}")
|
|
|
|
|
| def analyse(file, k, want_neighbours, rot):
|
| """
|
| Yields (verdict, results visibility, window rows, report, neighbour rows, basis, bar, note).
|
|
|
| A generator, not a function, because the bar has to move *during* the scan: the sections are
|
| scanned a slice at a time and the bar is repainted between slices, so a 40-section document
|
| shows forty steps of evidence that it is alive rather than one long freeze.
|
|
|
| A scan of a large document is the longer of the two waits - 41 sections is close to four
|
| minutes - so it carries the same rotating notes as the download, under its own bar.
|
| """
|
| hide = gr.update(visible=False)
|
| rot.update(new_rotation())
|
|
|
| def stop(msg, bar="", note=""):
|
| return msg, hide, [], "", [], "", bar, _panel(note), _nav(note)
|
|
|
| if not READY["ok"]:
|
| yield stop("### Press **Start** above first")
|
| return
|
| if file is None:
|
| yield stop("### Choose a PDF, then press **Check this PDF**")
|
| return
|
|
|
| t0 = time.time()
|
| yield stop("", _bar(0, 1, "Reading the PDF"), rotation_note(rot))
|
| doc = pipeline.pdf_to_text(Path(file).read_bytes())
|
| windows = pipeline.window_document(doc["masked"])
|
| if not windows:
|
| yield stop("### Nothing readable in this file\n\nIt may be an image-only scan.")
|
| return
|
|
|
| total = len(windows)
|
| yield stop("", _bar(0, total), rotation_note(rot))
|
|
|
|
|
|
|
|
|
| findings = []
|
| try:
|
| for i in range(0, total, GPU_SLICE):
|
| findings += _scan_slice(windows[i:i + GPU_SLICE])
|
| done = len(findings)
|
| yield stop("", _bar(done, total), rotation_note(rot))
|
| except Exception as e:
|
| yield stop(f"### Scan failed\n\n`{type(e).__name__}: {e}`")
|
| return
|
| finally:
|
| if ON_ZEROGPU:
|
| detector.release()
|
|
|
| yield stop("", _bar(total, total, "Building the report"), rotation_note(rot))
|
|
|
| summary = analysis.consolidate(findings, doc, model=READY["model"])
|
| report = analysis.report_markdown(summary)
|
|
|
| rows, basis_text = [], ""
|
| if want_neighbours:
|
| yield stop("", _bar(total, total, "Finding similar files"), rotation_note(rot))
|
| try:
|
| nb = embedder.neighbours_for(summary, windows, k=int(k))
|
| rows = embedder.neighbour_rows(nb)
|
| basis_text = f"Compared using {nb['basis']}."
|
| except Exception as e:
|
| basis_text = f"Similar-file lookup unavailable ({type(e).__name__})."
|
|
|
| if summary["injected"]:
|
| verdict_text = (f"## 🔴 Something is hidden in this PDF\n\n**{summary['n_regions']}** "
|
| f"affected {'part' if summary['n_regions'] == 1 else 'parts'} · "
|
| f"{summary['n_windows_flagged']} of {summary['n_windows']} sections "
|
| f"flagged · {time.time() - t0:.0f}s")
|
| else:
|
| verdict_text = (f"## 🟢 Nothing suspicious found\n\nAll **{summary['n_windows']}** "
|
| f"sections checked · {time.time() - t0:.0f}s")
|
|
|
|
|
| yield (verdict_text, gr.update(visible=True), analysis.window_rows(findings),
|
| report, rows, basis_text, "", _panel(""), _nav(""))
|
|
|
|
|
| def on_pick(model: str):
|
| """
|
| React to the picker: ask for a token only where one can do something.
|
|
|
| Returns (what this model costs you, token row, Start button, cleared verdict, checklist).
|
| Gemma replaces Start with Send token rather than showing both, so there is exactly one thing
|
| to press, and the checklist is redrawn because a gated model has a step the others do not.
|
| """
|
| spec = detector.DETECTORS[model]
|
| return (spec.summary,
|
| gr.update(visible=spec.gated),
|
| gr.update(visible=not spec.gated),
|
| _panel(""),
|
| _checklist(model, 0))
|
|
|
|
|
| def _lock():
|
| """Take the button away for the duration of a scan, and say why it is gone."""
|
| return gr.update(interactive=False, value="Checking…")
|
|
|
|
|
| def _unlock():
|
| return gr.update(interactive=True, value="Check this PDF")
|
|
|
|
|
|
|
|
|
|
|
|
|
| with gr.Blocks(title="PDF Injection Detector", **_BLOCKS_KW) as demo:
|
| gr.Markdown(
|
| "# PDF Injection Detector\n"
|
| "Upload a PDF and find out whether something has been hidden inside it - where it is, "
|
| "what kind of thing it is, and the exact text that gave it away.",
|
| elem_id="hero")
|
|
|
|
|
| with gr.Group(visible=True, elem_id="setup") as setup_panel:
|
| gr.Markdown("### Start")
|
| model_pick = gr.Dropdown(
|
| choices=[(spec.label, key) for key, spec in detector.DETECTORS.items()],
|
| value=DEFAULT_MODEL, label="Detection model")
|
| model_note = gr.Markdown(detector.DETECTORS[DEFAULT_MODEL].summary)
|
|
|
|
|
|
|
| with gr.Group(visible=False) as token_row:
|
| token_box = gr.Textbox(label="Hugging Face read token", type="password",
|
| placeholder="hf_...", lines=1)
|
| token_help = gr.Markdown(token_hint())
|
| send_token = gr.Button("Send token", variant="primary", size="lg")
|
|
|
| start = gr.Button("Start", variant="primary", size="lg")
|
|
|
| token_verdict = gr.Markdown(visible=False, elem_classes=["waiting"])
|
|
|
| steps_view = gr.Markdown(_checklist(DEFAULT_MODEL, 0), elem_id="steps")
|
| load_bar = gr.HTML()
|
| waiting_view = gr.Markdown(visible=False, elem_classes=["waiting"])
|
| with gr.Row(elem_classes=["wnav"], visible=False) as load_nav:
|
| load_prev = gr.Button("‹", size="sm", min_width=48)
|
| load_next = gr.Button("›", size="sm", min_width=48)
|
|
|
|
|
| with gr.Group(visible=False) as scanner:
|
| with gr.Row():
|
| with gr.Column(scale=3):
|
| file = gr.File(label="Your PDF", file_types=[".pdf"], type="filepath")
|
| with gr.Column(scale=2):
|
| want_neighbours = gr.Checkbox(value=True, label="Also show similar known files")
|
| k = gr.Slider(1, MAX_NEIGHBOURS, value=MAX_NEIGHBOURS, step=1,
|
| label="How many similar files")
|
| cost = gr.Markdown(elem_id="cost")
|
| run = gr.Button("Check this PDF", variant="primary", size="lg")
|
| bar = gr.HTML()
|
| scan_note = gr.Markdown(visible=False, elem_classes=["waiting"])
|
| with gr.Row(elem_classes=["wnav"], visible=False) as scan_nav:
|
| scan_prev = gr.Button("‹", size="sm", min_width=48)
|
| scan_next = gr.Button("›", size="sm", min_width=48)
|
| verdict = gr.Markdown(elem_id="verdict")
|
|
|
|
|
| with gr.Group(visible=False) as results:
|
| with gr.Tabs():
|
| with gr.Tab("Report"):
|
| report = gr.Markdown(elem_id="report")
|
| with gr.Tab("Section by section"):
|
| gr.Markdown("Every section of the document, in order, with what the model "
|
| "said about each one.")
|
| windows_table = gr.Dataframe(headers=analysis.WINDOW_COLUMNS, wrap=True)
|
| with gr.Tab("Similar known files"):
|
| gr.Markdown("The closest matches among the 1,100 files this system was "
|
| "tested on. These are the most *similar* files, not a second "
|
| "opinion on the verdict.")
|
| neighbours_table = gr.Dataframe(headers=embedder.NEIGHBOUR_COLUMNS, wrap=True)
|
| basis = gr.Markdown()
|
|
|
|
|
|
|
|
|
| rotation = gr.State({})
|
|
|
|
|
|
|
| _connect_out = [steps_view, setup_panel, scanner, waiting_view, token_verdict, load_bar,
|
| load_nav]
|
| model_pick.change(on_pick, model_pick,
|
| [model_note, token_row, start, token_verdict, steps_view],
|
| show_progress="hidden")
|
| start.click(connect, [model_pick, token_box, rotation], _connect_out, show_progress="hidden")
|
| send_token.click(connect, [model_pick, token_box, rotation], _connect_out,
|
| show_progress="hidden")
|
| file.change(describe, file, cost, show_progress="hidden")
|
|
|
| for button, delta, target in ((load_prev, -1, waiting_view), (load_next, 1, waiting_view),
|
| (scan_prev, -1, scan_note), (scan_next, 1, scan_note)):
|
| button.click(lambda rot, d=delta: step_rotation(rot, d), rotation, target,
|
| queue=False, show_progress="hidden")
|
|
|
|
|
|
|
|
|
|
|
| (run.click(_lock, None, run, queue=False, show_progress="hidden")
|
| .then(analyse, [file, k, want_neighbours, rotation],
|
| [verdict, results, windows_table, report, neighbours_table, basis, bar,
|
| scan_note, scan_nav], show_progress="hidden")
|
| .then(_unlock, None, run, queue=False, show_progress="hidden"))
|
|
|
|
|
| def launch(**kw):
|
| """
|
| Serve immediately, with the default model already on its way.
|
|
|
| `prefetch_default()` returns at once - the download runs on a daemon thread - so this is still
|
| the same instant launch, only with MiMo usually cached by the time anyone presses a button.
|
| """
|
| prefetch_default()
|
| return demo.launch(**{**_LAUNCH_KW, **kw})
|
|
|
|
|
| if __name__ == "__main__":
|
| prefetch_default()
|
| demo.launch(**_LAUNCH_KW,
|
| server_name="0.0.0.0" if os.environ.get("SPACE_ID") else "127.0.0.1") |
