scripts/embed-eval/knn-bench.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+"""Leave-one-out nearest-neighbour benchmark for an embedding model.
+
+Embeds a labelled corpus of Paperless documents and, for each document, predicts
+its correspondent and document type from the label of its nearest neighbour
+(cosine similarity, the document itself excluded). Accuracy is the share of
+documents whose predicted label equals the filed one.
+
+Corpus: JSON lines {"id", "correspondent", "doctype", "text"}, exported
+read-only from the paperless database, for example on the CNPG primary:
+
+    select json_build_object('id', d.id, 'correspondent', d.correspondent_id,
+        'doctype', d.document_type_id,
+        'text', left(regexp_replace(d.content, '\\s+', ' ', 'g'), 3000))
+    from documents_document d
+    where d.correspondent_id is not null and d.document_type_id is not null
+      and length(d.content) > 200
+    order by md5(d.id::text || '<seed>') limit 400
+
+Ceiling is the share of documents whose label occurs on at least one other
+document; no nearest-neighbour method can exceed it. Floor is the share of the
+most common label.
+
+Usage:
+    knn-bench.py --corpus corpus.jsonl -e https://api.mistral.ai/v1/embeddings \\
+        -m mistral-embed --key-env MISTRAL_API_KEY --batch 8
+    knn-bench.py --corpus corpus.jsonl -m bge-m3 \\
+        -e http://llama-embed.ai.svc.cluster.local:8080/v1/embeddings \\
+        --via-pod "deploy/llama-server -n ai"
+
+Standard library only. Vectors are cached per model next to the corpus.
+"""
+import argparse
+import json
+import math
+import os
+import re
+import shlex
+import subprocess
+import sys
+import time
+import urllib.error
+import urllib.request
+from collections import Counter
+
+
+def post(endpoint, payload, key, via_pod):
+    body = json.dumps(payload).encode()
+    if via_pod:
+        cmd = ["kubectl", "exec", "-i", *shlex.split(via_pod), "--", "curl", "-sS", "-m", "300",
+               "-H", "Content-Type: application/json", "-d", "@-", endpoint]
+        proc = subprocess.run(cmd, input=body, capture_output=True, check=False)
+        if proc.returncode != 0:
+            raise RuntimeError(f"kubectl exec failed: {proc.stderr.decode()}")
+        return json.loads(proc.stdout)
+    headers = {"Content-Type": "application/json"}
+    if key:
+        headers["Authorization"] = f"Bearer {key}"
+    req = urllib.request.Request(endpoint, data=body, headers=headers)
+    for attempt in range(6):
+        try:
+            with urllib.request.urlopen(req, timeout=300) as resp:
+                return json.load(resp)
+        except urllib.error.HTTPError as err:
+            if err.code == 429 and attempt < 5:
+                time.sleep(2 ** attempt)
+                continue
+            raise RuntimeError(f"HTTP {err.code}: {err.read().decode()}") from err
+    raise RuntimeError("retries exhausted")
+
+
+def embed_all(docs, args, key):
+    cache = f"{args.corpus}.{args.model}.{args.max_chars}.vecs.json"
+    if os.path.exists(cache):
+        with open(cache) as fh:
+            vecs = json.load(fh)
+        if len(vecs) == len(docs):
+            print(f"vectors loaded from {cache}")
+            return vecs
+    vecs = []
+    shrunk = 0
+    start = time.time()
+    for i in range(0, len(docs), args.batch):
+        chunk = [d["text"] for d in docs[i:i + args.batch]]
+        for _ in range(6):
+            try:
+                out = post(args.endpoint, {"model": args.model, "input": chunk}, key, args.via_pod)
+                break
+            except RuntimeError as err:
+                # Mistral reports "Input id N has T tokens, exceeding max M tokens": cut that text to fit.
+                hit = re.search(r"Input id (\d+) has (\d+) tokens, exceeding max (\d+)", str(err))
+                if not hit:
+                    raise
+                idx, have, limit = (int(g) for g in hit.groups())
+                chunk[idx] = chunk[idx][:int(len(chunk[idx]) * limit / have * 0.95)]
+                shrunk += 1
+        else:
+            raise RuntimeError(f"documents {i}-{i + len(chunk) - 1} still over the token limit")
+        if "data" not in out:
+            raise RuntimeError(f"no embeddings in response for documents {i}-{i + len(chunk) - 1}: {json.dumps(out)[:500]}")
+        data = sorted(out["data"], key=lambda item: item["index"])
+        if len(data) != len(chunk):
+            raise RuntimeError(f"asked for {len(chunk)} vectors, got {len(data)}")
+        vecs.extend(item["embedding"] for item in data)
+        done = len(vecs)
+        if done % 50 < args.batch:
+            print(f"  embedded {done}/{len(docs)} ({time.time() - start:.0f}s)", flush=True)
+    if shrunk:
+        print(f"  {shrunk} over-limit texts were shortened to fit the model's token limit")
+    with open(cache, "w") as fh:
+        json.dump(vecs, fh)
+    return vecs
+
+
+def normalise(vec):
+    norm = math.sqrt(sum(x * x for x in vec)) or 1.0
+    return [x / norm for x in vec]
+
+
+def nearest_neighbours(vecs):
+    n = len(vecs)
+    best = [(-2.0, -1)] * n
+    for i in range(n):
+        vi = vecs[i]
+        for j in range(i + 1, n):
+            s = sum(a * b for a, b in zip(vi, vecs[j]))
+            if s > best[i][0]:
+                best[i] = (s, j)
+            if s > best[j][0]:
+                best[j] = (s, i)
+    return [j for _, j in best]
+
+
+def score(docs, neighbours, field):
+    labels = [d[field] for d in docs]
+    counts = Counter(labels)
+    hits = sum(labels[i] == labels[neighbours[i]] for i in range(len(docs)))
+    ceiling = sum(counts[label] > 1 for label in labels)
+    floor = max(counts.values())
+    n = len(docs)
+    return 100 * hits / n, 100 * ceiling / n, 100 * floor / n, len(counts)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--corpus", required=True)
+    ap.add_argument("-e", "--endpoint", required=True)
+    ap.add_argument("-m", "--model", required=True)
+    ap.add_argument("--key-env", default="", help="environment variable holding the API key")
+    ap.add_argument("--via-pod", default="", help='run curl in this pod, e.g. "deploy/llama-server -n ai"')
+    ap.add_argument("--batch", type=int, default=1)
+    ap.add_argument("--max-chars", type=int, default=0, help="truncate each text to this many characters (0 = as exported)")
+    args = ap.parse_args()
+
+    key = os.environ.get(args.key_env, "") if args.key_env else ""
+    if args.key_env and not key:
+        sys.exit(f"{args.key_env} is not set")
+
+    with open(args.corpus) as fh:
+        docs = [json.loads(line) for line in fh if line.strip()]
+    if args.max_chars:
+        for d in docs:
+            d["text"] = d["text"][:args.max_chars]
+    print(f"model={args.model} documents={len(docs)} max_chars={args.max_chars} endpoint={args.endpoint}")
+
+    start = time.time()
+    vecs = [normalise(v) for v in embed_all(docs, args, key)]
+    embed_s = time.time() - start
+    neighbours = nearest_neighbours(vecs)
+
+    print(f"dimensions={len(vecs[0])} embed+load={embed_s:.0f}s")
+    for field in ("correspondent", "doctype"):
+        acc, ceiling, floor, classes = score(docs, neighbours, field)
+        print(f"{field:<14} accuracy={acc:5.1f}%  ceiling={ceiling:5.1f}%  floor={floor:5.1f}%  labels={classes}")
+
+
+if __name__ == "__main__":
+    main()

scripts/embed-eval/knn-compare.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+"""Paired comparison of two embedding models scored by knn-bench.py.
+
+Both runs must cover the same documents in the same order. For each label field
+it reports how the per-document outcomes overlap, an exact two-sided McNemar
+test on the documents where exactly one model is right, and how often the two
+models pick the same nearest neighbour.
+
+Usage:
+    knn-compare.py --corpus-a a.jsonl --vecs-a a.vecs.json \\
+                   --corpus-b b.jsonl --vecs-b b.vecs.json
+
+The vectors files are the caches knn-bench.py writes next to its corpus.
+Standard library only.
+"""
+import argparse
+import importlib.util
+import json
+import math
+import os
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+spec = importlib.util.spec_from_file_location("knn_bench", os.path.join(HERE, "knn-bench.py"))
+bench = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(bench)
+
+
+def load(corpus, vecs):
+    with open(corpus) as fh:
+        docs = [json.loads(line) for line in fh if line.strip()]
+    with open(vecs) as fh:
+        vectors = [bench.normalise(v) for v in json.load(fh)]
+    if len(docs) != len(vectors):
+        raise SystemExit(f"{vecs}: {len(vectors)} vectors for {len(docs)} documents")
+    return docs, bench.nearest_neighbours(vectors)
+
+
+def mcnemar_exact(only_a, only_b):
+    """Two-sided exact p-value for a fair coin over the discordant pairs."""
+    n = only_a + only_b
+    if n == 0:
+        return 1.0
+    k = min(only_a, only_b)
+    tail = sum(math.comb(n, i) for i in range(k + 1)) / 2 ** n
+    return min(1.0, 2 * tail)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    for side in "ab":
+        ap.add_argument(f"--corpus-{side}", required=True)
+        ap.add_argument(f"--vecs-{side}", required=True)
+    args = ap.parse_args()
+
+    docs_a, nn_a = load(args.corpus_a, args.vecs_a)
+    docs_b, nn_b = load(args.corpus_b, args.vecs_b)
+    if [d["id"] for d in docs_a] != [d["id"] for d in docs_b]:
+        raise SystemExit("the two corpora do not hold the same documents in the same order")
+
+    n = len(docs_a)
+    same_nn = sum(a == b for a, b in zip(nn_a, nn_b))
+    print(f"documents={n}  same nearest neighbour in both models: {same_nn} ({100 * same_nn / n:.1f}%)")
+    for field in ("correspondent", "doctype"):
+        ok_a = [docs_a[i][field] == docs_a[nn_a[i]][field] for i in range(n)]
+        ok_b = [docs_b[i][field] == docs_b[nn_b[i]][field] for i in range(n)]
+        both = sum(a and b for a, b in zip(ok_a, ok_b))
+        only_a = sum(a and not b for a, b in zip(ok_a, ok_b))
+        only_b = sum(b and not a for a, b in zip(ok_a, ok_b))
+        neither = n - both - only_a - only_b
+        p = mcnemar_exact(only_a, only_b)
+        print(f"{field:<14} both right={both}  only A={only_a}  only B={only_b}  both wrong={neither}  "
+              f"McNemar p={p:.3f}")
+
+
+if __name__ == "__main__":
+    main()