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