MiMo-7B PDF injection detector: CPU/GGUF runtime, verbatim Part A+B logic

README.md

@@ -1,13 +1,108 @@
 ---
-title: Mimo Injection Detector
-emoji: 🏢
+title: PDF Injection Detector (MiMo-7B)
+emoji: 🔍
 colorFrom: yellow
-colorTo: yellow
+colorTo: gray
 sdk: gradio
-sdk_version: 6.22.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
+license: mit
+short_description: MiMo-7B finds payloads hidden inside PDF files
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PDF Injection Detector — MiMo-7B on free CPU
+
+Upload a PDF. It is rendered to text with the extractor that built the project corpus, the regions
+carrying structural markers are ranked, and **MiMo-7B-RL** reads the most promising ones and says
+whether a payload is hidden there — naming the family and quoting the substring that convinced it.
+
+It 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.**
+
+## The three repos this is built on
+
+| Repo | What this Space takes from it |
+|---|---|
+| [Generated_Injected_PDFs_HARMLESS](https://huggingface.co/datasets/Cyber-security-final-project/Generated_Injected_PDFs_HARMLESS) | The 12 injection families and their structural signatures — the definition of what an attack looks like |
+| [HARMLESS_Synthetic_Injected_PDFs_EDA](https://huggingface.co/datasets/Cyber-security-final-project/HARMLESS_Synthetic_Injected_PDFs_EDA) | `build_skeleton`, `mask_leaks` and `payload_window` — how a PDF becomes the text a model reads |
+| [Evaluation_of_OpenSource_Models…](https://huggingface.co/datasets/Cyber-security-final-project/Evaluation_of_OpenSource_Models_for_PDF_Injection_Recognition) | Part A's embedding index and winning configuration; Part B's prompt, prefill and parser |
+
+The prompt, the MiMo prefill, the brace-counting parser and the whole text-extraction path are
+quoted **verbatim** from those notebooks. That is the correctness argument for quoting Part B's
+scores here at all: change how the text is extracted or how the question is asked, and the
+published numbers stop describing this program.
+
+## Fitting inside a free Space
+
+A free Space is 2 vCPUs and 16 GB of RAM, with no GPU. Three consequences, all of them visible in
+the interface rather than hidden:
+
+**MiMo runs on the CPU as a GGUF.** Part B ran the BF16 checkpoint quantised to 4-bit NF4 by
+`bitsandbytes`, which requires CUDA. Here the same base model runs as
+[`quantflex/MiMo-7B-RL-nomtp-Q4_K_M.gguf`](https://huggingface.co/quantflex/MiMo-7B-RL-nomtp-GGUF)
+(4.7 GB) through `llama.cpp`. That build has MiMo's multi-token-prediction layers removed, because
+`llama.cpp` cannot load them — MTP is a speculative-decoding accelerator that the ordinary forward
+pass does not use, so greedy output should be unaffected, but it is a real difference and it is
+stated rather than buried. **Read F1 0.945 as the figure for the configuration Part B measured,
+not as a measurement of this Space.**
+
+**It triages instead of scanning everything.** MiMo reads a 3,000-character window in roughly two
+minutes on 2 vCPUs, and a real PDF has dozens of windows. So the same marker alternation that
+located the payload in the corpus is run over the whole skeleton, every hit becomes a candidate
+window with the identical ±1,500-character shape, overlapping ones are merged, and the most
+marker-dense go to the model first. **The ranking decides reading order, never the verdict.** The
+report always says how many marked regions were left unread, so "clean" never overstates itself.
+
+A file with no marker anywhere yields exactly one candidate — the head of the document — which is
+byte-identical to what the corpus builder produced for a *clean* file. The triage only knows the
+twelve families this project generated, so a payload shaped like none of them is triaged as if the
+file were clean and MiMo sees the head of the file. That is a real hole, and the interface says so.
+
+**Nothing is downloaded until it is needed.** The page comes up first; the 4.7 GB GGUF, the
+550 MB embedding model and the 3 MB index arrive on the first scan and are cached after that.
+
+## The numbers
+
+On the 1,100-document corpus Part B measured, MiMo-7B-RL scored:
+
+| | |
+|---|---|
+| F1 | **0.945** |
+| precision | 0.988 |
+| recall | 0.906 |
+| names the family correctly | 43.3% of files it caught |
+| false alarms | 10 of 200 clean files (5%) |
+| unparsable answers | 155 of 1,100 |
+
+**A detector that calls every file malicious scores F1 0.900 on this corpus**, because 82% of it is
+injected. Read 0.945 against 0.900, not against zero — it is a 5% relative improvement on doing no
+work at all. Gemma-2-9B scored 0.969 and is the actual Part B winner; MiMo is used here because it
+is ungated, needs no token, and is 2.6× faster, which on a CPU is the difference between usable
+and not.
+
+Two limits worth stating plainly:
+
+- **The family is a suggestion, not a verdict** — right 43% of the time. The nearest known corpus
+  files are shown beside it so the two can disagree in public.
+- **The nearest-file lookup is weak on purpose to report.** Part A's winning embedder reaches
+  precision@5 of 35.6% against a 6.8% random baseline: fewer than 2 of the 5 files returned are the
+  same kind of attack. Far better than chance, and not good. It is labelled *resemblance*, never
+  *identification*.
+
+## The files
+
+| File | What it does |
+|---|---|
+| `app.py` | The Gradio interface and the document-level report. No detection logic. |
+| `corpus_text.py` | PDF bytes → skeleton → candidate windows. Everything above `Triage` is verbatim from the EDA notebook. |
+| `mimo.py` | The prompt, the prefill, the parser (verbatim from Part B) and the llama.cpp runtime. |
+| `neighbours.py` | Part A's embedding index and the nearest-neighbour lookup, with a provenance assertion. |
+
+`neighbours.check_provenance()` asserts the embedder repo, prefix, dimension, normalisation and
+input column against Part A's own `part_a_results.json` before any lookup runs — a query embedded
+with the wrong model lands in a different space and returns meaningless neighbours silently, with
+no error anywhere.
+
+## Secrets
+
+None. Every model used here is ungated.

app.py

@@ -0,0 +1,225 @@
+"""
+PDF Injection Detector - MiMo-7B on a free CPU Space.
+
+Upload a PDF, and this reads it the way the corpus was read, picks the parts most worth looking at,
+and asks MiMo-7B-RL whether a payload is hidden in them.
+
+`app.py` holds the interface 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 of those is quoted from the notebook that measured it.
+"""
+
+import traceback
+
+import gradio as gr
+
+import corpus_text
+import mimo
+import neighbours
+
+# Measured in Part B on a T4. Kept here to be shown next to the app's own timings, because the two
+# are wildly different and pretending otherwise would be the misleading thing.
+PART_B = {"f1": 0.945, "precision": 0.988, "recall": 0.906, "family_acc": 0.433,
+          "false_alarm_rate": 0.05, "unparsable": 155, "n": 1100}
+
+# A free Space is 2 vCPUs. A 7B model in Q4_K_M reads a 3,000-character window in roughly this
+# long there - measured loosely, and shown as a range because it is a range.
+SECONDS_PER_WINDOW = 130
+
+MAX_WINDOWS = 6
+
+WINDOW_COLUMNS = ["#", "where in the skeleton", "markers", "verdict", "family", "evidence"]
+
+
+def fmt_eta(n_windows: int) -> str:
+    lo = n_windows * SECONDS_PER_WINDOW * 0.6 / 60
+    hi = n_windows * SECONDS_PER_WINDOW * 1.5 / 60
+    return f"{lo:.0f}-{hi:.0f} minutes"
+
+
+def triage(path):
+    """Extract and rank, without loading any model. Runs in under a second, so it runs on upload."""
+    if not path:
+        return None, "Upload a PDF to see what will be read."
+
+    data = open(path, "rb").read()
+    skeleton, truncated, dropped = corpus_text.build_skeleton(data)
+    candidates = corpus_text.candidate_windows(skeleton)
+    families = corpus_text.detect_markers(data)
+
+    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 "."),
+    ]
+    if candidates[0]["is_head"]:
+        lines.append(
+            "**No structural marker found anywhere in the file.** The head of the document will be "
+            "read instead — which is exactly what a *clean* file in the corpus looks like. Note "
+            "that the triage only recognises the twelve families this project generated, so a "
+            "payload shaped like none of them lands here too.")
+    else:
+        lines.append(
+            f"**{len(candidates)} candidate region(s)** carry a marker. "
+            f"Structural signatures present in the raw file: `{'`, `'.join(families) or 'none'}`.")
+        lines.append(
+            "_Those signatures decide **reading order only**. The verdict is MiMo's alone — the "
+            "regex has been wrong about a file before and will be again._")
+    return (skeleton, candidates), "\n\n".join(lines)
+
+
+def on_upload(path, n_windows):
+    state, summary = triage(path)
+    if state is None:
+        return state, summary
+    n = min(int(n_windows), len(state[1]))
+    return state, summary + (
+        f"\n\nMiMo will read **{n} of {len(state[1])}** region(s): about **{fmt_eta(n)}** on this "
+        f"free CPU Space, plus a one-off download the first time anyone uses it.")
+
+
+def run(path, n_windows, want_neighbours, progress=gr.Progress()):
+    """Score the top windows and build the report. Yields so the page updates as answers arrive."""
+    if not path:
+        yield "Upload a PDF first.", [], [], ""
+        return
+
+    progress(0, desc="reading the PDF")
+    state, _ = triage(path)
+    skeleton, candidates = state
+    chosen = candidates[:int(n_windows)]
+
+    rows, results, log = [], [], []
+    for i, win in enumerate(chosen):
+        progress(i / len(chosen), desc=f"MiMo reading region {i + 1} of {len(chosen)}")
+        try:
+            r = mimo.judge(win["text"], progress=lambda m: progress(i / len(chosen), desc=m))
+        except Exception as e:
+            log.append(f"region {i + 1} failed: {type(e).__name__}: {e}")
+            continue
+
+        results.append((win, r))
+        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([i + 1, where, ", ".join(win["families"]) or "-", verdict,
+                     r["pred_family"], (r["evidence"] or "-")[:160]])
+        log.append(f"--- region {i + 1} ({where}, prompt via {r['prompt_route']}) ---\n{r['raw']}")
+        yield building(chosen, candidates, results), rows, [], "\n\n".join(log)
+
+    report = building(chosen, candidates, results, done=True)
+
+    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 chosen[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 read (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 below are the same "
+                       f"kind of attack. Read it as *resemblance*, not as 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())
+
+    yield report, rows, nb_rows, "\n\n".join(log)
+
+
+def building(chosen, candidates, results, done=False) -> str:
+    """The document-level report. One verdict from several window verdicts, and the arithmetic."""
+    if not results:
+        return "Reading…" if not done else "Nothing was read."
+
+    hits = [(w, r) for w, r in results if r["pred_injected"]]
+    unread = len(candidates) - len(chosen)
+    unparsed = sum(1 for _, r in results if not r["parse_ok"])
+
+    if hits:
+        fams = sorted({r["pred_family"] for _, r in hits if r["pred_family"] != "none"})
+        head = (f"## Payload found\n\nMiMo flagged **{len(hits)} of the {len(results)} region(s) "
+                f"it read**.")
+        if fams:
+            head += (f" It named the family as **{', '.join(fams)}** — correct 43% of the time in "
+                     f"Part B, so treat it as a suggestion.")
+        else:
+            head += " It did not commit to a family."
+    else:
+        head = (f"## Nothing found\n\nMiMo read **{len(results)} region(s)** and flagged none of "
+                f"them.")
+
+    caveats = []
+    if unread > 0:
+        caveats.append(f"**{unread} marked region(s) were not read** — the free CPU tier reads a "
+                       f"few regions, not all of them. 'Clean' here means 'clean in what was "
+                       f"read'.")
+    if unparsed:
+        caveats.append(f"{unparsed} answer(s) could not be parsed and count as *not injected*, "
+                       f"exactly as Part B scored them (155 of 1,100 there).")
+    if not done:
+        caveats.append("_Still reading…_")
+
+    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 {100 - 200 / 11:.0f}% injected — where a detector that flags "
+             f"everything without reading it scores F1 0.900. Read 0.945 against 0.900, not "
+             f"against zero.\n\n_{mimo.MODEL_CAVEAT}_")
+    return body
+
+
+INTRO = """
+# PDF Injection Detector — MiMo-7B
+
+Upload a PDF. It is rendered to text with the extractor that built the corpus, the regions carrying
+structural markers are ranked, and **MiMo-7B-RL** reads the most promising ones and says whether a
+payload is hidden there — with the substring that convinced it.
+
+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.
+
+It runs on a **free CPU Space**, so MiMo reads roughly one region every two minutes. That is the
+whole reason it triages instead of scanning everything.
+"""
+
+
+with gr.Blocks(title="PDF Injection Detector") as demo:
+    gr.Markdown(INTRO)
+    state = gr.State()
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            pdf = gr.File(label="PDF", file_types=[".pdf"], type="filepath")
+            n_windows = gr.Slider(1, MAX_WINDOWS, value=3, step=1,
+                                  label="Regions for MiMo to read",
+                                  info="Each one costs a couple of minutes on free CPU.")
+            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 PDF", variant="primary")
+            plan = gr.Markdown("Upload a PDF to see what will be read.")
+
+        with gr.Column(scale=2):
+            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"):
+                raw = gr.Textbox(lines=22, show_copy_button=True, interactive=False,
+                                 info="The prompt route and the untouched generation per region.")
+
+    for ev in (pdf.change, n_windows.change):
+        ev(on_upload, [pdf, n_windows], [state, plan])
+    go.click(run, [pdf, n_windows, want_nb], [report, window_table, nb_table, raw])
+
+if __name__ == "__main__":
+    demo.queue(max_size=8).launch()

corpus_text.py

@@ -0,0 +1,226 @@
+"""
+Stage 1 - turning an uploaded PDF into the exact kind of string the corpus was built from.
+
+Everything in the first block is **lifted verbatim from the EDA notebook that built the corpus**
+(`HARMLESS_Synthetic_Injected_PDFs_EDA/Final_project_V7_EDA.ipynb`, cells 88-90 and 107). That is
+not tidiness, it is the correctness argument for this whole Space: Part A's embedding index and
+Part B's F1 of 0.945 were both measured on `payload_window` strings produced by exactly this code.
+Extract the text even slightly differently and those published numbers stop describing this app.
+
+Do not "improve" anything above the `Triage` heading.
+"""
+
+import re
+import zlib
+
+import numpy as np
+
+# ---------------------------------------------------------------------------------------------
+# Verbatim from the corpus build - EDA cell 88
+# ---------------------------------------------------------------------------------------------
+
+STREAM_RE = re.compile(rb"(stream\r?\n)(.*?)(endstream)", re.S)
+
+INJECTION_MARKERS = {
+    "javascript_injection":     [rb"/S\s*/JavaScript", rb"/JS\s*\("],
+    "cross_site_scripting":     [rb"<script>", rb"fetch\("],
+    "ssrf":                     [rb"169\.254\.169\.254", rb"/latest/meta-data"],
+    "object_action_injection":  [rb"/S\s*/Launch", rb"cmd\.exe"],
+    "llm_prompt_injection":     [rb"IGNORE PREVIOUS INSTRUCTIONS", rb"LLM-INJECT"],
+    "shellcode_embedded_exe":   [rb"application#2Fx-msdownload", rb"MZ.{0,20}\\x90\\x90"],
+    "polyglot_file":            [rb"PK\\x03\\x04", rb"POLYGLOT ZIP\+PDF"],
+    "dde_template_injection":   [rb"remote\.dotm\?dde=", rb"=cmd\|"],
+    "steganographic_payload":   [rb"LSB-STEGO"],
+    "ransomware_simulation":    [rb"RANSIM TEST"],
+    "xfa_acroform_injection":   [rb"/Type\s*/XFA", rb"xfa:data"],
+    "uri_redirect_phishing":    [rb"phishing\.\w+\.test"],
+}
+FRAMEWORK_MARKERS = ["AMTSO", "WICAR", "OWASP", "AtomicRedTeam", "Metasploit",
+                     "Glasswall", "mindcrypt", "RanSim", "RANSIM", "custom"]
+BASE_PAYLOAD_MARKERS = ["EICAR-STANDARD-ANTIVIRUS-TEST-FILE", "AMTSO-POTENTIALLY-UNWANTED-TEST-FILE",
+                        "WICAR-BROWSER-TEST-PAYLOAD", "RANSIM-KNOWBE4-ENCRYPTION-SIMULATION"]
+
+# Brand / test-file identifiers a model could string-match on without reasoning about PDF
+# structure. The structural markers above are deliberately NOT masked: those ARE the attack, not a
+# giveaway label. Masking them would delete the thing the model is supposed to find.
+LEAK_STRINGS = sorted(
+    set(FRAMEWORK_MARKERS) | set(BASE_PAYLOAD_MARKERS)
+    | {"EICAR", "RANSIM", "eicar-standard-antivirus-test-file",
+       r"X5O!P%@AP[4\PZX54(P^)7CC)7}$"},
+    key=len, reverse=True)            # longest first: "EICAR-STANDARD-..." masks before "EICAR"
+LEAK_RE = re.compile("|".join(re.escape(s) for s in LEAK_STRINGS), re.IGNORECASE)
+
+# ---------------------------------------------------------------------------------------------
+# Verbatim from the corpus build - EDA cell 89
+# ---------------------------------------------------------------------------------------------
+
+STREAM_BODY_CAP      = 4_096      # chars kept from any single stream body
+SKELETON_CHAR_BUDGET = 120_000    # ~30k tokens
+HEAD_SHARE           = 0.45       # of the budget; the rest is the tail window
+SCAN_WINDOW          = 8_000_000  # bytes scanned from each end for markers
+
+CLEAN_RE  = re.compile(r"[^\x20-\x7e\n]")
+SPACES_RE = re.compile(r"[ ]{4,}")
+
+
+def _printable_frac(chunk: bytes, sample: int = 200_000) -> float:
+    """Fraction of bytes that are ordinary printable ASCII."""
+    if not chunk:
+        return 1.0
+    arr = np.frombuffer(chunk[:sample], dtype=np.uint8)
+    ok = ((arr >= 32) & (arr < 127)) | (arr == 9) | (arr == 10) | (arr == 13)
+    return float(ok.mean())
+
+
+def build_skeleton(data: bytes):
+    """Render a PDF as payload-preserving text. Returns (skeleton, was_truncated, n_binary_dropped)."""
+    dropped = 0
+
+    def replace(match):
+        nonlocal dropped
+        opener, body, closer = match.group(1), match.group(2), match.group(3)
+        try:                                            # most streams are FlateDecode
+            inflated = zlib.decompress(body)
+            if _printable_frac(inflated) > 0.6:
+                return opener + inflated[:STREAM_BODY_CAP] + b"\n" + closer
+        except zlib.error:
+            pass
+        if _printable_frac(body) > 0.6:                 # already plain text
+            return opener + body[:STREAM_BODY_CAP] + closer
+        dropped += 1                                    # genuinely binary (an image)
+        return opener + b"<<BINARY %d bytes>>" % len(body) + closer
+
+    text = STREAM_RE.sub(replace, data).decode("latin-1")
+
+    # Truncate FIRST, then clean: no point running two character-class substitutions over 76
+    # million characters only to discard 99.8% of the result.
+    truncated = len(text) > SKELETON_CHAR_BUDGET
+    if truncated:
+        # Head AND tail. The two insertion strategies put payloads at opposite ends of the file,
+        # so a plain head truncation would lose most of them.
+        head_n = int(SKELETON_CHAR_BUDGET * HEAD_SHARE)
+        tail_n = SKELETON_CHAR_BUDGET - head_n
+        elided = len(text) - SKELETON_CHAR_BUDGET
+        text = (text[:head_n * 2]
+                + f"\n<<... {elided} characters elided ...>>\n"
+                + text[-tail_n * 2:])
+
+    text = CLEAN_RE.sub(" ", text)                      # drop control/binary residue
+    text = SPACES_RE.sub("   ", text)
+
+    if len(text) > SKELETON_CHAR_BUDGET:                # enforce the budget after cleanup
+        head_n = int(SKELETON_CHAR_BUDGET * HEAD_SHARE)
+        text = text[:head_n] + text[-(SKELETON_CHAR_BUDGET - head_n):]
+        truncated = True
+    return text, truncated, dropped
+
+
+def mask_leaks(skeleton: str) -> str:
+    """Blank the brand identifiers, keep the structural shape of the injection intact."""
+    return LEAK_RE.sub(lambda m: "X" * len(m.group(0)), skeleton)
+
+
+def scan_window(data: bytes, window: int = SCAN_WINDOW) -> bytes:
+    """Head and tail of a file, for the regex marker scans."""
+    if len(data) <= 2 * window:
+        return data
+    return data[:window] + data[-window:]
+
+
+def detect_markers(blob) -> list:
+    """Which injection types are structurally present. Works on raw bytes or on a skeleton."""
+    data = blob.encode("latin-1", errors="replace") if isinstance(blob, str) else blob
+    data = scan_window(data)
+    return [name for name, patterns in INJECTION_MARKERS.items()
+            if any(re.search(p, data, re.IGNORECASE) for p in patterns)]
+
+
+# One alternation instead of 26 separate scans - EDA cell 107. The earliest match of the union is
+# by definition the earliest match of any individual pattern.
+ANY_MARKER_RE = re.compile(
+    b"|".join([p for ps in INJECTION_MARKERS.values() for p in ps]
+              + [re.escape(s).encode() for s in FRAMEWORK_MARKERS + BASE_PAYLOAD_MARKERS]),
+    re.IGNORECASE)
+
+WINDOW = 1_500        # characters kept either side of the payload - EDA cell 107
+
+
+def payload_window(text: str, half: int = WINDOW) -> str:
+    """
+    The neighbourhood of the injection, exactly as the corpus column of this name was built.
+
+    Clean files have no marker, so they fall back to the head of the document - which keeps them
+    comparable in length rather than empty. Kept here unchanged because it is the single-window
+    case, and because `candidate_windows` below must agree with it character for character.
+    """
+    m = ANY_MARKER_RE.search(text.encode("latin-1", errors="replace"))
+    if not m:
+        return text[:2 * half]
+    return text[max(0, m.start() - half): m.start() + half]
+
+
+# ---------------------------------------------------------------------------------------------
+# Triage - new here, and the one place this Space departs from the notebooks
+# ---------------------------------------------------------------------------------------------
+#
+# The notebooks scored one window per document, because they already knew where the payload was.
+# An uploaded file offers no such promise: a payload can sit anywhere, and on the free CPU tier
+# MiMo reads roughly one window every couple of minutes, so "score every window" is not on offer.
+#
+# So the same marker alternation that located the corpus payload is run over the *whole* skeleton
+# instead of stopping at the first hit. Every match becomes a candidate window with the identical
+# +/-1,500-character shape, they are merged where they overlap, and the most marker-dense ones go
+# to the model first. A file with no marker anywhere yields exactly one candidate - the head of
+# the document - which is byte-identical to what `payload_window` returns for a clean corpus file.
+#
+# What this is NOT: a detector. The ranking decides reading order, never the verdict. It also only
+# knows the twelve families' signatures, so a payload shaped like none of them is triaged as if it
+# were clean and the model sees the head of the file. That limit is stated in the UI, not buried.
+
+MAX_MATCH_SCAN = 4_000            # matches considered; a pathological file will not run forever
+
+
+def candidate_windows(skeleton: str, half: int = WINDOW) -> list:
+    """
+    Every +/-`half` neighbourhood around a marker in `skeleton`, merged and ranked.
+
+    Returns dicts with `start`, `end`, `text`, `n_markers`, `families` and `is_head`, most
+    marker-dense first. Never empty: with no markers at all it returns the head window.
+    """
+    blob = skeleton.encode("latin-1", errors="replace")
+
+    spans = []
+    for i, m in enumerate(ANY_MARKER_RE.finditer(blob)):
+        if i >= MAX_MATCH_SCAN:
+            break
+        spans.append((max(0, m.start() - half), m.start() + half, m.start()))
+
+    if not spans:
+        head = skeleton[:2 * half]
+        return [{"start": 0, "end": len(head), "text": head, "n_markers": 0,
+                 "families": [], "is_head": True}]
+
+    # Merge overlaps so two markers 200 characters apart are read once, not twice.
+    merged = []
+    for start, end, hit in spans:
+        if merged and start <= merged[-1]["end"]:
+            merged[-1]["end"] = max(merged[-1]["end"], end)
+            merged[-1]["hits"].append(hit)
+        else:
+            merged.append({"start": start, "end": end, "hits": [hit]})
+
+    out = []
+    for span in merged:
+        # A merged span can grow past one window. The model reads at most 2*half characters, so
+        # centre the slice on the first marker in the span rather than sending a longer string
+        # than any corpus row ever carried.
+        first = span["hits"][0]
+        start = max(0, first - half)
+        text = skeleton[start:start + 2 * half]
+        out.append({"start": start, "end": start + len(text), "text": text,
+                    "n_markers": len(span["hits"]),
+                    "families": detect_markers(text), "is_head": False})
+
+    # Density first, then position: an early hit is where the head-insertion strategy puts things.
+    out.sort(key=lambda w: (-len(w["families"]), -w["n_markers"], w["start"]))
+    return out

mimo.py

@@ -0,0 +1,240 @@
+"""
+Stage 2 - asking MiMo.
+
+The prompt, the prefill and the parser are **lifted verbatim from Part B**
+(`Evaluation_of_OpenSource_Models_for_PDF_Injection_Recognition/Final_Project_Evaluation_V1.ipynb`,
+cell 48). Same system message, same closed family list in the same order, same MiMo prefill, same
+brace-counting parser, same greedy decoding, same 200-token cap. A reworded prompt or a widened
+parser is a different experiment, and Part B's measured F1 of 0.945 would no longer describe it.
+
+The one thing that is genuinely different from Part B is **how the weights are executed**. Part B
+ran the BF16 checkpoint quantised to 4-bit NF4 through `bitsandbytes` on a T4. A free Hugging Face
+Space has no GPU and `bitsandbytes` requires CUDA, so this runs the same weights as a Q4_K_M GGUF
+through `llama.cpp` on the CPU. Both are 4-bit quantisations of `XiaomiMiMo/MiMo-7B-RL`, and they
+are not the same 4-bit quantisation - see MODEL_CAVEAT, which the UI prints rather than hides.
+"""
+
+import json
+import os
+import re
+import threading
+
+FAMILIES = ["cross_site_scripting", "dde_template_injection", "javascript_injection",
+            "llm_prompt_injection", "object_action_injection", "polyglot_file",
+            "ransomware_simulation", "shellcode_embedded_exe", "ssrf",
+            "steganographic_payload", "uri_redirect_phishing", "xfa_acroform_injection"]
+
+# ---------------------------------------------------------------------------------------------
+# Verbatim from Part B, cell 48
+# ---------------------------------------------------------------------------------------------
+
+SYSTEM = (
+    "You are a PDF security analyst. You are given the raw extracted text of a PDF file - object "
+    "definitions, stream contents and metadata, exactly as they appear in the file. Some of these "
+    "files have had a malicious payload injected into them; most, but not all, have. Your job is to "
+    "say which, and to point at your evidence.\n\n"
+    "Answer with a single JSON object and nothing else:\n"
+    '{"injected": true or false, '
+    '"injection_type": one of ' + json.dumps(FAMILIES) + ' or "none", '
+    '"evidence": the exact substring from the input that convinced you, at most 200 characters, '
+    'or "" if none, '
+    '"reasoning": one short sentence}\n\n'
+    "If the file looks clean, answer injected=false and injection_type=\"none\". Do not guess a "
+    "family when you do not believe there is an injection."
+)
+
+MAX_CHARS = 3000        # payload_window is already capped at this; belt and braces
+MAX_NEW = 200
+
+
+def build_messages(text: str):
+    """The chat turns for one window."""
+    return [{"role": "system", "content": SYSTEM},
+            {"role": "user", "content": "PDF extract:\n\n```\n" + text[:MAX_CHARS] + "\n```"}]
+
+
+# Text appended to the assistant turn, so the model resumes from it instead of starting free.
+# MiMo is reasoning-trained, opens every answer with `<think>`, and at MAX_NEW = 200 the budget is
+# gone before the block closes - not one of its 1,100 Part B answers contained a closing
+# `</think>`, so it never reached the JSON. An empty, already-closed block says the deliberation is
+# finished before it begins, and the opening brace puts it inside the answer.
+PREFILL = '<think>\n\n</think>\n\n{"injected":'
+
+
+def scan_objects(raw: str):
+    r"""
+    Every balanced {...} in the text, counting braces and skipping anything inside a string literal.
+
+    A regex cannot do this - matching balanced delimiters is outside what regular expressions can
+    express - and the naive `\{[^{}]*\}` this replaced was actively harmful: every injected file in
+    this corpus carries an EICAR-style marker containing a `}`, so the moment a model quoted its
+    evidence the match was truncated mid-string and the verdict was discarded. The bug fired
+    exactly when the model was RIGHT.
+    """
+    objs, depth, start, in_str, esc = [], 0, None, False, False
+    for i, ch in enumerate(raw or ""):
+        if in_str:
+            if esc:
+                esc = False
+            elif ch == "\\":
+                esc = True
+            elif ch == '"':
+                in_str = False
+            continue
+        if ch == '"':
+            in_str = True
+        elif ch == "{":
+            if depth == 0:
+                start = i
+            depth += 1
+        elif ch == "}":
+            depth -= 1
+            if depth == 0 and start is not None:
+                objs.append(raw[start:i + 1])
+                start = None
+            depth = max(depth, 0)
+    return objs
+
+
+def parse_response(raw: str) -> dict:
+    """
+    Pull the verdict out of whatever the model said.
+
+    The LAST balanced object is taken, not the first: reasoning-trained models restate the schema
+    while thinking and emit the real answer at the end. If nothing parses as JSON the two fields
+    that matter are lifted out individually rather than thrown away. Only a response with no
+    recoverable verdict counts as parse_ok=False, and that reads as 'not injected': a detector that
+    cannot make itself understood has caught nothing.
+
+    Kept deliberately narrow, exactly as Part B scored it. A wider salvage would recover more
+    verdicts and would also mean the F1 quoted in the UI describes a parser that is not this one.
+    """
+    for m in reversed(scan_objects(raw)):
+        try:
+            obj = json.loads(m)
+        except json.JSONDecodeError:
+            continue
+        if "injected" in obj:
+            inj = obj["injected"]
+            inj = inj if isinstance(inj, bool) else str(inj).strip().lower() in {"true", "yes", "1"}
+            fam = str(obj.get("injection_type", "none") or "none").strip().lower()
+            return {"parse_ok": True, "parsed_by": "balanced JSON",
+                    "pred_injected": int(inj),
+                    "pred_family": fam if fam in FAMILIES else "none",
+                    "evidence": str(obj.get("evidence", ""))[:200],
+                    "reasoning": str(obj.get("reasoning", ""))[:300]}
+
+    m = re.search(r'"injected"\s*:\s*(true|false)', raw or "", re.I)
+    if m:
+        f = re.search(r'"injection_type"\s*:\s*"([a-z_]+)"', raw, re.I)
+        fam = f.group(1).lower() if f else "none"
+        return {"parse_ok": True, "parsed_by": "field regex",
+                "pred_injected": int(m.group(1).lower() == "true"),
+                "pred_family": fam if fam in FAMILIES else "none",
+                "evidence": "", "reasoning": ""}
+
+    return {"parse_ok": False, "parsed_by": "unrecoverable", "pred_injected": 0,
+            "pred_family": "none", "evidence": "", "reasoning": ""}
+
+
+# ---------------------------------------------------------------------------------------------
+# Execution - CPU, llama.cpp, Q4_K_M
+# ---------------------------------------------------------------------------------------------
+
+BASE_REPO = "XiaomiMiMo/MiMo-7B-RL"
+
+# MiMo carries Multi-Token-Prediction layers that llama.cpp cannot load, so the only GGUF that runs
+# at all is one with those layers removed. MTP is a speculative-decoding accelerator - the ordinary
+# forward pass does not use it - so greedy output should be unaffected, but "should be" is doing
+# real work in that sentence and the quantiser says so too.
+GGUF_REPO = "quantflex/MiMo-7B-RL-nomtp-GGUF"
+GGUF_FILE = "MiMo-7B-RL-nomtp-Q4_K_M.gguf"
+
+MODEL_CAVEAT = (
+    "Part B measured MiMo-7B-RL on a GPU as BF16 weights quantised to 4-bit NF4 by bitsandbytes. "
+    "A free Space has no GPU, so this runs the same base model as a Q4_K_M GGUF through llama.cpp "
+    "on the CPU, from a build with MiMo's multi-token-prediction layers stripped (llama.cpp cannot "
+    "load them). The prompt, prefill, decoding and parser are byte-identical to Part B; the "
+    "arithmetic underneath is not. Treat F1 0.945 as the figure for the configuration Part B "
+    "measured, not as a measurement of this Space."
+)
+
+N_CTX = 4096                     # a 3,000-char window is ~1,000 tokens, plus 200 generated
+
+_llm = None
+_tokenizer = None
+_lock = threading.Lock()         # one 4.7 GB model, one request at a time
+
+
+def model_path(progress=None) -> str:
+    """Fetch the GGUF from the hub, cached on the Space's disk after the first run."""
+    from huggingface_hub import hf_hub_download
+    if progress:
+        progress(f"downloading {GGUF_FILE} (4.7 GB, first run only)")
+    return hf_hub_download(GGUF_REPO, GGUF_FILE)
+
+
+def load_tokenizer():
+    """
+    MiMo's own tokenizer, used **only** to apply its chat template.
+
+    The template is the wrapper of role tags around the prompt, and Part B let each model's own
+    tokenizer write it. Doing the same here rather than hand-rolling ChatML is what keeps the
+    rendered string identical to the one that was measured. Only the tokenizer files are
+    downloaded - a few megabytes, not the weights.
+    """
+    global _tokenizer
+    if _tokenizer is None:
+        from transformers import AutoTokenizer
+        _tokenizer = AutoTokenizer.from_pretrained(BASE_REPO, trust_remote_code=True)
+    return _tokenizer
+
+
+def render_prompt(text: str) -> str:
+    """
+    Apply MiMo's chat template to one window, then append the prefill.
+
+    If the tokenizer cannot be reached the ChatML fallback is used. MiMo is a ChatML model, so this
+    produces the same string in practice - but it is a reconstruction rather than the model's own
+    template, so the caller is told which route was taken instead of the difference being silent.
+    """
+    messages = build_messages(text)
+    try:
+        rendered = load_tokenizer().apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True)
+        route = "tokenizer template"
+    except Exception:
+        rendered = "".join(f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n"
+                           for m in messages) + "<|im_start|>assistant\n"
+        route = "ChatML fallback"
+    return rendered + PREFILL, route
+
+
+def load(progress=None):
+    """Load MiMo once and keep it. ~4.7 GB resident, well inside the free tier's 16 GB."""
+    global _llm
+    if _llm is None:
+        from llama_cpp import Llama
+        path = model_path(progress)
+        if progress:
+            progress("loading MiMo-7B into memory")
+        _llm = Llama(model_path=path, n_ctx=N_CTX, n_threads=os.cpu_count() or 2,
+                     n_batch=256, logits_all=False, verbose=False)
+    return _llm
+
+
+def judge(text: str, progress=None) -> dict:
+    """
+    One window in, one parsed verdict out.
+
+    Greedy (`temperature=0.0`), 200 new tokens, exactly as Part B decoded. The prefill was given to
+    the model but not generated by it, so it is put back before parsing or the opening brace of the
+    JSON is missing and every answer parses as unrecoverable.
+    """
+    llm = load(progress)
+    prompt, route = render_prompt(text)
+    with _lock:
+        out = llm.create_completion(prompt=prompt, max_tokens=MAX_NEW, temperature=0.0,
+                                    top_k=1, stop=["<|im_end|>", "<|endoftext|>"])
+    raw = PREFILL + out["choices"][0]["text"]
+    return {**parse_response(raw), "raw": raw.strip()[:2000], "prompt_route": route}

neighbours.py

@@ -0,0 +1,105 @@
+"""
+Stage 3 - where this window sits in the corpus Part A measured.
+
+Part A's bake-off picked `nomic-ai/nomic-embed-text-v1.5` over MiniLM and BGE, embedding the
+`payload_window` column with the `search_document: ` prefix its model card specifies for indexing,
+unit-normalised, 768 dimensions. Those four facts are recorded in `part_a_results.json` under
+`winner`, and `check_provenance()` asserts them against that file rather than trusting this
+docstring - because a query embedded with a different model, or without the prefix, lands in a
+different space and every neighbour returned is meaningless, silently, with no error anywhere.
+
+What it is worth: **precision@5 of 35.6%** against a 6.8% random baseline. Fewer than 2 of the 5
+files returned are the same kind of attack as the query. That is far better than chance and it is
+not good, which is why this is labelled "nearest files in the corpus" and never "the same attack".
+"""
+
+import json
+import re
+import urllib.request
+
+import numpy as np
+import pandas as pd
+
+REPO = "Cyber-security-final-project/Evaluation_of_OpenSource_Models_for_PDF_Injection_Recognition"
+BASE = f"https://huggingface.co/datasets/{REPO}/resolve/main/"
+INDEX_URL = BASE + "Part_A_Outputs/corpus_embeddings.parquet"
+RESULTS_URL = BASE + "Part_A_Outputs/part_a_results.json"
+
+# Part A's winner. Every field here is asserted against the results file before use.
+EMBED_REPO = "nomic-ai/nomic-embed-text-v1.5"
+PREFIX = "search_document: "
+DIMS = 768
+
+# The corpus filename encodes its own attack family ("javascript_injection_OWASP_0004.pdf"), and
+# the clean controls are hash-named. That is a label on a known corpus file, not a model input -
+# nothing here ever shows a filename to a model.
+FAMILY_RE = re.compile(r"^([a-z_]+?)_(?:AMTSO|WICAR|OWASP|AtomicRedTeam|Metasploit|Glasswall|"
+                       r"mindcrypt|RanSim|RANSIM|custom)_\d+\.pdf$")
+
+_index = None
+_model = None
+
+
+def family_of(file_id: str) -> str:
+    """The attack family a corpus file carries, read off its name. Clean files are hash-named."""
+    m = FAMILY_RE.match(file_id or "")
+    return m.group(1) if m else "clean (control)"
+
+
+def check_provenance() -> dict:
+    """Assert this module embeds queries the way the index was built. Cheap, and load-bearing."""
+    with urllib.request.urlopen(RESULTS_URL) as r:
+        winner = json.loads(r.read().decode())["winner"]
+    assert winner["repo"] == EMBED_REPO, f"index built by {winner['repo']}, app uses {EMBED_REPO}"
+    assert winner["prefix"] == PREFIX, "prefix differs from the one the index was built with"
+    assert winner["dims"] == DIMS, "dimension mismatch"
+    assert winner["normalised"], "index is not unit-normalised; the dot product is not cosine"
+    assert winner["input_column"] == "payload_window", "index was not built on payload windows"
+    return winner
+
+
+def load_index():
+    """The 1,100 x 768 index Part A exported, straight from the dataset repo. ~3 MB."""
+    global _index
+    if _index is None:
+        df = pd.read_parquet(INDEX_URL)
+        _index = {"ids": df["file_id"].to_numpy(),
+                  "matrix": df.drop(columns="file_id").to_numpy(dtype=np.float32)}
+    return _index
+
+
+def load_model():
+    """Nomic-embed-v1.5 on the CPU. ~550 MB, and a 3,000-character window embeds in a second."""
+    global _model
+    if _model is None:
+        from sentence_transformers import SentenceTransformer
+        _model = SentenceTransformer(EMBED_REPO, trust_remote_code=True, device="cpu")
+    return _model
+
+
+def embed(text: str) -> np.ndarray:
+    """Embed with the prefix and normalisation Part A used. Anything else is a different space."""
+    return load_model().encode([PREFIX + text], normalize_embeddings=True,
+                               show_progress_bar=False)[0]
+
+
+def neighbours(text: str, k: int = 5) -> list:
+    """
+    The k corpus files nearest one window.
+
+    Vectors are unit-normalised, so the dot product is the cosine similarity and the whole lookup
+    against 1,100 files is one matrix-vector product - no index structure needed at this size.
+    """
+    idx = load_index()
+    sims = idx["matrix"] @ embed(text)
+    top = np.argsort(-sims)[:k]
+    return [{"file_id": str(idx["ids"][i]),
+             "family": family_of(str(idx["ids"][i])),
+             "similarity": round(float(sims[i]), 3)} for i in top]
+
+
+NEIGHBOUR_COLUMNS = ["corpus file", "its attack family", "cosine similarity"]
+
+
+def neighbour_rows(text: str, k: int = 5) -> list:
+    return [[n["file_id"], n["family"], n["similarity"]] for n in neighbours(text, k)]

requirements.txt

@@ -0,0 +1,23 @@
+# CPU-only wheels, both of them deliberate.
+#
+# torch: sentence-transformers pulls torch, and the default PyPI wheel is the CUDA build - about
+# 2.5 GB of NVIDIA libraries that a CPU Space can never use, and enough to make the image build
+# fail on size alone. The cpu index gives a ~200 MB wheel instead.
+#
+# llama-cpp-python: the sdist compiles llama.cpp from source, which on a Space build takes tens of
+# minutes and often times out. abetlen's index publishes prebuilt CPU wheels; 0.3.19 is the newest
+# one there, and is new enough to load a Qwen2-architecture GGUF like MiMo.
+--extra-index-url https://download.pytorch.org/whl/cpu
+--extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu
+
+torch
+llama-cpp-python==0.3.19
+
+gradio>=4.44
+huggingface_hub>=0.23
+transformers>=4.45
+sentence-transformers>=3.0
+einops                    # nomic-embed-v1.5's remote code imports it
+numpy
+pandas
+pyarrow                   # read the Part A embedding index

test_fidelity.py

@@ -0,0 +1,73 @@
+"""
+Does this Space read a PDF the way the corpus was read?
+
+This is the test the whole design rests on. Part A's index and Part B's F1 of 0.945 were measured
+on `skeleton_masked` and `payload_window` strings produced by the EDA notebook's extractor. If
+`corpus_text.py` produces even slightly different text, those numbers stop describing this app and
+become decoration.
+
+So it pulls real PDFs out of the generation repo, runs them through `corpus_text.py`, and compares
+the result **character by character** against the published parquet. Not "close enough" — identical.
+
+    python test_fidelity.py [n_files]
+"""
+
+import sys
+
+import pandas as pd
+
+import corpus_text
+
+CORPUS = ("https://huggingface.co/datasets/Cyber-security-final-project/"
+          "HARMLESS_Synthetic_Injected_PDFs_EDA/resolve/main/Datasets/"
+          "synthetic_corpus_part2_clustered.parquet")
+GENERATION_REPO = "Cyber-security-final-project/Generated_Injected_PDFs_HARMLESS"
+
+
+def main(n=8):
+    from huggingface_hub import hf_hub_download
+
+    corpus = pd.read_parquet(CORPUS).set_index("file_id")
+
+    # Injected and clean both, and not all from one family: masking and binary-stream handling
+    # differ between them, and a test that only saw one would pass on a broken extractor.
+    ids = list(corpus.index)
+    sample = ids[::max(1, len(ids) // n)][:n]
+
+    failures = 0
+    for fid in sample:
+        row = corpus.loc[fid]
+        try:
+            path = hf_hub_download(GENERATION_REPO, f"Output_PDFs/{fid}", repo_type="dataset")
+        except Exception as e:
+            print(f"  ? {fid}: could not fetch ({type(e).__name__})")
+            continue
+
+        skeleton, _, _ = corpus_text.build_skeleton(open(path, "rb").read())
+        masked = corpus_text.mask_leaks(skeleton)
+        window = corpus_text.payload_window(skeleton)
+
+        checks = {"skeleton_masked": (masked, row["skeleton_masked"]),
+                  "payload_window": (window, row["payload_window"])}
+
+        # The single-window path must also be what the triage returns first for a marked file,
+        # or the model reads a different string here than the corpus was scored on.
+        top = corpus_text.candidate_windows(skeleton)[0]
+
+        bad = [k for k, (got, want) in checks.items() if got != want]
+        if bad:
+            failures += 1
+            print(f"  x {fid}: differs in {', '.join(bad)}")
+            for k in bad:
+                got, want = checks[k]
+                print(f"      {k}: got {len(got):,} chars, corpus has {len(want):,}")
+        else:
+            note = "head" if top["is_head"] else f"{len(top['families'])} family marker(s)"
+            print(f"  . {fid}: identical  (triage top = {note})")
+
+    print(f"\n{len(sample) - failures}/{len(sample)} identical to the published corpus")
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(int(sys.argv[1]) if len(sys.argv) > 1 else 8))