app.py

"""
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
document-shape features in `doc_features.py`, the embedding lookup in `neighbours.py` and the
family-naming model in `family_model.py` - each quoted from, or fitted in, the notebook that
measured it.
"""

import re
import traceback
from pathlib import Path

import gradio as gr

import corpus_text
import doc_features
import family_model
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 name**, and which name it is conditional on changed when the
# family-naming model was added. MiMo alone gets the family right 43% of the time, so most of the
# advice it selected was the remedy for a different attack; `family_model` chooses it now, at P@1
# 0.994 on files carrying a known signature and 0.691 on files carrying none. Neither figure is
# certainty, so the first line of every entry is still a containment step that holds whatever the
# family turns out to be, and the interface still says which model named it and how sure it was.
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 = (
    "This file was flagged but no family was named, 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}

# 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. That ceiling is what a
# batch is.
MAX_PER_BATCH = mimo.MAX_WINDOWS_GPU
DEFAULT_PER_BATCH = mimo.MAX_WINDOWS_GPU

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) -> str:
    seconds = n * mimo.SECONDS_PER_WINDOW
    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 document_answer(results):
    """
    One document-level verdict from a batch of per-region ones, for the family model.

    Part B scored MiMo once per document; this app runs it over several regions of one. The first
    flagged region is taken as the document's answer, falling back to the first region read when
    nothing was flagged - the same rule the neighbour lookup uses to pick its query, so the two
    features the family model consumes describe the same piece of text.
    """
    if not results:
        return None, None
    for window, answer in results:
        if answer["pred_injected"]:
            return window, answer
    return results[0][0], results[0][1]


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):
    """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))}")


def on_upload(path, per_batch, cover_all):
    """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)) 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("Pick a batch and press **Check this batch**. One batch is one run of the model, "
                 "sized to fit inside a single ZeroGPU grant; run as many 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_family, cover_all, progress=gr.Progress()):
    """Score one batch."""
    if not path:
        return "Upload a PDF first.", [], [], [], ""

    progress(0.05, desc="reading the PDF")
    try:
        data, skeleton, truncated, dropped, 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]

    progress(0.15, desc="loading MiMo")
    try:
        answers = mimo.judge_all([w["text"] for w in batch],
                                 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. 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}"), [], [], [], \
            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']}")

    # Stage 4. Both the family model and the neighbour table embed a region with Part A's model, so
    # they are enabled together - the 550 MB download is the cost of either one alone.
    family, family_rows, nb_rows, family_error = None, [], [], None
    query_window, doc_answer = document_answer(results)

    if want_family and results:
        progress(0.9, desc="embedding the region and naming the family")
        try:
            neighbours.check_provenance()
            stats = doc_features.describe(path, data, skeleton, truncated, dropped)
            family = family_model.predict(skeleton, doc_answer, stats, query_window["text"])
            family_rows = family_model.rows(family)
            nb_rows = neighbours.neighbour_rows(query_window["text"], k=5)
        except Exception as e:
            family_error = f"{type(e).__name__}: {e}"
            log.append(traceback.format_exc())

    report = build_report(results, idx, groups, len(candidates), family, family_error,
                          doc_answer, bool(want_family))

    if nb_rows:
        flagged = doc_answer and doc_answer["pred_injected"]
        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 Part A's "
                   f"winning configuration — the same vector the family model's neighbour-vote "
                   f"feature is built from. Precision@5 on this index is **35.6%** against a 6.8% "
                   f"random baseline: fewer than 2 of the 5 listed are the same kind of attack. "
                   f"Read it as *resemblance*, not identification.")

    return report, rows, family_rows, nb_rows, "\n\n".join(log)


def begin(per_batch):
    """
    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.
    """
    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)** — roughly {fmt_eta(n)} once "
            f"the model is in memory. The first scan after a restart also downloads the weights, "
            f"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 family_section(family, family_error, mimo_fams, want_family, mimo_flagged) -> tuple:
    """
    The family model's shortlist, and which family the treatment below should be written for.

    Returns `(markdown, [families to treat], source)`. The model is allowed to disagree with MiMo
    and usually should - naming the family is the job it was fitted for and the job MiMo is worst
    at - but a disagreement is printed rather than resolved silently.
    """
    if family is None:
        if not want_family:
            note = ("\n\n### Attack family\n\nThe family-naming model is switched off, so the "
                    "family below is **MiMo's own guess — right about 43% of the time**. Switch "
                    "*Name the attack family* on for the model's ranked answer instead.")
        else:
            note = (f"\n\n### Attack family\n\nThe family-naming model could not run "
                    f"(`{family_error}`), so the family below is **MiMo's own guess — right about "
                    f"43% of the time**. Everything else in this report is unaffected.")
        return note, mimo_fams, "MiMo"

    ranked = family["ranked"]
    top, confidence = ranked[0]
    shortlist = " · ".join(f"**{name}** {prob:.0%}" if i == 0 else f"{name} {prob:.0%}"
                           for i, (name, prob) in enumerate(ranked))

    lines = ["\n\n### Attack family", ""]
    if top == "none":
        stance = ("**against MiMo's flag**" if mimo_flagged else
                  "which **agrees with MiMo**, from a different four signals")
        lines.append(
            f"The family-naming model puts **{confidence:.0%}** of its confidence on this file "
            f"carrying **no injection at all** — {stance}. Ranked: {shortlist}.")
    else:
        lines.append(f"The family-naming model names this **{top}** at **{confidence:.0%}** "
                     f"confidence. Ranked: {shortlist}.")

    lines.append(
        f"\nIt fuses four signals — the structural signatures found in the file, MiMo's own verdict "
        f"and guess, seventeen document-shape measurements, and a weighted vote of the 20 nearest "
        f"corpus files in Part A's embedding index. Out of fold it reaches **P@1 "
        f"{family['p_at_1']:.3f}** in this regime ({family['regime']}).")

    disagreement = [f for f in mimo_fams if f != top]
    if disagreement and top != "none":
        lines.append(
            f"\n_MiMo said **{', '.join(disagreement)}**; the model says **{top}**. The model is "
            f"the one the advice below follows, because naming the family is the task it was "
            f"fitted for and the one MiMo is measurably worst at (43%)._")

    return "\n".join(lines), ([top] if top != "none" else []), "the family-naming model"


def build_report(results, idx, groups, n_candidates, family, family_error, doc_answer,
                 want_family) -> 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]

    mimo_fams = sorted({r["pred_family"] for _, r in hits if r["pred_family"] != "none"})
    fam_md, treat_fams, source = family_section(family, family_error, mimo_fams, want_family,
                                                bool(hits))

    if hits:
        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 — the verdict "
                f"it scores F1 0.945 on.")
        head += fam_md

        # Treatment goes directly under the family, because it is the only part of this report
        # anyone acts on. A flagged file with no name still gets the generic containment advice
        # rather than silence.
        head += "\n\n### What to do about it\n"
        if treat_fams:
            for name in treat_fams:
                head += f"\n**{name}**\n\n{TREATMENT[name]}\n"
            certainty = (f"{family['confidence']:.0%} confident, P@1 {family['p_at_1']:.3f} in "
                         f"this regime" if family else "right about 43% of the time")
            head += (f"\n_Chosen by **{source}** ({certainty}) — so this advice can still be the "
                     f"remedy for a different attack. The containment step in the first sentence "
                     f"holds either way._")
        elif family is not None and family["top"] == "none":
            head += (f"\n{TREATMENT_UNNAMED}\n\n_MiMo flagged a region; the family-naming model "
                     f"disagrees and puts the file in the **clean** class. Two models disagreeing "
                     f"is a reason to look at the file yourself, not a reason to trust either._")
        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.")
        # A clean verdict is the harder half of the claim, so the model's opinion is worth as much
        # here as it is over a hit - it either corroborates MiMo or contradicts it.
        if family is not None:
            head += fam_md

    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. The "
             f"family-naming model's own baseline is the same corpus's 18.2% clean share: always "
             f"answering *none* scores P@1 0.182.\n\n"
             f"_{mimo.RUNTIME_NOTE}_")
    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**,
**MiMo-7B-RL** reads the batch you choose and says whether a payload is hidden there, and a
**family-naming model** then decides *which kind* of attack it is.

Batching is what keeps a long document inside one ZeroGPU grant: 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.

**Two models, because they are good at different questions.** MiMo answers *is there a payload
here?* — F1 0.945 on the corpus it was measured against — and quotes the substring that convinced
it. It is poor at *which kind?*, naming the family correctly only **43%** of the time, which
mattered because the remediation advice is chosen by family. So the family is named instead by a
gradient-boosted model fitted on four cheap signals (structural signatures, MiMo's own opinion,
seventeen document-shape measurements, and a vote of the nearest corpus files): **P@1 0.994** on
files carrying one of the twelve known signatures, **0.691** on files carrying none. The report
says which of those two regimes your file fell into rather than quoting one number for both.

**None of it is guaranteed correct** — not the verdict, not the family, not the treatment. A clean
verdict is not proof of a clean file. 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 is no longer MiMo's job anyway.

Running MiMo in 4-bit NF4 — Part B's own configuration — at about {mimo.SECONDS_PER_WINDOW:g}s per
region. The family-naming model is 915 KB and runs on the CPU in under a second.
"""


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")
            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.")
            # One control for both, because they are the same download: the family model's
            # neighbour-vote feature and the neighbour table are two readings of the same query
            # vector, so paying 550 MB for one and not the other would be strange.
            want_family = gr.Checkbox(
                value=True, label="Name the attack family (and show nearest corpus files)",
                info=("Runs the family-naming model and Part A's lookup. Adds a one-off 550 MB "
                      "embedding-model download. Off, the family shown is MiMo's own guess."),
                interactive=family_model.AVAILABLE)
            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("Family shortlist"):
                family_table = gr.Dataframe(headers=family_model.ROW_COLUMNS, 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]
    for ev in (pdf.change, per_batch.change, sweep.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],
                    [go, report, window_table, family_table, nb_table, raw], queue=False)
    scan = scan.then(run, [pdf, per_batch, batch_pick, want_family, sweep],
                     [report, window_table, family_table, nb_table, raw], concurrency_limit=1)
    scan.then(finish, None, go, queue=False)

if __name__ == "__main__":
    demo.queue(max_size=8).launch()