doc eval

Whether the document tree finds better passages than flat chunks, measured on EU legislation

What this measures

litesearch.tree splits each document into a tree of sections (title, chapter, section, article) and records, for every chunk, the section it came from. The argument for building it was that a chunk which knows where it lives should retrieve better. This notebook tests that argument against the litesearch that came before the tree.

The corpus is the one recorded in the outputs below: 11 PDFs and 524 pages, almost all of it EU directives and regulations, plus the Attention paper. Legislation was chosen because it has real hierarchy (TITLE, CHAPTER, SECTION, Article), which is the kind of document the tree was built for.

Four systems share one corpus, one encoder and one query set:

system	what it is
flat-page	litesearch before the tree: `chunk_markdown` on each page, searched with `db.search`
flat-fine	the tree’s own chunk texts in a plain store, with no sections and no headings
tree	`db.add_doc` and `db.doc_search`: chunks linked to sections, embedded with their heading, adjacent hits merged into spans
sections	`db.sections`: the same search rolled up so that each result is a whole section

flat-fine is the control that matters. The tree also chunks more finely than one chunk per page, and without this control it would take credit for an effect that any store gets by using smaller chunks.

The encoder is a weak stand-in. The machine this ran on could not reach huggingface.co, so no ONNX or model2vec model could be downloaded. The vector side of each search uses LSA instead: TF-IDF word weights reduced to 256 dimensions with a truncated SVD, fitted on this corpus. Every system uses it, so the comparisons between systems hold, but absolute scores understate anything that depends on the meaning of the text, heading context most of all.

import glob, math, random, re, time
from collections import Counter
from pathlib import Path
import numpy as np
from litesearch import database
from litesearch.data import pdf_parse, chunk_markdown

PDFS = sorted(glob.glob('../examples/pdfs/*.pdf')) + ['pdfs/attention_is_all_you_need.pdf']
pages_by_doc = {p: list(enumerate(pdf_parse(p))) for p in PDFS}
print(f'{len(pages_by_doc)} documents, {sum(len(v) for v in pages_by_doc.values())} pages')

Dictionary used where Stream expected, treating as empty stream
Dictionary used where Stream expected, treating as empty stream
Dictionary used where Stream expected, treating as empty stream
Dictionary used where Stream expected, treating as empty stream
Dictionary used where Stream expected, treating as empty stream
Dictionary used where Stream expected, treating as empty stream
Dictionary used where Stream expected, treating as empty stream
Dictionary used where Stream expected, treating as empty stream

11 documents, 524 pages

class LSA:
    "TF-IDF + truncated SVD, fitted on the corpus. Stands in for an ONNX model — see the caveat above."
    def __init__(self, dim=256):
        from sklearn.feature_extraction.text import TfidfVectorizer
        from sklearn.decomposition import TruncatedSVD
        from sklearn.pipeline import make_pipeline
        from sklearn.preprocessing import Normalizer
        self.p = make_pipeline(TfidfVectorizer(sublinear_tf=True, min_df=2, max_df=0.5,
                                               stop_words='english', ngram_range=(1,2)),
                               TruncatedSVD(n_components=dim, random_state=0), Normalizer(copy=False))
    def fit(self, t): self.p.fit(list(t)); return self
    def __call__(self, t, **kw): return self.p.transform(list(t)).astype(np.float16)

all_text = [t for pg in pages_by_doc.values() for _, t in pg if t and t.strip()]
enc = LSA(256).fit(all_text)
E = lambda q: enc([q])[0].tobytes()

Build the four stores

# tree: one add_doc per document
tdb = database(); tdb.get_tree('store')
for p, pages in pages_by_doc.items():
    tdb.add_doc(pages, title=Path(p).stem, source=p, kind='pdf', emb_fn=enc)
trows = tdb.t.store(select='rowid as rowid, content, node_id, page')
nodes = {n['id']: n for n in tdb.t.nodes()}

# tree without the heading path folded into the embedding
ndb = database(); ndb.get_tree('store')
for p, pages in pages_by_doc.items():
    ndb.add_doc(pages, title=Path(p).stem, source=p, kind='pdf', emb_fn=enc, with_heading=False)

# flat-page: litesearch before this work
fdb = database(); fst = fdb.get_store('flat', hash=True, ann=True, doc=str, page=int)
fst.sync([dict(content=c, doc=Path(p).name, page=pg)
          for p, pages in pages_by_doc.items() for pg, txt in pages
          for c in chunk_markdown(txt) if len(c.strip()) >= 40], emb_fn=enc)

# flat-fine: the tree's chunks with every structural signal removed
gdb = database(); gst = gdb.get_store('fine', hash=True, ann=True)
gst.sync([dict(content=r['content']) for r in trows], emb_fn=enc)

print(f'flat-page {len(fst()):>5} chunks')
print(f'flat-fine {len(gst()):>5} chunks')
print(f'tree      {len(trows):>5} chunks, {len(nodes)} nodes')

flat-page   512 chunks
flat-fine  1475 chunks
tree       1475 chunks, 2290 nodes

The trees these PDFs produce are worth a look on their own. pdf-oxide marks 16 to 31 lines per page as markdown headings, all at the same depth, which is not structure. detect_mode notices many headings with no variation in depth and switches to the chapter mode, which finds structure from words such as CHAPTER, Article and TITLE in the text and works out the order of levels for each document separately. struct_levels reports that order: for the VAT directive (dir_2006_112) it is title (or part, or annex) at level 1, then chapter, section and article.

from litesearch.tree import detect_mode, struct_levels, build_tree
for p in list(pages_by_doc)[:5]:
    pages = pages_by_doc[p]
    t = build_tree(pages, title=Path(p).stem)
    lv = Counter(n.level for n in t)
    print(f'{Path(p).name[:31]:<31} {detect_mode(pages):8} {len(t):>5} nodes  '
          f'levels={dict(sorted(lv.items()))}  ladder={struct_levels(pages)}')

attention_is_all_you_need.pdf   markdown     5 nodes  levels={0: 1, 2: 1, 4: 3}  ladder={}
dir_1993_13_2022-05-28_eng.pdf  chapter     14 nodes  levels={0: 1, 1: 13}  ladder={'article': 1}
dir_1993_83_2019-06-06_eng.pdf  chapter     20 nodes  levels={0: 1, 1: 4, 2: 15}  ladder={'chapter': 1, 'article': 2}
dir_1996_9_2019-06-06_eng.pdf   chapter     22 nodes  levels={0: 1, 1: 4, 2: 17}  ladder={'chapter': 1, 'article': 2}
dir_2006_112_2022-07-01_eng.pdf chapter   1983 nodes  levels={0: 1, 1: 124, 2: 126, 3: 53, 4: 1679}  ladder={'title': 1, 'chapter': 2, 'section': 3, 'annex': 1, 'part': 1, 'article': 4}

Queries and ground truth

Each query is one sentence of 9 to 26 words copied from an Article, at most one per Article and 150 in all. Each sentence is asked in two forms:

verbatim, the sentence as printed. FTS5 answers this almost trivially, so it only checks that nothing got worse.
degraded, with the rarest 45% of its words removed. This is not a paraphrase, but it removes the exact-match shortcut and leaves something like what a person types when they do not know the wording.

A result counts as a hit when the text a system returns contains the original sentence. That test does not depend on how the system chunked the document.

_ART, _SENT = re.compile(r'^\s*#*\s*Article\s+\d+', re.I), re.compile(r'[^.!?\n]{60,240}[.!?]')
arts = {i: n for i, n in nodes.items() if _ART.match(n['title'] or '')}
def owner(nid):
    seen = set()
    while nid and nid not in seen:
        seen.add(nid)
        if nid in arts: return nid
        nid = (nodes.get(nid) or {}).get('parent_id')
by_node = {}
for r in trows: by_node.setdefault(r['node_id'], []).append(r['content'])

rnd = random.Random(0); aids = list(arts); rnd.shuffle(aids)
qs = []
for a in aids:
    for m in _SENT.finditer(' '.join(by_node.get(a, []))):
        s = ' '.join(m.group(0).split()).strip('. ')
        if 9 <= len(s.split()) <= 26 and not s.startswith(('http','- [','|')): qs.append((s, a)); break
    if len(qs) >= 150: break

df = Counter(w.lower().strip('.,;:()') for t in all_text for w in t.split())
def degrade(q, keep=0.55):
    ws = q.split()
    if len(ws) < 6: return q
    rare = sorted(range(len(ws)), key=lambda i: df.get(ws[i].lower().strip('.,;:()'), 0))
    drop = set(rare[:max(1, int(len(ws)*(1-keep)))])
    return ' '.join(w for i, w in enumerate(ws) if i not in drop)
dqs = [(degrade(q), a, q) for q, a in qs]
print(f'{len(qs)} queries, {len(arts)} Article nodes indexed')
print(f'  verbatim : {qs[0][0][:86]}')
print(f'  degraded : {dqs[0][0][:86]}')

150 queries, 1903 Article nodes indexed
  verbatim : ‘Supply of services’ shall mean any transaction which does not constitute a supply of 
  degraded : of shall any which not a supply of goods

def nrm(s): return ' '.join((s or '').lower().split())
def score_text(runs, k=10):
    r1=r5=rk=0; mrr=0.0
    for q, ts in runs:
        r = next((i for i,t in enumerate(ts[:k]) if nrm(q) in nrm(t)), None)
        if r is None: continue
        rk += 1; mrr += 1/(r+1); r1 += r==0; r5 += r<5
    n = max(1,len(runs)); return dict(r1=r1/n, r5=r5/n, r10=rk/n, mrr=mrr/n)
def score_art(runs, k=5):
    r1=r3=rk=0; mrr=0.0
    for gt, got in runs:
        r = next((i for i,a in enumerate(got[:k]) if a==gt), None)
        if r is None: continue
        rk += 1; mrr += 1/(r+1); r1 += r==0; r3 += r<3
    n = max(1,len(runs)); return dict(r1=r1/n, r3=r3/n, r5=rk/n, mrr=mrr/n)
def show(nm, s): print(f'  {nm:<24} ' + '  '.join(f'{k}={v:.3f}' for k,v in s.items()))

plain = lambda db, tbl, q, k=10: [h['content'] for h in (db.search(q, E(q), columns=['content'],
                        limit=k, table_name=tbl, dtype=np.float16) or [])]
tree_hits = lambda db, q, k=10, **kw: [h['content'] for h in db.doc_search(q, E(q), limit=k, dtype=np.float16, **kw)]
SYS = {'flat-page (before)':   lambda q: plain(fdb,'flat',q),
       'flat-fine (chunking)': lambda q: plain(gdb,'fine',q),
       'tree doc_search':      lambda q: tree_hits(tdb,q),
       'tree -heading ctx':    lambda q: tree_hits(ndb,q),
       'tree -spans':          lambda q: tree_hits(tdb,q,spans=False)}

A. Verbatim queries, a regression check

Every system scores close to 1. The scores are r1, r5 and r10, the share of queries whose sentence appears in the top 1, 5 and 10 results, and mrr, the mean reciprocal rank (1 for a hit at rank 1, 0.5 at rank 2, and so on, averaged over queries). An exact sentence is what FTS5 is for, so this task cannot separate retrieval strategies. It shows that the tree costs nothing on the easy case.

for nm, fn in SYS.items(): show(nm, score_text([(q, fn(q)) for q,_ in qs]))

  flat-page (before)       r1=0.940  r5=0.973  r10=0.973  mrr=0.956
  flat-fine (chunking)     r1=0.947  r5=0.973  r10=0.973  mrr=0.960
  tree doc_search          r1=0.967  r5=0.993  r10=0.993  mrr=0.980
  tree -heading ctx        r1=0.960  r5=0.993  r10=0.993  mrr=0.975
  tree -spans              r1=0.947  r5=0.973  r10=0.973  mrr=0.960

B. Degraded queries, where the difference should show up

for nm, fn in SYS.items(): show(nm, score_text([(o, fn(d)) for d,_,o in dqs]))

  flat-page (before)       r1=0.267  r5=0.460  r10=0.553  mrr=0.354
  flat-fine (chunking)     r1=0.460  r5=0.620  r10=0.707  mrr=0.532
  tree doc_search          r1=0.420  r5=0.667  r10=0.740  mrr=0.517
  tree -heading ctx        r1=0.427  r5=0.653  r10=0.733  mrr=0.519
  tree -spans              r1=0.420  r5=0.653  r10=0.713  mrr=0.510

Moving from flat-page to flat-fine accounts for the whole gain. MRR rises from 0.354 to 0.532, half again as high, from smaller chunks alone. Moving from flat-fine to the tree gives 0.517, slightly behind once chunk size is held equal. Removing heading context gives 0.519, no change.

On this corpus, the retrieval gain credited to document structure is a chunking gain. The tree is what produces the better chunks, because a section is narrower than a page, but heading context and span merging each move MRR by less than 0.01.

C. Which Article answers this

This is the question a reader asks of a legal corpus. Each system’s results are mapped to Articles in the same way, so the mapping is not a variable: an exact chunk match first, then any Article whose text overlaps the returned passage. Here r1, r3 and r5 are the share of queries whose source Article is in the top 1, 3 or 5 Articles returned.

c2a, a2text = {}, {}
for r in trows:
    if (a := owner(r['node_id'])):
        c2a[nrm(r['content'])] = a
        a2text.setdefault(a, []).append(r['content'])
a2text = {a: nrm(' '.join(v)) for a, v in a2text.items()}
a_keys = [(a, t[:200]) for a, t in a2text.items() if len(t) > 60]

def to_arts(texts, k=5):
    out = []
    for t in texts:
        n = nrm(t)
        got = [c2a[n]] if n in c2a else [a for a, key in a_keys if key in n or n[:200] in a2text[a]]
        for a in got:
            if a not in out: out.append(a)
        if len(out) >= k: break
    return out[:k]
def sec_arts(q, k=5, **kw):
    out = []
    for s in tdb.sections(q, E(q), limit=k*2, per=2, dtype=np.float16, **kw):
        a = s['node_id'] if s['node_id'] in arts else (c2a.get(nrm(s['snippets'][0])) if s['snippets'] else None)
        if a and a not in out: out.append(a)
        if len(out) >= k: break
    return out

for label, qset in (('verbatim', [(q,a) for q,a in qs]), ('degraded', [(d,a) for d,a,_ in dqs])):
    print(f'[{label}]')
    for nm, fn in (('flat-page (before)',   lambda q: to_arts(plain(fdb,'flat',q,10))),
                   ('flat-fine (chunking)', lambda q: to_arts(plain(gdb,'fine',q,10))),
                   ('tree doc_search',      lambda q: to_arts(tree_hits(tdb,q,10))),
                   ('sections score=max',   lambda q: sec_arts(q)),
                   ('sections score=sum',   lambda q: sec_arts(q, score='sum'))):
        show(nm, score_art([(a, fn(q)) for q, a in qset]))

[verbatim]
  flat-page (before)       r1=0.267  r3=0.827  r5=0.947  mrr=0.538
  flat-fine (chunking)     r1=0.933  r3=0.993  r5=1.000  mrr=0.961
  tree doc_search          r1=0.933  r3=0.993  r5=1.000  mrr=0.961
  sections score=max       r1=0.927  r3=0.993  r5=1.000  mrr=0.957
  sections score=sum       r1=0.727  r3=0.973  r5=1.000  mrr=0.852
[degraded]
  flat-page (before)       r1=0.087  r3=0.313  r5=0.360  mrr=0.194
  flat-fine (chunking)     r1=0.447  r3=0.580  r5=0.640  mrr=0.519
  tree doc_search          r1=0.407  r3=0.593  r5=0.680  mrr=0.507
  sections score=max       r1=0.387  r3=0.600  r5=0.720  mrr=0.505
  sections score=sum       r1=0.180  r3=0.387  r5=0.553  mrr=0.304

sections() can score a section by the best of its hits (score='max'), their mean, or their sum. Summing favours long sections: a long chapter with five weak hits outranks the short Article that answers the question. With max, a section is as relevant as its best evidence, and section search becomes level with chunk search (verbatim MRR 0.957 against 0.961) while returning a whole section. That is why max is the default.

flat-page shows what page-sized chunks cannot do. On verbatim queries it puts the right text first 94% of the time, but it names the right Article first only 27% of the time, because one page spans several Articles and the chunk cannot say which. That is the case for the tree that these numbers support.

What this changed

finding	action taken
pdf-oxide marks every bold line as a level-1 heading, and the VAT directive came out as 5,242 one-line sections	`detect_mode` now rejects many headings with no depth and builds levels from CHAPTER, Article and TITLE instead: 1,983 sections over 4 levels
summing hit scores in `sections()` favours long sections	the default is now `score='max'`: MRR +0.11 on verbatim queries and +0.20 on degraded ones
`adaptive_weights`, which weights the keyword and vector sides by the shape of the query, never fired on prose, and a replacement based on FTS coverage was worse	`doc_search(adaptive=...)` defaults to off and is documented rather than removed
chunk size outweighs every structural feature	chunk size moved to the top of the work list; later evals made 512-byte chunks the `Index` default
heading context measures as noise with a weak encoder	kept, since it is cheap, and documented as unproven until a real encoder judges it