File size: 665 Bytes
1817730 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | from sovlm.dictionary import Dictionary
from sovlm.inverted_index import InvertedIndex
from sovlm.corpus import tokenize
def test_bm25_retrieval():
d = Dictionary()
inv = InvertedIndex()
chunks = [
"the cat sat on the mat",
"the dog ran in the park",
"quantum mechanics governs atomic transitions",
]
for i, text in enumerate(chunks):
ids = [d.add(tok) for tok in tokenize(text)]
inv.add_chunk(i, text, ids)
d.freeze()
query = [d["the"], d["cat"]]
results = inv.retrieve(query, top_k=2)
assert len(results) > 0
assert results[0][1] == 0 # chunk 0 has "cat"
|