Joeven — Learn Autonomous AI Agents

Python loads on first run

import math

def cosine(a, b):
    na = math.sqrt(sum(x * x for x in a)) or 1.0
    nb = math.sqrt(sum(x * x for x in b)) or 1.0
    return sum(x * y for x, y in zip(a, b)) / (na * nb)

docs = []
for i in range(64):
    vec = [((i * 3 + d * 7) % 11) / 10 - 0.5 for d in range(4)]
    docs.append({"id": i, "vec": vec, "bucket": 0 if vec[0] >= 0 else 1})

query = [0.4, -0.1, 0.2, 0.0]
q_bucket = 0 if query[0] >= 0 else 1

def brute(q, k=5):
    ranked = sorted(docs, key=lambda d: cosine(q, d["vec"]), reverse=True)
    return ranked[:k]

def ann_bucket(q, k=5):
    subset = [d for d in docs if d["bucket"] == q_bucket]
    ranked = sorted(subset, key=lambda d: cosine(q, d["vec"]), reverse=True)
    return ranked[:k]

b = brute(query)
a = ann_bucket(query)
print("brute ids:", [d["id"] for d in b])
print("ann ids:  ", [d["id"] for d in a])
overlap = {d["id"] for d in b} & {d["id"] for d in a}
print("recall@5 vs brute:", len(overlap) / 5)
print("ann scanned", sum(1 for d in docs if d["bucket"] == q_bucket), "of", len(docs))

Run to execute this in your browser. Nothing is sent to a server.