Spaces:
Running on Zero
Running on Zero
File size: 1,249 Bytes
6a46e44 | 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 | """
tests/test_bm25.py
──────────────────
Unit tests for the BM25 keyword retriever.
"""
from __future__ import annotations
from llama_index.core.schema import TextNode
from docling_pdf_processor.retrieval.bm25 import BM25, KeywordSearchRetriever
def test_bm25_tokenize():
bm25 = BM25()
assert bm25.tokenize("Hello, world!") == ["hello", "world"]
def test_bm25_search_ranking():
bm25 = BM25()
docs = [
"the quick brown fox",
"the lazy dog sleeps",
"foxes are quick and brown",
]
bm25.fit(docs)
results = bm25.search("quick brown fox", top_k=3)
assert len(results) == 3
# The most relevant doc should be the one containing quick + brown + fox
top_idx = results[0][0]
assert top_idx in (0, 2)
def test_keyword_search_retriever():
nodes = [
TextNode(text="Python is great for data science."),
TextNode(text="JavaScript runs in the browser."),
TextNode(text="Python and machine learning go hand in hand."),
]
retriever = KeywordSearchRetriever(nodes)
results = retriever.retrieve("machine learning", top_k=2)
assert len(results) >= 1
assert "machine learning" in results[0].node.text.lower()
|