65 lines
2.9 KiB
Python
65 lines
2.9 KiB
Python
"""Synthetic, isolated exact-search comparison; does not access the user Vault."""
|
|
import heapq
|
|
import json
|
|
import math
|
|
import random
|
|
import sqlite3
|
|
import statistics
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import sqlite_vec
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
from app.retrieval import space_index
|
|
from app.retrieval.routed_vectors import RemoteEmbeddings, _unit_vector
|
|
|
|
|
|
def main():
|
|
rng = random.Random(42)
|
|
count, dimensions = 4000, 384
|
|
with tempfile.TemporaryDirectory(prefix='notes-vec-bench-') as temporary:
|
|
conn = sqlite3.connect(Path(temporary) / 'vectors.db')
|
|
conn.row_factory = sqlite3.Row
|
|
conn.enable_load_extension(True)
|
|
sqlite_vec.load(conn)
|
|
conn.enable_load_extension(False)
|
|
conn.execute('CREATE TABLE blocks(block_id TEXT PRIMARY KEY, embedding_local_only INTEGER)')
|
|
conn.execute('CREATE TABLE routed_block_vectors(space_id TEXT,block_id TEXT,dimensions INTEGER,vector TEXT,PRIMARY KEY(space_id,dimensions,block_id))')
|
|
vectors = [_unit_vector([rng.uniform(-1, 1) for _ in range(dimensions)], dimensions) for _ in range(count)]
|
|
conn.executemany('INSERT INTO blocks VALUES (?,0)', [(str(i),) for i in range(count)])
|
|
conn.executemany('INSERT INTO routed_block_vectors VALUES (?,?,?,?)', [('benchmark', str(i), dimensions, json.dumps(v)) for i, v in enumerate(vectors)])
|
|
start = time.perf_counter()
|
|
space_index.ensure(conn, 'benchmark', dimensions)
|
|
migration_ms = (time.perf_counter() - start) * 1000
|
|
conn.commit()
|
|
query = vectors[0]
|
|
batch = RemoteEmbeddings('benchmark', dimensions, [query])
|
|
def legacy():
|
|
def hits():
|
|
for row in conn.execute('SELECT block_id,vector FROM routed_block_vectors'):
|
|
vector = _unit_vector(json.loads(row[1]), dimensions)
|
|
yield row[0], max(0., min(1., math.fsum(a*b for a,b in zip(query,vector))))
|
|
return heapq.nlargest(20, hits(), key=lambda hit:hit[1])
|
|
def native():
|
|
return [(hit.id,hit.score) for hit in space_index.search(conn,batch,20)]
|
|
measurements = {}
|
|
results = {}
|
|
for name, operation in [('python_json_scan', legacy), ('sqlite_vec',native)]:
|
|
elapsed = []
|
|
for _ in range(5):
|
|
start = time.perf_counter()
|
|
results[name] = operation()
|
|
elapsed.append((time.perf_counter()-start)*1000)
|
|
measurements[name] = {'median_ms':statistics.median(elapsed), 'samples_ms':elapsed}
|
|
assert [hit[0] for hit in results['python_json_scan']] == [hit[0] for hit in results['sqlite_vec']]
|
|
print(json.dumps({'blocks':count,'dimensions':dimensions,'top_k':20,'migration_ms':migration_ms,
|
|
'same_top_k':True,'measurements':measurements},indent=2))
|
|
conn.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|