diff --git a/scripts/embed-eval/knn-bench.py b/scripts/embed-eval/knn-bench.py new file mode 100644 index 00000000..94286aba --- /dev/null +++ b/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 || '') 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() diff --git a/scripts/embed-eval/knn-compare.py b/scripts/embed-eval/knn-compare.py new file mode 100644 index 00000000..df4e5d65 --- /dev/null +++ b/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()