Spaces:
Sleeping
Sleeping
File size: 4,477 Bytes
c42f4ef | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Fast tests: no neural models except the bundled OCR."""
from pathlib import Path
from mrag.config import Config
from mrag.generate import ImageHit, TextHit, build_prompt
from mrag.ingest import chunk_text, iter_files, load_any
from mrag.pipeline import rrf
def test_config_env_and_overrides(monkeypatch):
monkeypatch.setenv("MRAG_TOP_K_TEXT", "9")
monkeypatch.setenv("MRAG_LLM_VISION", "false")
monkeypatch.setenv("MRAG_CLIP_MIN_SCORE", "0.3")
c = Config.load(llm_model="gemma3:4b")
assert c.top_k_text == 9 and c.llm_vision is False and c.clip_min_score == 0.3 and c.llm_model == "gemma3:4b"
def test_chunking_covers_all_words_with_overlap():
words = [f"w{i}" for i in range(1000)]
chunks = chunk_text(" ".join(words), chunk_words=100, overlap=20)
assert all(len(c.split()) <= 100 for c in chunks)
seen = {w for c in chunks for w in c.split()}
assert seen == set(words)
assert chunks[0].split()[-20:] == chunks[1].split()[:20]
def test_chunking_paragraphs():
text = "\n\n".join(" ".join(["p%d" % i] * 30) for i in range(10))
chunks = chunk_text(text, chunk_words=100, overlap=0)
assert len(chunks) == 4 and sum(len(c.split()) for c in chunks) == 300
def test_chunking_overlap_with_a_full_paragraph():
text = " ".join(["a"] * 170) + "\n\n" + " ".join(["b"] * 180)
chunks = chunk_text(text, chunk_words=180, overlap=40)
assert all(len(c.split()) <= 180 for c in chunks)
assert sum(c.split().count("b") for c in chunks) == 180
def test_rrf_prefers_agreement():
fused = rrf([["a", "b", "c"], ["b", "c", "a"], ["b"]])
assert fused[0][0] == "b"
def test_pdf_loader_text_chart_and_scanned_page(corpus):
ex = load_any(corpus / "report.pdf")
assert [t.page for t in ex.texts] == [1, 2] # page 3 has no text layer
assert "Meera Iyer" in ex.texts[0].text
assert "41.2" not in " ".join(t.text for t in ex.texts) # chart value only exists in the image
pages = sorted((i.page, i.meta.get("scanned_page", False)) for i in ex.images)
assert (2, False) in pages and (3, True) in pages
def test_docx_loader_table_and_image(corpus):
ex = load_any(corpus / "handbook.docx")
kinds = {t.kind for t in ex.texts}
assert kinds == {"text", "table"}
table = next(t.text for t in ex.texts if t.kind == "table")
assert "Electric forklift | 2500 kg" in table
assert len(ex.images) == 1 and ex.images[0].image.size == (600, 400)
def test_iter_files_skips_hidden(tmp_path):
(tmp_path / ".hidden").mkdir()
(tmp_path / ".hidden" / "x.txt").write_text("secret")
(tmp_path / "a.txt").write_text("hello")
(tmp_path / "b.xyz").write_text("nope")
assert [p.name for p in iter_files([tmp_path])] == ["a.txt"]
def test_ocr_reads_scanned_page_and_chart(corpus):
from mrag.models import OCR
ocr = OCR()
ex = load_any(corpus / "report.pdf")
by_page = {i.page: ocr(i.image) for i in ex.images}
assert "36-month" in by_page[3] and "warranty" in by_page[3].lower()
assert "41.2" in by_page[2] and "Asia-Pacific" in by_page[2]
notice = load_any(corpus / "images" / "notice_board.png").images[0].image
assert "Monday" in ocr(notice)
def test_prompt_numbering():
p = build_prompt("q?", [TextHit("1", "alpha", "/x/a.pdf", 2, "text", 0.5),
TextHit("2", "beta", "/x/b.png", None, "ocr", 0.4)],
[ImageHit("i", "/tmp/i.png", "/x/b.png", None, "beta", 0.3)])
assert "[1] (a.pdf, p.2)\nalpha" in p and "[2] (b.png, OCR of an image)" in p
assert "[Img 1] b.png | text in image: beta" in p and p.rstrip().endswith("Answer:")
def test_pdf_loader_keeps_vector_drawn_figures(tmp_path):
"""A plotted line is vector artwork, invisible to page.get_images()."""
import pymupdf
pdf = tmp_path / "vector.pdf"
doc = pymupdf.open()
page = doc.new_page()
page.insert_text((70, 70), "Plot A demonstrates the relationship between pressure and temperature.")
page.draw_line((75, 300), (400, 300), color=(0, 0, 0), width=2)
page.draw_line((75, 300), (75, 120), color=(0, 0, 0), width=2)
page.draw_line((75, 280), (350, 140), color=(0, 0, 1), width=3)
assert page.get_drawings() and not page.get_images(full=True)
doc.save(pdf)
doc.close()
out = load_any(pdf)
assert out.texts and "Plot A demonstrates" in out.texts[0].text
assert any(im.meta.get("page_render") and im.page == 1 for im in out.images)
|