mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-14 04:12:05 +00:00
135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
"""Derived BM25 length stays consistent across mutations and old snapshots."""
|
|
|
|
# pylint: disable=protected-access,missing-function-docstring
|
|
|
|
import random
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from reme.components.keyword_index import BM25Index
|
|
from reme.components.tokenizer import RegexTokenizer
|
|
|
|
|
|
class ScanningIndex(BM25Index):
|
|
"""Original implementation used as a scoring oracle."""
|
|
|
|
@property
|
|
def total_len(self):
|
|
return int(self._doc_lens[~self._deleted].sum())
|
|
|
|
|
|
def assert_length(index):
|
|
assert index.total_len == int(index._doc_lens[~index._deleted].sum())
|
|
|
|
|
|
def make_index(cls=BM25Index):
|
|
index = cls()
|
|
index.tokenizer = RegexTokenizer(filter_stopwords=False)
|
|
return index
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mutations_and_scores_match_scanning_oracle(tmp_path, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
index, reference = make_index(), make_index(ScanningIndex)
|
|
rng = random.Random(417)
|
|
for iteration in range(80):
|
|
doc_id = f"doc-{rng.randrange(12)}"
|
|
if iteration % 4 == 0:
|
|
for target in (index, reference):
|
|
await target.delete_docs([doc_id, doc_id, "absent"])
|
|
else:
|
|
text = " ".join(rng.choices(["alpha", "beta", "gamma"], k=rng.randrange(8)))
|
|
for target in (index, reference):
|
|
await target.add_docs({doc_id: text})
|
|
if iteration % 9 == 0:
|
|
for target in (index, reference):
|
|
await target.optimize_index()
|
|
assert_length(index)
|
|
assert index.total_len == reference.total_len
|
|
assert index.avg_len == reference.avg_len
|
|
assert await index.retrieve("alpha beta", 6) == await reference.retrieve("alpha beta", 6)
|
|
selected = [f"doc-{i}" for i in range(5)]
|
|
assert await index.retrieve_filtered("gamma beta", 3, selected) == await reference.retrieve_filtered(
|
|
"gamma beta",
|
|
3,
|
|
selected,
|
|
)
|
|
await index.clear()
|
|
assert index.total_len == 0
|
|
await index.add_docs({"new": "alpha beta"})
|
|
await index.delete_docs(["new"])
|
|
await index.optimize_index()
|
|
assert index.total_len == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_snapshot_roundtrip_rebuilds_length_without_new_fields(tmp_path, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
index = make_index()
|
|
await index.add_docs({"old": "alpha beta", "live": "gamma beta alpha"})
|
|
await index.delete_docs(["old"])
|
|
snapshot = index._snapshot()
|
|
assert "total_len" not in snapshot and "_total_len" not in snapshot
|
|
restored = make_index()
|
|
await restored.add_docs({"unrelated": "word " * 100})
|
|
restored._restore(snapshot)
|
|
assert restored.total_len == 3
|
|
assert await restored.retrieve("beta") == await index.retrieve("beta")
|
|
await index.dump()
|
|
loaded = make_index()
|
|
await loaded.load()
|
|
assert loaded.total_len == 3
|
|
await loaded.add_docs({"live": "alpha"})
|
|
assert loaded.total_len == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_batch_counts_only_published_arrays(tmp_path, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
index = make_index()
|
|
await index.add_docs({"old": "alpha beta", "keep": "gamma"})
|
|
tokenize = index._tokenize
|
|
|
|
def fail(text):
|
|
if text == "fail":
|
|
raise ValueError("tokenizer failed")
|
|
return tokenize(text)
|
|
|
|
with patch.object(index, "_tokenize", side_effect=fail):
|
|
with pytest.raises(ValueError, match="tokenizer failed"):
|
|
await index.add_docs({"pending": "alpha alpha", "old": "fail"})
|
|
# Existing partial-batch behavior retires old but does not append pending.
|
|
assert_length(index)
|
|
assert index.total_len == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_postings_failure_still_counts_appended_arrays(tmp_path, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
index = make_index()
|
|
with patch.object(index, "_extend_postings", side_effect=RuntimeError("posting failed")):
|
|
with pytest.raises(RuntimeError, match="posting failed"):
|
|
await index.add_docs({"new": "alpha beta"})
|
|
assert_length(index)
|
|
assert index.total_len == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_never_scans_lengths_with_live_mask():
|
|
class NoBooleanScan(np.ndarray):
|
|
"""Integer posting lookups are allowed, corpus-wide mask scans are not."""
|
|
|
|
def __getitem__(self, key):
|
|
if isinstance(key, np.ndarray) and key.dtype == bool:
|
|
raise AssertionError("query scanned all document lengths")
|
|
return super().__getitem__(key)
|
|
|
|
index = make_index()
|
|
await index.add_docs({"one": "alpha beta", "two": "beta beta"})
|
|
index._doc_lens = index._doc_lens.view(NoBooleanScan)
|
|
assert index.total_len == 4
|
|
assert await index.retrieve("beta")
|
|
assert await index.retrieve_filtered("beta", 1, ["two"])
|