my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	#!/usr/bin/env python3
#2	"""
#3	Mnemosyne vs Vector DBs: Head-to-Head BEAM Retrieval Benchmark
#4	===============================================================
#5	Compares Mnemosyne (FTS5+vec hybrid) against FAISS (flat L2) and
#6	ChromaDB (HNSW) on the same BEAM dataset. Measures:
#7
#8	- Ingestion throughput (msgs/sec)
#9	- Storage size (MB)
#10	- Retrieval latency (ms) - avg, p50, p95, p99
#11	- Recall@10 (fraction of gold messages found)
#12
#13	Usage:
#14	cd /root/.hermes/projects/mnemosyne
#15	python3 tests/benchmark_vs_vectordbs.py --scales 100K,500K,1M
#16	"""
#17
#18	# Raise fd limit before any imports
#19	import resource
#20	resource.setrlimit(resource.RLIMIT_NOFILE, (65536, 65536))
#21
#22	import argparse
#23	import ast
#24	import gc
#25	import json
#26	import os
#27	import statistics
#28	import sys
#29	import tempfile
#30	import time
#31	from collections import defaultdict
#32	from pathlib import Path
#33	from typing import Dict, List
#34
#35	import numpy as np
#36
#37	PROJECT_ROOT = Path(__file__).resolve().parent.parent
#38	sys.path.insert(0, str(PROJECT_ROOT))
#39
#40	from mnemosyne.core.beam import BeamMemory, init_beam
#41
#42	# --- Vector DB imports ---
#43	try:
#44	import faiss
#45	HAS_FAISS = True
#46	except ImportError:
#47	HAS_FAISS = False
#48	print("FAISS not installed. Skipping FAISS comparison.")
#49
#50	try:
#51	import chromadb
#52	from chromadb.config import Settings
#53	HAS_CHROMA = True
#54	except ImportError:
#55	HAS_CHROMA = False
#56	print("ChromaDB not installed. Skipping ChromaDB comparison.")
#57
#58	# --- Config ---
#59	DEFAULT_SCALES = ["100K"]
#60	DEFAULT_TOP_K = 10
#61	BENCHMARK_QUERIES_PER_SCALE = 100
#62
#63	# --- Embedding model for FAISS/Chroma ---
#64	# Use Mnemosyne's embedder (already loaded, no extra install)
#65	EMBEDDER = None
#66
#67	def get_embedder():
#68	global EMBEDDER
#69	if EMBEDDER is None:
#70	from mnemosyne.core import embeddings as _emb
#71	EMBEDDER = _emb
#72	return EMBEDDER
#73
#74
#75	def load_beam_data(scales: List[str], max_convs: int = 3) -> Dict:
#76	"""Load BEAM conversations and their gold-standard question/answer pairs."""
#77	from datasets import load_dataset
#78
#79	data = {}
#80	for scale in scales:
#81	print(f" Loading BEAM {scale}...")
#82	conversations = []
#83
#84	if scale == "10M":
#85	ds = load_dataset("Mohammadta/BEAM-10M", streaming=True)
#86	split = "10M" if "10M" in ds else list(ds.keys())[0]
#87
#88	for i, sample in enumerate(ds[split]):
#89	if max_convs and i >= max_convs:
#90	break
#91
#92	# Extract messages
#93	plans = sample.get("plans", [])
#94	messages = []
#95	for plan in plans:
#96	blocks = plan.get("chat", []) if isinstance(plan, dict) else []
#97	for block in blocks:
#98	if isinstance(block, list):
#99	for msg in block:
#100	if isinstance(msg, dict):
#101	messages.append(msg.get("content", ""))
#102
#103	# Extract questions + gold message IDs
#104	probing = sample.get("probing_questions", {})
#105	if isinstance(probing, str):
#106	probing = ast.literal_eval(probing)
#107
#108	questions = []
#109	for ability, qs in probing.items():
#110	if isinstance(qs, list):
#111	for q in qs:
#112	if isinstance(q, dict):
#113	source_ids = q.get("source_chat_ids", {})
#114	all_sources = []
#115	if isinstance(source_ids, dict):
#116	for v in source_ids.values():
#117	all_sources.extend(v if isinstance(v, list) else [v])
#118
#119	questions.append({
#120	"question": q.get("question", ""),
#121	"gold_message_indices": [int(s) for s in all_sources if str(s).isdigit()],
#122	})
#123
#124	conversations.append({
#125	"id": str(i),
#126	"messages": messages,
#127	"questions": questions[:BENCHMARK_QUERIES_PER_SCALE],
#128	})
#129
#130	ds.cleanup_cache_files() if hasattr(ds, 'cleanup_cache_files') else None
#131	del ds
#132	gc.collect()
#133	else:
#134	ds = load_dataset("Mohammadta/BEAM", streaming=True)
#135	if scale not in ds:
#136	continue
#137
#138	for i, sample in enumerate(ds[scale]):
#139	if max_convs and i >= max_convs:
#140	break
#141
#142	# Extract messages
#143	blocks = sample.get("chat", [])
#144	messages = []
#145	for block in blocks:
#146	if isinstance(block, list):
#147	for msg in block:
#148	if isinstance(msg, dict):
#149	messages.append(msg.get("content", ""))
#150
#151	# Extract questions
#152	probing = sample.get("probing_questions", "{}")
#153	if isinstance(probing, str):
#154	probing = ast.literal_eval(probing)
#155
#156	questions = []
#157	for ability, qs in probing.items():
#158	if isinstance(qs, list):
#159	for q in qs:
#160	if isinstance(q, dict):
#161	source_ids = q.get("source_chat_ids", {})
#162	all_sources = []
#163	if isinstance(source_ids, dict):
#164	for v in source_ids.values():
#165	all_sources.extend(v if isinstance(v, list) else [v])
#166
#167	questions.append({
#168	"question": q.get("question", ""),
#169	"gold_message_indices": [int(s) for s in all_sources if str(s).isdigit()],
#170	})
#171
#172	conversations.append({
#173	"id": str(i),
#174	"messages": messages,
#175	"questions": questions[:BENCHMARK_QUERIES_PER_SCALE],
#176	})
#177
#178	ds.cleanup_cache_files() if hasattr(ds, 'cleanup_cache_files') else None
#179	del ds
#180	gc.collect()
#181
#182	data[scale] = conversations
#183	print(f" {len(conversations)} conversations, {sum(len(c['messages']) for c in conversations)} msgs, {sum(len(c['questions']) for c in conversations)} questions")
#184
#185	return data
#186
#187
#188	# ============================================================
#189	# Benchmark: Mnemosyne
#190	# ============================================================
#191
#192	def benchmark_mnemosyne(conversations: List[dict], top_k: int = DEFAULT_TOP_K) -> dict:
#193	"""Benchmark Mnemosyne ingestion + retrieval on BEAM data."""
#194	results = {
#195	"name": "Mnemosyne (FTS5+vec)",
#196	"ingest_time_s": 0,
#197	"db_size_mb": 0,
#198	"latencies_ms": [],
#199	"recalls": [],
#200	"total_queries": 0,
#201	}
#202
#203	with tempfile.TemporaryDirectory() as tmpdir:
#204	db_path = Path(tmpdir) / "mnemosyne_bench.db"
#205	init_beam(db_path)
#206	beam = BeamMemory(session_id="bench", db_path=db_path)
#207
#208	# Ingest all conversations
#209	t0 = time.perf_counter()
#210	total_msgs = 0
#211	for conv in conversations:
#212	for msg_content in conv["messages"]:
#213	if msg_content.strip():
#214	beam.remember(msg_content, source="bench", importance=0.5)
#215	total_msgs += 1
#216
#217	# Consolidate some to episodic
#218	try:
#219	cursor = beam.conn.cursor()
#220	cursor.execute("SELECT id, content FROM working_memory WHERE session_id = ? ORDER BY timestamp ASC LIMIT 200", (beam.session_id,))
#221	rows = cursor.fetchall()
#222	if rows:
#223	ids = [r["id"] for r in rows]
#224	summary = " \| ".join(r["content"][:80] for r in rows[:5])
#225	beam.consolidate_to_episodic(summary=summary[:500], source_wm_ids=ids, source="bench_consolidation", importance=0.4, scope="global")
#226	beam.conn.commit()
#227	except Exception:
#228	pass
#229
#230	results["ingest_time_s"] = time.perf_counter() - t0
#231	results["db_size_mb"] = os.path.getsize(db_path) / (1024 * 1024)
#232
#233	# Query
#234	for conv in conversations:
#235	for q in conv["questions"]:
#236	question = q["question"]
#237	gold_indices = set(q.get("gold_message_indices", []))
#238
#239	if not question or not gold_indices:
#240	continue
#241
#242	t0 = time.perf_counter()
#243	try:
#244	memories = beam.recall(question, top_k=top_k)
#245	except Exception:
#246	memories = []
#247	latency = (time.perf_counter() - t0) * 1000
#248
#249	# Check recall: how many gold messages were found?
#250	# We match by content substring since we don't have per-message IDs
#251	hits = 0
#252	for mem in memories[:top_k]:
#253	mem_content = mem.get("content", "")
#254	# Fuzzy match: if any gold message content appears in the memory content
#255	for gidx in gold_indices:
#256	if gidx < len(conv["messages"]):
#257	gold_content = conv["messages"][gidx][:100]
#258	if gold_content and gold_content[:50] in mem_content:
#259	hits += 1
#260	break
#261
#262	recall = hits / len(gold_indices) if gold_indices else 0
#263
#264	results["latencies_ms"].append(latency)
#265	results["recalls"].append(recall)
#266	results["total_queries"] += 1
#267
#268	beam.conn.close()
#269
#270	return results
#271
#272
#273	# ============================================================
#274	# Benchmark: FAISS
#275	# ============================================================
#276
#277	def benchmark_faiss(conversations: List[dict], top_k: int = DEFAULT_TOP_K) -> dict:
#278	"""Benchmark FAISS flat L2 index on BEAM data."""
#279	if not HAS_FAISS:
#280	return {"name": "FAISS", "error": "not installed"}
#281
#282	results = {
#283	"name": "FAISS (flat L2)",
#284	"ingest_time_s": 0,
#285	"db_size_mb": 0,
#286	"latencies_ms": [],
#287	"recalls": [],
#288	"total_queries": 0,
#289	}
#290
#291	embedder = get_embedder()
#292
#293	# Build message index
#294	all_messages = []
#295	t0 = time.perf_counter()
#296
#297	for conv in conversations:
#298	all_messages.extend(msg for msg in conv["messages"] if msg.strip())
#299
#300	# Embed all messages
#301	embeddings_list = embedder.embed(all_messages)
#302	if embeddings_list is None:
#303	return {"name": "FAISS (flat L2)", "error": "embedding failed"}
#304	embeddings = np.array(embeddings_list) if not isinstance(embeddings_list, np.ndarray) else embeddings_list
#305
#306	# Build FAISS index
#307	dim = embeddings.shape[1]
#308	index = faiss.IndexFlatL2(dim)
#309	index.add(embeddings.astype(np.float32))
#310
#311	results["ingest_time_s"] = time.perf_counter() - t0
#312
#313	# Estimate size
#314	results["db_size_mb"] = (embeddings.nbytes + sum(len(m.encode()) for m in all_messages)) / (1024 * 1024)
#315
#316	# Query
#317	msg_offset = 0
#318	for conv in conversations:
#319	for q in conv["questions"]:
#320	question = q["question"]
#321	gold_indices = set(q.get("gold_message_indices", []))
#322
#323	if not question:
#324	continue
#325
#326	t0 = time.perf_counter()
#327	q_emb = embedder.embed_query(question)
#328	if q_emb is None:
#329	continue
#330	distances, indices = index.search(q_emb.reshape(1, -1).astype(np.float32), top_k)
#331	latency = (time.perf_counter() - t0) * 1000
#332
#333	# Check recall
#334	retrieved_global_indices = indices[0].tolist()
#335	retrieved_local = [i - msg_offset for i in retrieved_global_indices if msg_offset <= i < msg_offset + len(conv["messages"])]
#336
#337	hits = len(set(retrieved_local) & gold_indices)
#338	recall = hits / len(gold_indices) if gold_indices else 0
#339
#340	results["latencies_ms"].append(latency)
#341	results["recalls"].append(recall)
#342	results["total_queries"] += 1
#343
#344	msg_offset += len(conv["messages"])
#345
#346	return results
#347
#348
#349	# ============================================================
#350	# Benchmark: ChromaDB
#351	# ============================================================
#352
#353	def benchmark_chroma(conversations: List[dict], top_k: int = DEFAULT_TOP_K) -> dict:
#354	"""Benchmark ChromaDB on BEAM data."""
#355	if not HAS_CHROMA:
#356	return {"name": "ChromaDB", "error": "not installed"}
#357
#358	results = {
#359	"name": "ChromaDB (HNSW)",
#360	"ingest_time_s": 0,
#361	"db_size_mb": 0,
#362	"latencies_ms": [],
#363	"recalls": [],
#364	"total_queries": 0,
#365	}
#366
#367	embedder = get_embedder()
#368
#369	with tempfile.TemporaryDirectory() as tmpdir:
#370	client = chromadb.Client(Settings(
#371	chroma_db_impl="duckdb+parquet",
#372	persist_directory=tmpdir,
#373	anonymized_telemetry=False,
#374	))
#375
#376	collection = client.create_collection(name="beam_bench")
#377
#378	# Ingest
#379	t0 = time.perf_counter()
#380	msg_offset = 0
#381
#382	for conv in conversations:
#383	messages = [m for m in conv["messages"] if m.strip()]
#384	if not messages:
#385	continue
#386
#387	ids = [f"msg_{msg_offset + i}" for i in range(len(messages))]
#388	embeddings_list = embedder.embed(messages)
#389	if embeddings_list is None:
#390	continue
#391	embeddings = np.array(embeddings_list) if not isinstance(embeddings_list, np.ndarray) else embeddings_list
#392
#393	collection.add(
#394	embeddings=embeddings.tolist(),
#395	documents=messages,
#396	ids=ids,
#397	)
#398	msg_offset += len(messages)
#399
#400	results["ingest_time_s"] = time.perf_counter() - t0
#401
#402	# Check DB size
#403	db_size = sum(
#404	os.path.getsize(os.path.join(root, f))
#405	for root, _, files in os.walk(tmpdir)
#406	for f in files
#407	)
#408	results["db_size_mb"] = db_size / (1024 * 1024)
#409
#410	# Query
#411	msg_offset = 0
#412	for conv in conversations:
#413	for q in conv["questions"]:
#414	question = q["question"]
#415	gold_indices = set(q.get("gold_message_indices", []))
#416
#417	if not question:
#418	continue
#419
#420	t0 = time.perf_counter()
#421	q_emb = embedder.embed_query(question)
#422	if q_emb is None:
#423	continue
#424	chroma_results = collection.query(
#425	query_embeddings=[q_emb.tolist()],
#426	n_results=top_k,
#427	)
#428	latency = (time.perf_counter() - t0) * 1000
#429
#430	# Check recall
#431	retrieved_ids = chroma_results["ids"][0]
#432	retrieved_local = [
#433	int(rid.split("_")[1]) - msg_offset
#434	for rid in retrieved_ids
#435	if rid.startswith("msg_")
#436	]
#437
#438	hits = len(set(retrieved_local) & gold_indices)
#439	recall = hits / len(gold_indices) if gold_indices else 0
#440
#441	results["latencies_ms"].append(latency)
#442	results["recalls"].append(recall)
#443	results["total_queries"] += 1
#444
#445	msg_offset += len(conv["messages"])
#446
#447	return results
#448
#449
#450	# ============================================================
#451	# Report
#452	# ============================================================
#453
#454	def print_report(all_results: Dict[str, Dict[str, dict]]):
#455	"""Print comparison report."""
#456	print("\n" + "=" * 100)
#457	print(" MNEMOSYNE vs VECTOR DATABASES — BEAM Retrieval Benchmark")
#458	print("=" * 100)
#459
#460	for scale, systems in all_results.items():
#461	print(f"\n --- Scale: {scale} ---")
#462	print(f" {'System':<30} {'Ingest':>10} {'DB Size':>10} {'Avg Lat':>10} {'P50 Lat':>10} {'P95 Lat':>10} {'Recall':>10}")
#463	print(f" {'-'30} {'-'10} {'-'10} {'-'10} {'-'10} {'-'10} {'-'*10}")
#464
#465	for sys_name, r in systems.items():
#466	if "error" in r:
#467	print(f" {sys_name:<30} {'SKIPPED':>10} ({r['error']})")
#468	continue
#469
#470	ingest = f"{r['ingest_time_s']:.1f}s"
#471	db = f"{r['db_size_mb']:.1f}MB"
#472
#473	if r["latencies_ms"]:
#474	avg_lat = f"{statistics.mean(r['latencies_ms']):.1f}ms"
#475	p50 = f"{statistics.median(r['latencies_ms']):.1f}ms"
#476	p95 = f"{sorted(r['latencies_ms'])[int(len(r['latencies_ms'])*0.95)]:.1f}ms" if len(r['latencies_ms']) > 1 else "N/A"
#477	else:
#478	avg_lat = p50 = p95 = "N/A"
#479
#480	if r["recalls"]:
#481	avg_recall = f"{statistics.mean(r['recalls'])*100:.1f}%"
#482	else:
#483	avg_recall = "N/A"
#484
#485	print(f" {sys_name:<30} {ingest:>10} {db:>10} {avg_lat:>10} {p50:>10} {p95:>10} {avg_recall:>10}")
#486
#487	print("\n" + "=" * 100)
#488
#489
#490	def main():
#491	parser = argparse.ArgumentParser()
#492	parser.add_argument("--scales", default="100K", help="Comma-separated scales")
#493	parser.add_argument("--convs", type=int, default=3, help="Conversations per scale")
#494	args = parser.parse_args()
#495
#496	scales = [s.strip() for s in args.scales.split(",")]
#497
#498	print(f"\n{'='*100}")
#499	print(f" Mnemosyne vs Vector DBs Head-to-Head")
#500	print(f" Scales: {scales} \| Conversations/scale: {args.convs}")
#501	print(f"{'='*100}")
#502
#503	# Load data
#504	print("\n[1/3] Loading BEAM data...")
#505	data = load_beam_data(scales, max_convs=args.convs)
#506
#507	all_results = {}
#508
#509	for scale in scales:
#510	if scale not in data:
#511	continue
#512
#513	conversations = data[scale]
#514	all_results[scale] = {}
#515
#516	# Benchmark each system
#517	print(f"\n[2/3] Benchmarking {scale}...")
#518
#519	print(" Mnemosyne...")
#520	all_results[scale]["Mnemosyne"] = benchmark_mnemosyne(conversations)
#521
#522	if HAS_FAISS:
#523	print(" FAISS...")
#524	all_results[scale]["FAISS"] = benchmark_faiss(conversations)
#525
#526	print(f"\n[3/3] Report:")
#527	print_report(all_results)
#528
#529	# Save results
#530	out_path = PROJECT_ROOT / "results" / "vectordb_comparison.json"
#531	os.makedirs(out_path.parent, exist_ok=True)
#532
#533	# Convert to serializable format
#534	serializable = {}
#535	for scale, systems in all_results.items():
#536	serializable[scale] = {}
#537	for sys_name, r in systems.items():
#538	sr = dict(r)
#539	if "latencies_ms" in sr:
#540	sr["avg_latency_ms"] = statistics.mean(sr["latencies_ms"]) if sr["latencies_ms"] else 0
#541	sr["p50_latency_ms"] = statistics.median(sr["latencies_ms"]) if sr["latencies_ms"] else 0
#542	sr["p95_latency_ms"] = sorted(sr["latencies_ms"])[int(len(sr["latencies_ms"])*0.95)] if len(sr["latencies_ms"]) > 1 else 0
#543	if "recalls" in sr:
#544	sr["avg_recall"] = statistics.mean(sr["recalls"]) if sr["recalls"] else 0
#545	sr.pop("latencies_ms", None)
#546	sr.pop("recalls", None)
#547	serializable[scale][sys_name] = sr
#548
#549	with open(out_path, "w") as f:
#550	json.dump(serializable, f, indent=2)
#551
#552	print(f"\n Results saved to: {out_path}")
#553
#554
#555	if __name__ == "__main__":
#556	main()
#557

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public