Spaces:
Running on Zero
Running on Zero
| """ | |
| PDF Injection Detector - MiMo-7B. | |
| Upload a PDF, and this reads it the way the corpus was read, cuts it into batches of regions that | |
| fit one run of the model, and asks MiMo-7B-RL whether a payload is hidden in the batch you choose. | |
| `app.py` holds the interface, the batching and the aggregation. It contains no detection logic of | |
| its own: the text extraction lives in `corpus_text.py`, the prompt and parser in `mimo.py`, the | |
| embedding lookup in `neighbours.py`, and each is quoted from the notebook that measured it. | |
| """ | |
| import re | |
| import traceback | |
| from pathlib import Path | |
| import gradio as gr | |
| import corpus_text | |
| import mimo | |
| import neighbours | |
| # What to do about each family, shown under a payload verdict. | |
| # | |
| # One entry per family in `mimo.FAMILIES`, and the keys are asserted against that list at import so | |
| # a renamed family cannot silently lose its advice. Each is written against the generator's own | |
| # object shapes rather than the general security literature, so a reader who opens a flagged file | |
| # finds the thing the text describes. | |
| # | |
| # **This is advice conditional on a guess.** MiMo names the family correctly 43% of the time, so | |
| # more often than not the paragraph shown here is the remedy for a different attack than the one | |
| # present. It is worth showing anyway - the first line of every entry is a containment step that | |
| # holds whatever the family turns out to be - but the interface says so rather than implying the | |
| # app knows what it is looking at. | |
| TREATMENT = { | |
| "javascript_injection": ( | |
| "Quarantine the file immediately. Do not open in Adobe Reader or any JS-enabled viewer. " | |
| "Strip the /JS and /JavaScript PDF objects using a PDF sanitiser (e.g. qpdf or mutool). " | |
| "Disable JavaScript in PDF readers organisation-wide via Group Policy."), | |
| "cross_site_scripting": ( | |
| "Do not open in a browser-based PDF viewer. Use a sandboxed desktop reader only. " | |
| "Submit to your security team for content stripping before redistribution."), | |
| "ssrf": ( | |
| "Block outbound HTTP requests from any server that processes this file. " | |
| "Do not open on cloud infrastructure without egress filtering - the payload targets " | |
| "the AWS/GCP metadata endpoint (169.254.169.254). Flag for security review."), | |
| "object_action_injection": ( | |
| "Open only in read-only sandboxed mode. Strip /Launch, /OpenAction and /AA objects " | |
| "using a PDF sanitiser. Alert your SOC team - this payload attempts to execute a " | |
| "program outside the PDF viewer, with the opening user's own permissions."), | |
| "llm_prompt_injection": ( | |
| "Do not feed this document to any AI/LLM pipeline without human review first. " | |
| "The payload attempts to override AI instructions. Sanitise or reject the file " | |
| "before any automated processing."), | |
| "shellcode_embedded_exe": ( | |
| "Quarantine immediately and isolate the machine. Run through antivirus before any " | |
| "attachment is extracted, and do not let anyone double-click what comes out. " | |
| "Report the binary stream hash to your threat intelligence team."), | |
| "dde_template_injection": ( | |
| "Do not copy any field contents into Excel or Word - the payload activates as a DDE " | |
| "formula once it reaches Office, not while the PDF is open. Strip /AcroForm field " | |
| "values and any remote template reference, and block outbound requests to the linked " | |
| "template host."), | |
| "polyglot_file": ( | |
| "Treat this as two files, not one. Scan it again as an archive as well as a PDF: " | |
| "whichever format your scanner did not choose was never inspected at all. Do not let " | |
| "it through a filter that decides file type from the extension or the leading bytes " | |
| "alone."), | |
| "ransomware_simulation": ( | |
| "Isolate the host before anything else, and do not open the file. Confirm backups are " | |
| "offline and restorable, then hand the sample to incident response. In this corpus the " | |
| "payload is a harmless RANSIM/EICAR simulation - an unknown file of this shape should " | |
| "be treated as the real thing until proven otherwise."), | |
| "steganographic_payload": ( | |
| "Do not forward the file. Extract and inspect the embedded image streams separately: " | |
| "the payload hides in pixel data and survives filters that look for code or links. " | |
| "Re-encode or drop the images before any redistribution."), | |
| "uri_redirect_phishing": ( | |
| "Do not click anywhere in the document - the link annotation covers the entire page, " | |
| "so any click opens it. Verify the destination host out of band, strip /URI actions " | |
| "with a PDF sanitiser, and report the URL to your phishing intake."), | |
| "xfa_acroform_injection": ( | |
| "Strip the XFA form definition before the file is processed further. The payload lives " | |
| "in an XML document inside the PDF, so scanners that parse only PDF objects never see " | |
| "it. Do not open in a reader with XFA support enabled."), | |
| } | |
| # A verdict can be `injected=true` with `injection_type="none"` - the model is sure something is | |
| # wrong and will not say what. That still deserves an answer. | |
| TREATMENT_UNNAMED = ( | |
| "MiMo flagged this region but would not name a family, so no specific remediation applies. " | |
| "Treat the file as untrusted: do not open it on a machine that matters, do not forward it, " | |
| "and pass it to whoever handles security for you.") | |
| assert set(TREATMENT) == set(mimo.FAMILIES), ( | |
| f"TREATMENT does not cover mimo.FAMILIES: " | |
| f"missing {sorted(set(mimo.FAMILIES) - set(TREATMENT))}, " | |
| f"unknown {sorted(set(TREATMENT) - set(mimo.FAMILIES))}") | |
| # Measured in Part B on a T4, on the project's own 1,100-document corpus. | |
| PART_B = {"f1": 0.945, "precision": 0.988, "recall": 0.906, "family_acc": 0.433, | |
| "false_alarm_rate": 0.05, "unparsable": 155, "n": 1100} | |
| GPU = mimo.BACKEND == "gpu" | |
| # How many regions fit in ONE run is not a taste decision. On ZeroGPU a single grant is capped at | |
| # 300 seconds and the whole scan plus the 4-bit load must fit inside it; on 2 vCPUs a region costs | |
| # two minutes and a browser will not wait for many. That ceiling is what a batch is. | |
| MAX_PER_BATCH = mimo.MAX_WINDOWS_GPU if GPU else 6 | |
| DEFAULT_PER_BATCH = 8 if GPU else 3 | |
| WINDOW_COLUMNS = ["#", "where in the skeleton", "kind", "markers", "verdict", "family", "evidence"] | |
| # One example per injection family, plus one clean control, taken from the generation repo. | |
| # | |
| # The files on disk are named Example_N and nothing else, and that stays true: the uploader shows | |
| # a neutral filename, and nothing about the file announces its own answer. The *buttons* are | |
| # labelled by family so the demo can be driven deliberately ("show me ransomware"), which costs | |
| # nothing on the model side - a filename never reaches the prompt, only extracted text does. | |
| # | |
| # Labels come from `examples/manifest.json`, written by the same script that copies the PDFs, so | |
| # a button cannot end up pointing at the wrong family. `examples/Example_Key.txt` is the same | |
| # mapping in prose, for whoever is marking this. | |
| EXAMPLE_DIR = Path(__file__).resolve().parent / "examples" | |
| def load_examples(): | |
| """[(label, path)], families first in the model's own order, clean last. Empty if absent.""" | |
| manifest = EXAMPLE_DIR / "manifest.json" | |
| if not manifest.exists(): | |
| return [] | |
| import json | |
| mapping = json.loads(manifest.read_text(encoding="utf-8")) | |
| order = {f: i for i, f in enumerate(mimo.FAMILIES)} | |
| rows = [(fam, EXAMPLE_DIR / name) for name, fam in mapping.items() | |
| if (EXAMPLE_DIR / name).exists()] | |
| # "clean" is not in FAMILIES, so it sorts last - which is where it belongs: it is the control, | |
| # read after you have seen what a hit looks like. | |
| return sorted(rows, key=lambda r: order.get(r[0], len(order))) | |
| EXAMPLES = load_examples() | |
| def fmt_eta(n: int, runtime: str = None) -> str: | |
| seconds = n * mimo.SECONDS[runtime or mimo.BACKEND] | |
| if seconds < 90: | |
| return f"about {max(20, int(seconds))}s" | |
| return f"{seconds * 0.6 / 60:.0f}-{seconds * 1.5 / 60:.0f} min" | |
| def extract(path, cover_all=True): | |
| """PDF bytes to ranked candidate regions. No model touched, so this runs on upload.""" | |
| with open(path, "rb") as fh: | |
| data = fh.read() | |
| skeleton, truncated, dropped = corpus_text.build_skeleton(data) | |
| return (data, skeleton, truncated, dropped, | |
| corpus_text.candidate_windows(skeleton, cover_all=bool(cover_all))) | |
| def batches_of(candidates, per_batch): | |
| """The ranked regions cut into runnable chunks. Batch 1 is the most marker-dense.""" | |
| per_batch = max(1, int(per_batch)) | |
| return [candidates[i:i + per_batch] for i in range(0, len(candidates), per_batch)] | |
| def batch_label(batch, i, n_batches, runtime=None): | |
| """What one batch is, in a line, so the choice is informed rather than a number.""" | |
| first, last = batch[0], batch[-1] | |
| if first["is_head"]: | |
| where = "head of document" | |
| else: | |
| where = f"chars {min(w['start'] for w in batch):,}-{max(w['end'] for w in batch):,}" | |
| fams = sorted({f for w in batch for f in w["families"]}) | |
| n_marker = sum(1 for w in batch if w["source"] == "marker") | |
| kind = ("marker regions" if n_marker == len(batch) else | |
| "sweep of the document" if n_marker == 0 else | |
| f"{n_marker} marker + {len(batch) - n_marker} sweep") | |
| # "signature:" and not the bare family name. The regex has only seen a token like `/JS (` in | |
| # the text; plenty of harmless PDFs contain one. Printing "javascript_injection" on its own | |
| # before the model has read anything reads as a verdict the app has not made. | |
| tail = f" · signature: {', '.join(fams)}" if fams else "" | |
| return (f"Batch {i + 1} of {n_batches} — {len(batch)} region(s), {kind} · {where}{tail} · " | |
| f"~{fmt_eta(len(batch), runtime)}") | |
| def on_upload(path, per_batch, cover_all, runtime): | |
| """Extract, rank and batch. Fast enough to run on every upload and every slider move.""" | |
| if not path: | |
| return (None, "Upload a PDF to see what will be read.", | |
| gr.update(choices=[], value=None, interactive=False)) | |
| try: | |
| data, skeleton, truncated, dropped, candidates = extract(path, cover_all) | |
| except Exception as e: | |
| return (None, f"Could not read that file: `{type(e).__name__}: {e}`", | |
| gr.update(choices=[], value=None, interactive=False)) | |
| groups = batches_of(candidates, per_batch) | |
| # Plain label strings, NOT (label, index) pairs. The dropdown allows custom values so the API | |
| # can name a batch before any PDF has been uploaded, and that turns it into a free-text | |
| # combobox: a programmatic integer value of 0 is falsy, so the displayed text refused to | |
| # refresh and the box kept showing the previous split after the slider moved. Strings are | |
| # never falsy here, and `resolve_batch` reads the index straight back out of "Batch N of M". | |
| choices = [batch_label(b, i, len(groups), runtime) for i, b in enumerate(groups)] | |
| lines = [ | |
| f"**{len(data):,} bytes** on disk, rendered to a **{len(skeleton):,}-character skeleton**" | |
| + (f" (truncated to the {corpus_text.SKELETON_CHAR_BUDGET:,}-character budget)" | |
| if truncated else "") | |
| + (f", {dropped} binary stream(s) dropped." if dropped else "."), | |
| ] | |
| n_marker = sum(1 for w in candidates if w["source"] == "marker") | |
| fams = corpus_text.detect_markers(data) | |
| if n_marker: | |
| lines.append( | |
| f"**{n_marker} region(s) carry a marker**, and the remaining " | |
| f"{len(candidates) - n_marker} cover the rest of the document — " | |
| f"**{len(candidates)} in total, cut into {len(groups)} batch(es)** of at most " | |
| f"{int(per_batch)}. Structural signatures in the raw file: " | |
| f"`{'`, `'.join(fams) or 'none'}`.") | |
| lines.append( | |
| "_Ranking decides **reading order only** — batch 1 is the most marker-dense, not the " | |
| "guilty one. The verdict is MiMo's alone._") | |
| elif cover_all: | |
| lines.append( | |
| f"**No structural marker anywhere in the file.** The triage only recognises the twelve " | |
| f"families this project generated, so this means either a clean file or a payload " | |
| f"shaped like none of them — which is why the batches below sweep the **whole** " | |
| f"skeleton rather than stopping here: **{len(candidates)} region(s) in " | |
| f"{len(groups)} batch(es)**.") | |
| else: | |
| lines.append( | |
| "**No structural marker anywhere in the file**, and the sweep is switched off — so " | |
| "only the head of the document will be read, which is exactly what the corpus builder " | |
| "produced for a *clean* file. Switch the sweep on to look at the rest of it.") | |
| lines.append(f"Pick a batch and press **Check this batch**. One batch is one run of the model, " | |
| f"sized to fit the **{(runtime or mimo.BACKEND).upper()}** runtime's limit; run as many " | |
| f"batches as you like, one at a time.") | |
| return ((skeleton, candidates), "\n\n".join(lines), | |
| gr.update(choices=choices, value=choices[0], interactive=True)) | |
| def resolve_batch(value, n_batches: int) -> int: | |
| """ | |
| Whatever the dropdown handed back, as a batch index that exists. | |
| Because the dropdown allows custom values it can arrive as the integer index, as `None` before | |
| anything was picked, or as the label string itself. All three mean something, and none of them | |
| should be an exception in front of a user who just pressed a button. | |
| """ | |
| if isinstance(value, (int, float)): | |
| idx = int(value) | |
| else: | |
| digits = re.search(r"\d+", str(value or "")) | |
| idx = int(digits.group()) - 1 if digits else 0 # labels are 1-based, indices are not | |
| return idx if 0 <= idx < n_batches else 0 | |
| def run(path, per_batch, batch_index, want_neighbours, cover_all, runtime, | |
| progress=gr.Progress()): | |
| """Score one batch.""" | |
| if not path: | |
| return "Upload a PDF first.", [], [], "" | |
| progress(0.05, desc="reading the PDF") | |
| try: | |
| _, _, _, _, candidates = extract(path, cover_all) | |
| except Exception as e: | |
| return f"Could not read that file: `{type(e).__name__}: {e}`", [], [], "" | |
| groups = batches_of(candidates, per_batch) | |
| idx = resolve_batch(batch_index, len(groups)) | |
| batch = groups[idx] | |
| runtime = runtime or mimo.BACKEND | |
| progress(0.15, desc=f"loading MiMo ({runtime})") | |
| try: | |
| answers = mimo.judge_all([w["text"] for w in batch], backend=runtime, | |
| progress=lambda m: progress(0.4, desc=m)) | |
| except Exception as e: | |
| hint = "" | |
| if "quota" in str(e).lower(): | |
| hint = ("\n\nA free account gets about **five minutes of ZeroGPU per day**, and the " | |
| "scheduler reserves a whole run up front. The quota is per visitor, so this " | |
| "affects your account rather than the Space.") | |
| hint += ("\n\nSwitch the runtime to **cpu** — much slower, but no quota — or come back " | |
| "tomorrow." if mimo.CPU_AVAILABLE else | |
| " Fewer regions per batch costs less of it; otherwise it resets in 24 hours.") | |
| return (f"## MiMo could not run\n\n`{type(e).__name__}: {e}`{hint}\n\nRuntime selected: " | |
| f"**{runtime}**."), [], [], traceback.format_exc() | |
| offset = idx * max(1, int(per_batch)) | |
| rows, log, results = [], [], list(zip(batch, answers)) | |
| for i, (win, r) in enumerate(results): | |
| where = "head of document" if win["is_head"] else f"chars {win['start']:,}-{win['end']:,}" | |
| verdict = ("PAYLOAD" if r["pred_injected"] else "clean") + ( | |
| "" if r["parse_ok"] else " (unreadable answer)") | |
| rows.append([offset + i + 1, where, win["source"], ", ".join(win["families"]) or "-", | |
| verdict, r["pred_family"], (r["evidence"] or "-")[:160]]) | |
| log.append(f"--- region {offset + i + 1} ({where}, prompt via {r['prompt_route']}) ---\n" | |
| f"{r['raw']}") | |
| report = build_report(results, idx, groups, len(candidates), runtime) | |
| nb_rows = [] | |
| if want_neighbours and results: | |
| progress(0.95, desc="embedding and looking up the corpus") | |
| flagged = next((w for w, r in results if r["pred_injected"]), None) | |
| query = flagged or batch[0] | |
| try: | |
| neighbours.check_provenance() | |
| nb_rows = neighbours.neighbour_rows(query["text"], k=5) | |
| basis = ("the first flagged region" if flagged else | |
| "the first region in this batch (nothing was flagged)") | |
| report += (f"\n\n### Nearest files in the corpus\n\nEmbedded from **{basis}** with " | |
| f"Part A's winning configuration. Precision@5 on this index is **35.6%** " | |
| f"against a 6.8% random baseline: fewer than 2 of the 5 listed are the same " | |
| f"kind of attack. Read it as *resemblance*, not identification.") | |
| except Exception as e: | |
| report += (f"\n\n### Nearest files in the corpus\n\nUnavailable: " | |
| f"`{type(e).__name__}: {e}`") | |
| log.append(traceback.format_exc()) | |
| return report, rows, nb_rows, "\n\n".join(log) | |
| def begin(per_batch, runtime): | |
| """ | |
| Lock the button and say so, before the slow part starts. | |
| A second press while MiMo is mid-scan is worse here than in most apps: on ZeroGPU it queues a | |
| second grant against a daily quota that only affords one or two, so the cost of a stray click | |
| is the rest of the day. The button is disabled for the duration and `mimo._lock` serialises | |
| the model itself, so neither the UI nor the server can be made to run two scans at once. | |
| Stale results are cleared at the same time - leaving the previous batch's verdict on screen | |
| under a "loading" banner is how someone reads the wrong answer for the wrong file. | |
| """ | |
| runtime = runtime or mimo.BACKEND | |
| n = max(1, int(per_batch or 1)) | |
| return (gr.update(interactive=False, value="Checking…"), | |
| f"### Loading…\n\nMiMo is reading up to **{n} region(s)** on the **{runtime}** " | |
| f"runtime — roughly {fmt_eta(n, runtime)} once the model is in memory. The first scan " | |
| f"after a restart also downloads the weights, which takes several minutes longer.\n\n" | |
| f"_Leave this tab open; the report replaces this message when it is done._", | |
| [], [], "") | |
| def finish(): | |
| """Give the button back. Chained with `.then()`, so it runs even if the scan raised.""" | |
| return gr.update(interactive=True, value="Check this batch") | |
| def build_report(results, idx, groups, n_candidates, runtime) -> str: | |
| """The verdict for this batch, and an explicit account of what is still unread.""" | |
| if not results: | |
| return "Nothing was read." | |
| hits = [(w, r) for w, r in results if r["pred_injected"]] | |
| unparsed = sum(1 for _, r in results if not r["parse_ok"]) | |
| read = len(results) | |
| unread = n_candidates - read | |
| others = [i for i in range(len(groups)) if i != idx] | |
| if hits: | |
| fams = sorted({r["pred_family"] for _, r in hits if r["pred_family"] != "none"}) | |
| head = (f"## Payload found in batch {idx + 1}\n\n" | |
| f"> ### ⚠️ Instructions: do not open the file. Erase immediately!\n\n" | |
| f"MiMo flagged **{len(hits)} of the {read} region(s)** in this batch.") | |
| head += (f" It named the family as **{', '.join(fams)}** — correct 43% of the time in " | |
| f"Part B, so treat it as a suggestion." if fams | |
| else " It did not commit to a family.") | |
| # Treatment goes directly under the verdict, because it is the only part of this report | |
| # anyone acts on. One block per named family; a flagged region with no family still gets | |
| # the generic containment advice rather than silence. | |
| head += "\n\n### What to do about it\n" | |
| if fams: | |
| for family in fams: | |
| head += f"\n**{family}**\n\n{TREATMENT[family]}\n" | |
| named = "family" if len(fams) == 1 else "families" | |
| head += (f"\n_The {named} named above {'is' if len(fams) == 1 else 'are'} MiMo's " | |
| f"guess, right about 43% of the time — so this advice may be the remedy for a " | |
| f"different attack. The containment step in each first sentence holds either " | |
| f"way._") | |
| else: | |
| head += f"\n{TREATMENT_UNNAMED}\n" | |
| else: | |
| head = (f"## Nothing found in batch {idx + 1}\n\nMiMo read **{read} region(s)** in this " | |
| f"batch and flagged none of them.") | |
| caveats = [] | |
| if unread > 0: | |
| caveats.append( | |
| f"**This is 1 of {len(groups)} batches.** {unread} region(s) across " | |
| f"{len(others)} other batch(es) have not been read. Whatever this batch says, it says " | |
| f"it about {read} of the file's {n_candidates} candidate regions — nothing more.") | |
| if unparsed: | |
| sweep_bad = sum(1 for w, r in results if not r["parse_ok"] and w["source"] == "sweep") | |
| note = (f"{unparsed} answer(s) could not be parsed and count as *not injected*, exactly as " | |
| f"Part B scored them (155 of 1,100 there).") | |
| if sweep_bad: | |
| note += ( | |
| f" **{sweep_bad} of those were sweep regions**, and that is expected rather than " | |
| f"surprising: Part B only ever showed MiMo marker-centred windows or the head of a " | |
| f"document, never arbitrary mid-file content streams. Given a page of font " | |
| f"positioning operators it tends to carry on copying the input instead of " | |
| f"answering. Sweep regions buy coverage of text that would otherwise never be " | |
| f"looked at; they do not inherit Part B's accuracy, and a *clean* verdict on one " | |
| f"is close to no evidence at all.") | |
| caveats.append(note) | |
| body = head | |
| if caveats: | |
| body += "\n\n" + "\n\n".join("- " + c for c in caveats) | |
| body += (f"\n\n---\n\n**On the corpus Part B measured**, MiMo scored F1 {PART_B['f1']}, " | |
| f"precision {PART_B['precision']}, recall {PART_B['recall']} on {PART_B['n']:,} " | |
| f"documents that were 82% injected — where a detector that flags everything without " | |
| f"reading it scores F1 0.900. Read 0.945 against 0.900, not against zero.\n\n" | |
| f"_Runtime: **{runtime}**. {mimo.CAVEATS[runtime]}_") | |
| return body | |
| INTRO = f""" | |
| # PDF Injection Detector — MiMo-7B | |
| Upload a PDF. It is rendered to text with the extractor that built the project corpus, the regions | |
| carrying structural markers are ranked and cut into **batches that fit one run of the model**, and | |
| **MiMo-7B-RL** reads the batch you choose — reporting whether a payload is hidden there, and the | |
| substring that convinced it. | |
| Batching is what keeps a long document inside the runtime's limit: one batch is one run, and you | |
| decide how many runs to spend. The report always states how much of the file is still unread. | |
| **None of what it tells you is guaranteed correct — not the verdict, not the family, not the | |
| treatment.** On the corpus it was measured against, MiMo got the injected/clean call right often | |
| enough to score F1 0.945, but it named the attack family correctly only **43%** of the time — so | |
| more often than not the family shown, and therefore the remediation advice attached to it, belongs | |
| to a different attack. A clean verdict is not proof of a clean file either. Read every output as a | |
| prompt to look closer yourself, never as a decision that has already been made. | |
| This is a coursework artefact built on a synthetic corpus of 1,100 PDFs carrying harmless | |
| EICAR/AMTSO/WICAR/RANSIM test markers. **It is not a general malware scanner**, and real malware | |
| does not announce itself the way these samples do. | |
| **Why MiMo and not Gemma?** Part B's winner was Gemma-2-9B at F1 0.969, against MiMo's 0.945. But | |
| Gemma is gated behind a licence and a token, and it is 2.6× slower per window (10.95 s vs 4.18 s). | |
| On free ZeroGPU — one grant capped at 300 s, and roughly five minutes of GPU per day — that is the | |
| difference between a working demo and one that refuses strangers at the door and then runs out of | |
| quota. The cost of the swap is 0.024 F1 and family-naming dropping from 63% to 43%. | |
| Running on the **{mimo.BACKEND.upper()}** runtime | |
| ({'4-bit NF4 — Part B’s own configuration' if GPU else 'Q4_K_M GGUF via llama.cpp'}), | |
| about {mimo.SECONDS_PER_WINDOW:g}s per region. | |
| """ | |
| with gr.Blocks(title="PDF Injection Detector") as demo: | |
| gr.Markdown(INTRO) | |
| state = gr.State() | |
| with gr.Row(): | |
| # The example rail, down the left edge. Plain buttons rather than `gr.Examples`: the | |
| # built-in renders a horizontal table of filenames, and what is wanted here is one | |
| # labelled tab per attack type that loads its document on click. Buttons also sidestep | |
| # `gr.Examples`' caching, which would run a full scan of all thirteen files at startup | |
| # and spend the whole day's ZeroGPU quota before anyone opened the page. | |
| with gr.Column(scale=1, min_width=170): | |
| if EXAMPLES: | |
| gr.Markdown("### Examples\nOne document per attack type.") | |
| example_buttons = [(gr.Button(label.replace("_", " "), size="sm"), path) | |
| for label, path in EXAMPLES] | |
| else: | |
| example_buttons = [] | |
| with gr.Column(scale=2): | |
| pdf = gr.File(label="PDF", file_types=[".pdf"], type="filepath") | |
| # Offered rather than decided, because the GPU here is the scarce resource: a free | |
| # account gets ~5 minutes of ZeroGPU a day and one batch reserves most of a run. The | |
| # CPU path is ~30x slower and has no quota at all, which makes it the right answer | |
| # once the day's GPU is gone - the Space should not go dark until midnight. | |
| runtime_pick = gr.Radio( | |
| choices=mimo.BACKENDS, value=mimo.BACKEND, label="Runtime", | |
| visible=len(mimo.BACKENDS) > 1, | |
| info=("gpu = 4-bit NF4, Part B's own configuration, ~4s/region, limited by your " | |
| "daily ZeroGPU quota. cpu = Q4_K_M GGUF via llama.cpp, ~2min/region, " | |
| "unlimited.")) | |
| per_batch = gr.Slider(1, MAX_PER_BATCH, value=DEFAULT_PER_BATCH, step=1, | |
| label="Regions per batch", | |
| info=(f"How many regions one run of the model reads — it just " | |
| f"re-cuts the same list, so fewer per batch means more " | |
| f"batches. A run reserves the same GPU time whatever this " | |
| f"is set to, so lowering it inspects less of the file for " | |
| f"the same quota. Leave it at {MAX_PER_BATCH} unless you " | |
| f"want a faster single run.")) | |
| # allow_custom_value: the choices are empty until a PDF is uploaded, and without this | |
| # Gradio validates any incoming value against that empty list and rejects it - which | |
| # makes the batch un-selectable over the API even though the UI had populated it. | |
| # `resolve_batch` below is what actually decides the index, from the file itself. | |
| batch_pick = gr.Dropdown(label="Batch to check", choices=[], interactive=False, | |
| allow_custom_value=True, | |
| info="Each batch is a separate run — spend as many as you " | |
| "like.") | |
| sweep = gr.Checkbox( | |
| value=True, label="Sweep the rest of the document too", | |
| info="Off = marker regions only, which is the shape Part B measured. On = the " | |
| "batches cover the whole file, at the cost of regions MiMo often will not " | |
| "answer about.") | |
| want_nb = gr.Checkbox(value=True, label="Also show the nearest files in the corpus", | |
| info="Adds a one-off 550 MB embedding-model download.") | |
| go = gr.Button("Check this batch", variant="primary") | |
| plan = gr.Markdown("Upload a PDF to see what will be read.") | |
| with gr.Column(scale=4): | |
| with gr.Tab("Report"): | |
| report = gr.Markdown() | |
| with gr.Tab("Regions read"): | |
| window_table = gr.Dataframe(headers=WINDOW_COLUMNS, wrap=True, interactive=False) | |
| with gr.Tab("Nearest corpus files"): | |
| nb_table = gr.Dataframe(headers=neighbours.NEIGHBOUR_COLUMNS, interactive=False) | |
| with gr.Tab("What MiMo actually said"): | |
| # No `show_copy_button`: gradio 6 removed it, and this Space should survive an | |
| # sdk_version bump rather than crash at startup on a cosmetic argument. | |
| raw = gr.Textbox(lines=22, interactive=False, | |
| label="The prompt route and untouched generation per region") | |
| # Each example button just drops its path into the file component. That fires `pdf.change` | |
| # below, so an example goes through exactly the same triage as a real upload - there is no | |
| # second code path for demo files, and nothing about an example is pre-computed. | |
| for button, path in example_buttons: | |
| button.click(lambda p=str(path): p, None, pdf) | |
| inputs = [pdf, per_batch, sweep, runtime_pick] | |
| for ev in (pdf.change, per_batch.change, sweep.change, runtime_pick.change): | |
| ev(on_upload, inputs, [state, plan, batch_pick]) | |
| # Three chained steps: lock and show "Loading…", scan, unlock. `.then()` rather than | |
| # `.success()` for the last one, because the button must come back even when the scan raised - | |
| # a quota refusal that left the app permanently disabled would look like a crash. | |
| # concurrency_limit=1 is the server-side half of the same guarantee. | |
| scan = go.click(begin, [per_batch, runtime_pick], | |
| [go, report, window_table, nb_table, raw], queue=False) | |
| scan = scan.then(run, [pdf, per_batch, batch_pick, want_nb, sweep, runtime_pick], | |
| [report, window_table, nb_table, raw], concurrency_limit=1) | |
| scan.then(finish, None, go, queue=False) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=8).launch() | |