embedder.py

"""
Stage 2 - Embedding.

Part A, applied to an uploaded file: embed it, and return the corpus files that sit nearest to it.

**Exactly what Part A measured, or the numbers do not transfer:** `nomic-ai/nomic-embed-text-v1.5`,
the `search_document: ` prefix its model card specifies for indexing, unit-normalised vectors, and
the same 1,100-file index that notebook exported. All four facts are recorded in
`part_a_results.json` under `winner`, and `check_provenance()` asserts them against the file rather
than trusting this docstring.

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 the GUI labels this "most similar files in the corpus" and never "the same
attack".
"""

import json
import re

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"

EMBEDDERS = {
    "nomic": {"repo": "nomic-ai/nomic-embed-text-v1.5", "prefix": "search_document: ",
              "trust_remote_code": True, "dims": 768,
              "label": "Nomic-embed-v1.5 (Part A winner, family purity@20 26.0%)"},
    "minilm": {"repo": "sentence-transformers/all-MiniLM-L6-v2", "prefix": "",
               "trust_remote_code": False, "dims": 384,
               "label": "MiniLM-L6-v2 (21.0%, 24x faster - index would need rebuilding)"},
}
ACTIVE = "nomic"

# The corpus filename encodes its own attack family ("javascript_injection_OWASP_0004.pdf").
# That is a label on a known corpus file, not a model input - the models never see a filename.
FAMILY_RE = re.compile(r"^([a-z_]+?)_(?:AMTSO|WICAR|OWASP|AtomicRedTeam|Metasploit|Glasswall|"
                       r"mindcrypt|RanSim|RANSIM|custom)_\d+\.pdf$")

_index = {}      # cached (file_ids, matrix)
_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"


def load_index(url: str = INDEX_URL):
    """The 1,100 x 768 index Part A exported, straight from the dataset repo."""
    global _index
    if not _index:
        df = pd.read_parquet(url)
        ids = df["file_id"].to_numpy()
        mat = df.drop(columns="file_id").to_numpy(dtype=np.float32)
        _index = {"ids": ids, "matrix": mat}
    return _index


def check_provenance(url: str = RESULTS_URL) -> dict:
    """
    Assert this module embeds queries the way the index was built.

    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.
    """
    import urllib.request
    with urllib.request.urlopen(url) as r:
        winner = json.loads(r.read().decode())["winner"]

    spec = EMBEDDERS[ACTIVE]
    assert winner["repo"] == spec["repo"], f"index built by {winner['repo']}, app uses {spec['repo']}"
    assert winner["prefix"] == spec["prefix"], "prefix differs from the one the index was built with"
    assert winner["dims"] == spec["dims"], "dimension mismatch"
    assert winner["normalised"], "index is not unit-normalised; dot product is not cosine"
    return winner


def load_model(name: str = None):
    global _model
    name = name or ACTIVE
    if _model is None:
        from sentence_transformers import SentenceTransformer
        spec = EMBEDDERS[name]
        _model = SentenceTransformer(spec["repo"], trust_remote_code=spec["trust_remote_code"])
    return _model


def embed(texts, name: str = None) -> np.ndarray:
    """Embed with the prefix and normalisation Part A used. Anything else is a different space."""
    spec = EMBEDDERS[name or ACTIVE]
    model = load_model(name)
    texts = [texts] if isinstance(texts, str) else list(texts)
    return model.encode([spec["prefix"] + t for t in texts],
                        normalize_embeddings=True, show_progress_bar=False)


def neighbours(text: str, k: int = 5, name: str = None) -> list:
    """
    The k corpus files nearest one window of an uploaded PDF.

    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()
    q = embed(text, name)[0]
    sims = idx["matrix"] @ q
    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]


def neighbours_for(summary: dict, windows: list, k: int = 5, name: str = None) -> dict:
    """
    Neighbours for the part of the document that matters.

    If the detector flagged something, the query is the first flagged region - "what known attacks
    does this resemble". If it found nothing, the query is the head of the document, which is what
    Part A embedded for clean corpus files, so the comparison stays like-for-like.
    """
    if summary["regions"]:
        r = summary["regions"][0]
        query = next((w["text"] for w in windows if w["index"] == r["windows"][0]), windows[0]["text"])
        basis = f"region 1 (characters {r['start']:,}-{r['end']:,})"
    else:
        query = windows[0]["text"] if windows else ""
        basis = "the head of the document (nothing was flagged)"
    return {"basis": basis, "neighbours": neighbours(query, k=k, name=name) if query else []}


def neighbour_rows(result: dict) -> list:
    return [[n["file_id"], n["family"], n["similarity"]] for n in result["neighbours"]]


NEIGHBOUR_COLUMNS = ["corpus file", "its attack family", "cosine similarity"]