Spaces:
Runtime error
Runtime error
File size: 2,292 Bytes
64a1c56 | 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 | from app.bm25 import BM25Index
from app.chunking import chunk_text
from app.query import extract_pages, match_docs
def test_chunks_respect_size_and_overlap():
text = " ".join(f"This is sentence number {i} about topic {i % 5}." for i in range(80))
chunks = chunk_text(text, size=300, overlap=80)
assert len(chunks) > 5
assert all(len(c) <= 300 + 80 for c in chunks)
# overlap: the end of one chunk reappears at the start of the next
assert chunks[0].split(". ")[-1] in chunks[1]
def test_chunk_handles_giant_unbroken_text():
chunks = chunk_text("word " * 2000, size=500, overlap=50)
assert chunks and all(len(c) <= 560 for c in chunks)
def test_extract_pages_variants():
assert extract_pages("what is written on page 20?") == [20]
assert extract_pages("page no. 7 pe kya hai") == [7]
assert extract_pages("20 no ke page pe kya likha hai") == [20]
assert extract_pages("20th page summary") == [20]
assert extract_pages("explain pages 20-23") == [20, 21, 22, 23]
assert extract_pages("page 5 to 7") == [5, 6, 7]
assert extract_pages("20 se 22 pages") == [20, 21, 22]
assert extract_pages("page 3, 9 and 12") == [3, 9, 12]
assert extract_pages("what is the refund policy?") == []
assert extract_pages("who was born in 1990?") == []
def test_match_docs_by_filename():
docs = [
{"doc_id": "a", "filename": "HR_Policy.pdf"},
{"doc_id": "b", "filename": "annual-report-2024.pdf"},
]
assert match_docs("what does the hr policy say about leave", docs) == ["a"]
assert match_docs("summary of annual report 2024", docs) == ["b"]
assert match_docs("what about leave", docs) == []
def test_bm25_scopes_by_doc_and_page():
idx = BM25Index()
idx.build(
[
{"id": "1", "doc_id": "A", "page": 1, "text": "zebra migration across the plains"},
{"id": "2", "doc_id": "B", "page": 2, "text": "zebra stripes and lions"},
{"id": "3", "doc_id": "B", "page": 3, "text": "completely unrelated tax filing"},
]
)
idx._dirty = False
assert {h["id"] for h in idx.search("zebra", limit=5)} == {"1", "2"}
assert [h["id"] for h in idx.search("zebra", doc_ids=["B"])] == ["2"]
assert idx.search("zebra", pages=[1])[0]["id"] == "1"
|