179 lines
7.1 KiB
Python
179 lines
7.1 KiB
Python
#!/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()
|