From 5d2aacab6d3a08f2f544953053011e7b4aeedcd5 Mon Sep 17 00:00:00 2001 From: Trenton Holmes <797416+stumpylog@users.noreply.github.com> Date: Fri, 12 Jun 2026 06:59:28 -0700 Subject: [PATCH] Removes the benchmarking file --- src/bench_vector_store.py | 324 -------------------------------------- 1 file changed, 324 deletions(-) delete mode 100644 src/bench_vector_store.py diff --git a/src/bench_vector_store.py b/src/bench_vector_store.py deleted file mode 100644 index cf7e678b6..000000000 --- a/src/bench_vector_store.py +++ /dev/null @@ -1,324 +0,0 @@ -#!/usr/bin/env python3 -"""Head-to-head benchmark: PaperlessLanceVectorStore vs PaperlessSqliteVecVectorStore. - -Run from src/ with: - uv run python bench_vector_store.py [OPTIONS] - -Phase 1 (skipped if bench_data.pkl already exists): generate fake documents with -Faker and embed chunks via Ollama; save to disk for reuse. -Phase 2: benchmark both stores against identical data and print a comparison table. - -Requires both classes to coexist in paperless_ai.vector_store (Task 3 Phase A). -After Phase B replaces the file, the Lance import fails gracefully and only the -sqlite-vec half runs. -""" -from __future__ import annotations - -import argparse -import pickle -import statistics -import tempfile -import time -import uuid -from pathlib import Path -from typing import TypedDict - -import httpx -from faker import Faker -from llama_index.core.schema import TextNode -from llama_index.core.vector_stores.types import ( - FilterOperator, - MetadataFilter, - MetadataFilters, - VectorStoreQuery, -) - -try: - from paperless_ai.vector_store import PaperlessLanceVectorStore - - _LANCE_OK = True -except ImportError: - _LANCE_OK = False - -from paperless_ai.vector_store import PaperlessSqliteVecVectorStore - -DEFAULT_OLLAMA_URL = "http://192.168.1.87:11434" -DEFAULT_EMBED_MODEL = "qwen3-embedding:4b" -DEFAULT_DATA_FILE = "bench_data.pkl" -DEFAULT_N_DOCS = 2000 -DEFAULT_CHUNKS_PER_DOC = 3 -DEFAULT_QUERY_ITERS = 50 -_BATCH = 32 - - -def _embed(texts: list[str], url: str, model: str) -> list[list[float]]: - r = httpx.post( - f"{url}/api/embed", - json={"model": model, "input": texts}, - timeout=120.0, - ) - r.raise_for_status() - return r.json()["embeddings"] - - -def warm_up(url: str, model: str) -> int: - """Fire one embed call to load the model into GPU; return embedding dim.""" - print(f"Warming up {model}...", end=" ", flush=True) - dim = len(_embed(["warm"], url, model)[0]) - print(f"dim={dim}") - return dim - - -def generate_and_save( - n_docs: int, - chunks_per_doc: int, - url: str, - model: str, - out: str, -) -> list[dict]: - fake = Faker() - Faker.seed(42) - print(f"Generating {n_docs} docs ({chunks_per_doc} chunks each)...") - docs = [] - for i in range(n_docs): - body = "\n\n".join(fake.paragraph(nb_sentences=8) for _ in range(3)) - clen = max(1, len(body) // chunks_per_doc) - chunks = [] - for j in range(chunks_per_doc): - s = j * clen - e = s + clen if j < chunks_per_doc - 1 else len(body) - chunks.append({"node_id": str(uuid.uuid4()), "text": body[s:e], "embedding": None}) - docs.append({ - "doc_id": str(i + 1), - "title": fake.catch_phrase(), - "modified": fake.date_time_this_decade().isoformat(), - "chunks": chunks, - }) - - all_texts = [c["text"] for d in docs for c in d["chunks"]] - print(f"Embedding {len(all_texts)} chunks in batches of {_BATCH}...") - embeddings: list[list[float]] = [] - for i in range(0, len(all_texts), _BATCH): - embeddings.extend(_embed(all_texts[i : i + _BATCH], url, model)) - print(f" {min(i + _BATCH, len(all_texts))}/{len(all_texts)}", end="\r", flush=True) - print() - - idx = 0 - for d in docs: - for c in d["chunks"]: - c["embedding"] = embeddings[idx] - idx += 1 - - with open(out, "wb") as f: - pickle.dump(docs, f) - print(f"Saved to {out}") - return docs - - -def _build_nodes(docs: list[dict]) -> list[TextNode]: - nodes = [] - for d in docs: - for c in d["chunks"]: - n = TextNode( - id_=c["node_id"], - text=c["text"], - metadata={"document_id": d["doc_id"], "modified": d["modified"]}, - ) - n.relationships = {} - n.embedding = c["embedding"] - nodes.append(n) - return nodes - - -def _in_filter(ids: list[str]) -> MetadataFilters: - return MetadataFilters( - filters=[MetadataFilter(key="document_id", operator=FilterOperator.IN, value=ids)] - ) - - -def _dir_bytes(path: str) -> int: - return sum(f.stat().st_size for f in Path(path).rglob("*") if f.is_file()) - - -def _sqlite_bytes(uri: str) -> int: - p = Path(uri) / "llmindex.db" - return p.stat().st_size if p.exists() else 0 - - -def run_bench( - store, - nodes: list[TextNode], - docs: list[dict], - q_iters: int, - is_lance: bool, -) -> dict: - doc_ids = [d["doc_id"] for d in docs] - filter_ids = doc_ids[: max(1, len(doc_ids) // 5)] - q_vecs = [nodes[i * 10 % len(nodes)].embedding for i in range(q_iters)] - by_doc: dict[str, list[TextNode]] = {} - for n in nodes: - by_doc.setdefault(n.metadata["document_id"], []).append(n) - uri = store._uri # noqa: SLF001 - - # insert - t0 = time.perf_counter() - store.add(list(nodes)) - r: dict = {"insert": time.perf_counter() - t0} - - # query plain - times = [] - for emb in q_vecs: - t0 = time.perf_counter() - store.query(VectorStoreQuery(query_embedding=emb, similarity_top_k=10)) - times.append(time.perf_counter() - t0) - r["qp50"] = statistics.median(times) - r["qp95"] = sorted(times)[int(len(times) * 0.95)] - - # query filtered - times = [] - flt = _in_filter(filter_ids) - for emb in q_vecs: - t0 = time.perf_counter() - store.query(VectorStoreQuery(query_embedding=emb, similarity_top_k=10, filters=flt)) - times.append(time.perf_counter() - t0) - r["qfp50"] = statistics.median(times) - r["qfp95"] = sorted(times)[int(len(times) * 0.95)] - - # get_modified_times - times = [] - for _ in range(20): - t0 = time.perf_counter() - store.get_modified_times() - times.append(time.perf_counter() - t0) - r["gmt_p50"] = statistics.median(times) - - # upsert (fresh node IDs, same embeddings) - times = [] - for doc in docs[:q_iters]: - orig = by_doc.get(doc["doc_id"], []) - if not orig: - continue - fresh = [] - for o in orig: - fn = TextNode( - id_=str(uuid.uuid4()), - text=o.text, - metadata=o.metadata.copy(), - ) - fn.relationships = {} - fn.embedding = o.embedding - fresh.append(fn) - t0 = time.perf_counter() - store.upsert_document(doc["doc_id"], fresh) - times.append(time.perf_counter() - t0) - r["up50"] = statistics.median(times) if times else 0.0 - r["up95"] = sorted(times)[int(len(times) * 0.95)] if times else 0.0 - - r["size_pre"] = _dir_bytes(uri) if is_lance else _sqlite_bytes(uri) - - # compact - t0 = time.perf_counter() - if is_lance: - store.compact(retention_seconds=0) - else: - store.compact(force=True) - r["compact"] = time.perf_counter() - t0 - - r["size_post"] = _dir_bytes(uri) if is_lance else _sqlite_bytes(uri) - return r - - -def _pct(lv: float | None, sv: float) -> str: - if lv is None or lv == 0: - return "N/A" - p = (sv - lv) / lv * 100 - return f"{'+' if p > 0 else ''}{p:.0f}%" - - -def print_results(nodes: list[TextNode], q_iters: int, lance: dict | None, sq: dict) -> None: - W = 30 - n, dim = len(nodes), len(nodes[0].embedding) - print(f"\n=== Vector Store Benchmark ===") - print(f"Nodes: {n} | Dim: {dim} | Query iters: {q_iters}\n") - lh = "LanceDB" if lance else "LanceDB (N/A)" - print(f"{'Operation':<{W}} {lh:<22} {'sqlite-vec':<22} {'Delta'}") - print("-" * (W + 66)) - - def _s(v: float) -> str: - return f"{v:.3f}s" - - def _ms(v: float) -> str: - return f"{v * 1000:.1f}ms" - - def _mb(v: float) -> str: - return f"{v / 1e6:.1f} MB" - - def row(label: str, lv: float | None, sv: float, fmt) -> None: - ls = fmt(lv) if lv is not None else "N/A" - print(f"{label:<{W}} {ls:<22} {fmt(sv):<22} {_pct(lv, sv)}") - - def row2(label: str, lv1: float | None, lv2: float | None, sv1: float, sv2: float) -> None: - def ms_pair(a: float, b: float) -> str: - return f"{_ms(a)} / {_ms(b)}" - ls = ms_pair(lv1, lv2) if lv1 is not None else "N/A" - print(f"{label:<{W}} {ls:<22} {ms_pair(sv1, sv2):<22} {_pct(lv1, sv1)}") - - L = lance - row(f"insert ({n} nodes)", L["insert"] if L else None, sq["insert"], _s) - row2("query plain p50/p95", - L["qp50"] if L else None, L["qp95"] if L else None, sq["qp50"], sq["qp95"]) - row2("query filtered p50/p95", - L["qfp50"] if L else None, L["qfp95"] if L else None, sq["qfp50"], sq["qfp95"]) - row("get_modified_times p50", L["gmt_p50"] if L else None, sq["gmt_p50"], _ms) - row2("upsert p50/p95", - L["up50"] if L else None, L["up95"] if L else None, sq["up50"], sq["up95"]) - row("compact", L["compact"] if L else None, sq["compact"], _s) - row("file size pre-compact", L["size_pre"] if L else None, sq["size_pre"], _mb) - row("file size post-compact", L["size_post"] if L else None, sq["size_post"], _mb) - - -def main() -> None: - ap = argparse.ArgumentParser(description="Vector store head-to-head benchmark") - ap.add_argument("--n-docs", type=int, default=DEFAULT_N_DOCS) - ap.add_argument("--chunks-per-doc", type=int, default=DEFAULT_CHUNKS_PER_DOC) - ap.add_argument("--data-file", default=DEFAULT_DATA_FILE) - ap.add_argument("--regenerate", action="store_true") - ap.add_argument("--ollama-url", default=DEFAULT_OLLAMA_URL) - ap.add_argument("--embed-model", default=DEFAULT_EMBED_MODEL) - ap.add_argument("--query-iters", type=int, default=DEFAULT_QUERY_ITERS) - args = ap.parse_args() - - warm_up(args.ollama_url, args.embed_model) - - data_path = Path(args.data_file) - if args.regenerate or not data_path.exists(): - docs = generate_and_save( - args.n_docs, args.chunks_per_doc, args.ollama_url, args.embed_model, args.data_file - ) - else: - print(f"Loading {args.data_file}...") - with open(data_path, "rb") as f: - docs = pickle.load(f) - print(f"Loaded {len(docs)} docs ({sum(len(d['chunks']) for d in docs)} nodes)") - - all_nodes = _build_nodes(docs) - - lance_r = None - if _LANCE_OK: - print("\nBenchmarking LanceDB...") - with tempfile.TemporaryDirectory() as d: - store = PaperlessLanceVectorStore(uri=d) - lance_r = run_bench(store, all_nodes, docs, args.query_iters, is_lance=True) - else: - print("Skipping LanceDB (PaperlessLanceVectorStore not importable).") - - print("\nBenchmarking sqlite-vec...") - with tempfile.TemporaryDirectory() as d: - store = PaperlessSqliteVecVectorStore(uri=d) - sqlite_r = run_bench(store, all_nodes, docs, args.query_iters, is_lance=False) - - print_results(all_nodes, args.query_iters, lance_r, sqlite_r) - - -if __name__ == "__main__": - main()