File size: 2,062 Bytes
dbb04e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import time
import uuid
import numpy as np
import sys
import os
# Add src to path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))
from core.engine import HAIMEngine
from core.node import MemoryNode
from core.binary_hdv import BinaryHDV
from core.engine import HAIMEngine
from core.node import MemoryNode
from core.binary_hdv import BinaryHDV
from core.config import get_config, reset_config
import os
import asyncio
async def benchmark_query():
print("Initializing Engine...")
os.environ["HAIM_ENCODING_MODE"] = "binary"
reset_config()
engine = HAIMEngine()
await engine.initialize()
# Mock data generation
count = 2000
print(f"Generating {count} dummy memories...")
engine.tier_manager.hot.clear()
dim = engine.dimension
start_gen = time.time()
for i in range(count):
# Create random binary vector
data = np.random.randint(0, 2, size=(dim // 8), dtype=np.uint8)
hdv = BinaryHDV(data=data, dimension=dim)
node = MemoryNode(
id=f"mem_{i}",
hdv=hdv,
content=f"Dummy content {i}",
metadata={}
)
node.tier = "hot"
engine.tier_manager.hot[node.id] = node
# Add to FAISS too
engine.tier_manager._add_to_faiss(node)
print(f"Generation took {time.time() - start_gen:.4f}s")
# Benchmark Query
print("Benchmarking Query...")
query_text = "test query"
# Warmup
await engine.query(query_text, top_k=5)
start_time = time.time()
iterations = 50
for _ in range(iterations):
await engine.query(query_text, top_k=5)
total_time = time.time() - start_time
avg_time = total_time / iterations
print(f"Average Query Time over {iterations} runs: {avg_time*1000:.2f} ms")
print(f"Total Time: {total_time:.4f}s")
if __name__ == "__main__":
asyncio.run(benchmark_query())
|