mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-09-10 11:48:00 +00:00
Removes the benchmarking file
This commit is contained in:
Trenton Holmes
committed by
stumpylog
stumpylog
parent
ebd590725c
commit
5d2aacab6d
@@ -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()
|
||||
Reference in New Issue
Block a user