melbinjp/DocQA

1"""The splitter's job is to not cut a fact in half.2 3The old 500-character window did exactly that, and it cost correct answers:4Table 3 of *Attention Is All You Need* puts its column headers and its row5values far enough apart that no 500-character chunk held both, so the API said6it had no information about the big model's parameter count while the number7sat in the document.8"""9 10from utils.splitter import (11    DEFAULT_MAX_CHARS,12    DEFAULT_OVERLAP,13    split_pages,14    split_text,15)16 17 18def test_chunks_are_large_enough_to_hold_a_table():19    """The regression that started this. 500 was too small; guard the floor."""20    assert DEFAULT_MAX_CHARS >= 100021    assert DEFAULT_OVERLAP < DEFAULT_MAX_CHARS22 23 24def test_a_short_document_is_one_chunk():25    assert split_text("Short enough to stay whole.") == ["Short enough to stay whole."]26 27 28def test_empty_and_whitespace_produce_nothing():29    assert split_text("") == []30    assert split_text("   \n\n  \t ") == []31 32 33def test_chunks_respect_the_hard_limit():34    text = " ".join(f"word{i}" for i in range(4000))35    for chunk in split_text(text, max_chars=300, overlap=50):36        assert len(chunk) <= 30037 38 39def test_chunks_start_on_a_word_not_mid_token():40    text = " ".join(f"token{i}" for i in range(2000))41    chunks = split_text(text, max_chars=200, overlap=40)42    for chunk in chunks[1:]:43        assert chunk.startswith("token"), f"chunk begins mid-word: {chunk[:30]!r}"44 45 46def test_sentences_are_preferred_as_boundaries():47    sentence = "The quick brown fox jumped over the lazy dog. "48    chunks = split_text(sentence * 20, max_chars=300, overlap=50)49    # Every chunk but the last should end where a sentence ended.50    for chunk in chunks[:-1]:51        assert chunk.endswith("."), f"cut mid-sentence: {chunk[-40:]!r}"52 53 54def test_a_fact_split_by_a_boundary_survives_whole_somewhere():55    """What overlap is actually for."""56    filler = "padding text here. " * 4057    fact = "The big model has 213 million parameters and the base model has 65."58    text = filler + fact + " " + filler59    chunks = split_text(text, max_chars=400, overlap=120)60    assert any(fact in c for c in chunks), "the fact was cut and appears in no chunk whole"61 62 63def test_no_unbounded_growth_on_text_with_no_breaks():64    """One long run of characters must still terminate, and must still advance."""65    chunks = split_text("x" * 5000, max_chars=400, overlap=100)66    assert len(chunks) < 4067    assert "".join(chunks).count("x") >= 500068 69 70def test_paragraph_breaks_are_kept_for_use_as_boundaries():71    """The old splitter flattened every newline, throwing away the best break."""72    chunks = split_text("First para.\n\nSecond para.", max_chars=DEFAULT_MAX_CHARS)73    assert "\n\n" in chunks[0]74 75 76def test_pages_do_not_bleed_into_each_other():77    pages = [(1, "Page one text. " * 30), (2, "Page two text. " * 30)]78    out = split_pages(pages, max_chars=200, overlap=40)79    for item in out:80        if "one" in item["text"]:81            assert item["page"] == 182        if "two" in item["text"]:83            assert item["page"] == 284    assert {i["page"] for i in out} == {1, 2}85 86 87def test_overlap_at_or_above_max_is_rejected_rather_than_looping():88    import pytest89 90    with pytest.raises(ValueError):91        split_text("some text", max_chars=100, overlap=100)92 93 94# --- tables are one fact, and must not be cut into numbers with no names ---95 96TABLE = (97    "|model|layers|params|\n"98    "|---|---|---|\n"99    "|base|6|65|\n"100    "|big|6|213|\n"101)102 103 104def test_a_small_table_stays_in_one_chunk():105    """The regression: `params` and `213` ended up in different chunks, and the106    answer became 'the provided text does not contain information'."""107    out = split_pages([(9, "Some prose about results.\n" + TABLE)])108    table_chunks = [c["text"] for c in out if "|---|" in c["text"] or c["text"].lstrip().startswith("|")]109    assert len(table_chunks) == 1110    whole = table_chunks[0]111    assert "params" in whole and "213" in whole and "65" in whole112 113 114def test_a_table_too_big_for_one_chunk_repeats_its_header():115    rows = "".join(f"|row{i}|6|{i}|\n" for i in range(200))116    out = split_pages([(1, "|model|layers|params|\n|---|---|---|\n" + rows)], max_chars=400)117    table_chunks = [c["text"] for c in out if "|---|" in c["text"] or c["text"].lstrip().startswith("|")]118    assert len(table_chunks) > 1, "this table should have needed splitting"119    for chunk in table_chunks:120        assert "params" in chunk, "a row block lost its header, so its numbers are unnamed"121 122 123def test_prose_around_a_table_is_still_chunked_normally():124    prose = "This is a sentence about the results. " * 60125    out = split_pages([(3, prose + "\n" + TABLE + "\n" + prose)], max_chars=500)126    assert len([c for c in out if not "|---|" in c["text"] or c["text"].lstrip().startswith("|")]) > 1127    assert len([c for c in out if "|---|" in c["text"] or c["text"].lstrip().startswith("|")]) == 1128    assert {c["page"] for c in out} == {3}129 130 131def test_a_page_that_is_only_a_table_still_produces_a_chunk():132    out = split_pages([(2, TABLE)])133    assert len(out) == 1 and "213" in out[0]["text"] and out[0]["page"] == 2134 135 136def test_a_table_caption_is_not_severed_from_its_grid():137    """The caption is the only natural language a grid has.138 139    Measured on page 9 of Attention Is All You Need: with the caption filed as140    prose, the chunk holding `base 65` and `big 213` began `|Col1|train<br>N d d141    h` and was never retrieved at all for "how does the parameter count of the142    big model compare to the base model". The caption says "Unlisted values are143    identical to those of the base model", which is the language the question is144    asked in.145    """146    page = (147        "Some earlier prose that belongs to the page.\n"148        "\n"149        "Table 3: Variations on the architecture, against the base model.\n"150        "Columns: model, layers, params.\n"151        "|model|layers|params|\n"152        "|---|---|---|\n"153        "|base|6|65|\n"154        "|big|6|213|\n"155    )156    out = split_pages([(9, page)])157    holding = [c["text"] for c in out if "213" in c["text"]]158    assert len(holding) == 1159    chunk = holding[0]160    assert "Table 3" in chunk, "the caption was cut away from the grid"161    assert "base model" in chunk162    assert "params" in chunk and "65" in chunk163 164 165def test_a_long_paragraph_is_not_swallowed_as_a_caption():166    """Only the lines directly above a grid attach, and only a few of them."""167    prose = "\n".join(f"Sentence number {i} of ordinary prose." for i in range(20))168    out = split_pages([(1, prose + "\n|a|b|\n|---|---|\n|1|2|\n")], max_chars=2000)169    table = [c["text"] for c in out if "|---|" in c["text"]][0]170    assert "Sentence number 0" not in table, "the whole paragraph was pulled into the table"171 172 173def test_a_table_is_indexed_by_its_caption_not_its_grid():174    """A vector built from a grid is mostly numbers and stops being findable.175 176    Measured 2026-09-01: with the caption inside the chunk but the whole chunk177    embedded, the Table 3 chunk was still not retrieved for a question its own178    caption answers, while neighbouring chunks scored 0.36 to 0.47. The caption179    is what a question looks like; the grid is what the answer is in.180    """181    page = (182        "Preceding prose.\n"183        "\n"184        "Table 3: Variations on the architecture, against the base model.\n"185        "Columns: model, layers, params.\n"186        "|model|layers|params|\n"187        "|---|---|---|\n"188        "|big|6|213|\n"189    )190    table = [c for c in split_pages([(9, page)]) if "213" in c["text"]][0]191    assert "embed_text" in table, "the table would be indexed by its grid"192    assert "Table 3" in table["embed_text"] and "base model" in table["embed_text"]193    assert "|213|" not in table["embed_text"], "the grid leaked into what gets matched"194    # What is returned and cited is still the whole thing.195    assert "213" in table["text"] and "|---|" in table["text"]196 197 198def test_prose_chunks_carry_no_embed_text():199    out = split_pages([(1, "Just ordinary prose, no grid anywhere in it.")])200    assert all("embed_text" not in c for c in out)201