mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-13 04:01:35 +00:00
106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
"""Compare live-length lookup and real BM25 queries with the previous scan.
|
|
|
|
Run from the repository: python -m scripts.benchmark_bm25_length --repeats 15
|
|
Setup is excluded; no files, models or network are used. JSON goes to stdout.
|
|
"""
|
|
|
|
# pylint: disable=protected-access
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import platform
|
|
import statistics
|
|
import time
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
|
|
from reme.components.keyword_index import BM25Index
|
|
from reme.components.tokenizer import RegexTokenizer
|
|
|
|
|
|
def scanning_length(index):
|
|
"""Original length calculation, including lazy-deleted slots."""
|
|
return int(index._doc_lens[~index._deleted].sum())
|
|
|
|
|
|
async def measure(operation, repeats):
|
|
"""Warm up and report milliseconds per operation."""
|
|
for _ in range(3):
|
|
await operation()
|
|
samples = []
|
|
for _ in range(repeats):
|
|
start = time.perf_counter()
|
|
await operation()
|
|
samples.append((time.perf_counter() - start) * 1000)
|
|
return {"median_ms": statistics.median(samples), "p95_ms": float(np.percentile(samples, 95))}
|
|
|
|
|
|
async def benchmark(size, repeats):
|
|
"""Exercise a rare-term query globally and within 100 selected documents."""
|
|
index = BM25Index()
|
|
index.tokenizer = RegexTokenizer(filter_stopwords=False)
|
|
for start in range(0, size, 10000):
|
|
docs = {}
|
|
for i in range(start, min(start + 10000, size)):
|
|
docs[str(i)] = "alpha beta rare" if i % 1000 == 0 else "alpha beta"
|
|
await index.add_docs(docs)
|
|
selected = [str(i) for i in range(min(size, 100))]
|
|
|
|
async def length():
|
|
return index.total_len
|
|
|
|
async def query():
|
|
return await index.retrieve("rare", 10)
|
|
|
|
async def filtered():
|
|
return await index.retrieve_filtered("rare", 10, selected)
|
|
|
|
results = []
|
|
for deleted_fraction in (0.0, 0.5):
|
|
if deleted_fraction:
|
|
await index.delete_docs([str(i) for i in range(1, size, 2)])
|
|
for name, operation in (("length", length), ("retrieve", query), ("retrieve_filtered", filtered)):
|
|
expected = await operation()
|
|
cached = await measure(operation, repeats)
|
|
with patch.object(BM25Index, "total_len", property(scanning_length)):
|
|
assert await operation() == expected
|
|
baseline = await measure(operation, repeats)
|
|
results.append(
|
|
{
|
|
"slots": size,
|
|
"deleted_fraction": deleted_fraction,
|
|
"operation": name,
|
|
"baseline": baseline,
|
|
"incremental": cached,
|
|
},
|
|
)
|
|
return results
|
|
|
|
|
|
async def main():
|
|
"""Print reproducible environment and benchmark results."""
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--sizes", nargs="+", type=int, default=[10000, 100000, 1000000])
|
|
parser.add_argument("--repeats", type=int, default=15)
|
|
args = parser.parse_args()
|
|
rows = []
|
|
for size in args.sizes:
|
|
rows.extend(await benchmark(size, args.repeats))
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"python": platform.python_version(),
|
|
"platform": platform.platform(),
|
|
"numpy": np.__version__,
|
|
"repeats": args.repeats,
|
|
"results": rows,
|
|
},
|
|
indent=2,
|
|
),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|