Refine the Flash tree with merge

README.md

@@ -205,7 +205,7 @@ python3 run_pageindex.py --md_path /path/to/your/document.md
 > python3 run_pageindex.py --flash --pdf_path /path/to/your/document.pdf
 > ```
 >
-> Add `--optimize` to refine the tree structure for more efficient retrieval (`--optimize merge` skips the LLM expansion pass).
+> Add `--optimize` to refine the tree structure for more efficient retrieval (with an LLM expansion pass).
 
 ## 🚀 Agentic Vectorless RAG: An Example
 

pageindex/flash/README.md

@@ -30,6 +30,7 @@ missing, non-PDF, encrypted, empty, or unreadable file.
             "node_id": str,       # 4-digit, zero-padded
             "start_index": int,
             "end_index": int,
+            "key_items": [str],   # titles of merged-away subsections; absent when none
             "nodes": [...],       # absent on leaf nodes
         }
     ],

pageindex/flash/api.py

@@ -68,9 +68,10 @@ def _validate_pdf(pdf):
     return pdf
 
 
-def _thin(structure):
-    from ..utils import page_level_thinning, write_node_id
-    page_level_thinning(structure)
+def _merge(structure):
+    from ..tree_optimize import merge_tree
+    from ..utils import write_node_id
+    merge_tree(structure)
     write_node_id(structure)
 
 
@@ -82,9 +83,9 @@ async def _summarize(structure, page_list, model, concurrency=None):
 def _optimize(structure, page_texts, do_expand, model):
     """Merge/expand refinement between extraction and summaries.
 
-    Supersedes ``_thin``: merge collapses everything thinning would, but keeps
-    the dropped titles as ``key_items``. Summaries run after, so they describe
-    the final tree. Expand reads the same page text the summaries use.
+    Beyond the merge the default path runs anyway, this adds LLM expand and
+    reports before/after search-cost metrics. Summaries run after, so they
+    describe the final tree. Expand reads the same page text the summaries use.
     """
     import asyncio
     from ..tree_optimize import optimize
@@ -95,21 +96,24 @@ def _optimize(structure, page_texts, do_expand, model):
                                    do_expand=do_expand,
                                    page_count=len(page_texts)))
     return {"merges": outcome["merges"], "expands": outcome["expands"],
+            "same_page_merges": outcome["same_page_merges"],
+            "same_page_dropped": outcome["same_page_dropped"],
+            "kept_collapsed": outcome["kept_collapsed"],
             "before": outcome["before"], "after": outcome["after"]}
 
 
 def page_index_flash(pdf, summary=True, summary_model=None,
                      optimize=False, optimize_expand=True,
                      optimize_model=None, summary_concurrency=None) -> dict:
-    """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: if True, refine the tree for search cost (merge + expand) before summaries. optimize_expand: if False, optimization only performs deterministic merge; summary generation is unchanged. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
+    """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: if True, additionally expand oversized sections with an LLM and report search-cost metrics; a deterministic merge always runs, collapsing subtrees whose structure does not beat a linear scan and keeping the removed titles on the parent as ``key_items``. optimize_expand: if False, skip the LLM expansion and only report merge metrics. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
     result = extract_toc(_validate_pdf(pdf))
     structure = result.get("structure", [])
     if optimize and structure:
         result["optimize"] = _optimize(structure, result.get("page_texts") or [],
                                        optimize_expand,
                                        optimize_model or summary_model)
     elif structure:
-        _thin(structure)
+        _merge(structure)
     if summary and structure:
         import asyncio
         from ..utils import ConfigLoader
@@ -122,6 +126,9 @@ def page_index_flash(pdf, summary=True, summary_model=None,
                                concurrency=summary_concurrency))
     else:
         result.pop("page_texts", None)
+        if structure:
+            from ..utils import strip_internal_keys
+            strip_internal_keys(structure)   # summarize_tree does this on its way out
     return result
 
 

pageindex/page_index.py

@@ -5,6 +5,7 @@
 import random
 import re
 from .utils import *
+from .tree_optimize import merge_tree
 import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
@@ -1246,7 +1247,7 @@ def page_index_main(doc, opt=None):
 
     async def page_index_builder():
         structure = await tree_parser(page_list, opt, doc=doc, logger=logger)
-        page_level_thinning(structure)
+        merge_tree(structure)
         if opt.if_add_node_id == 'yes':
             write_node_id(structure)
         if opt.if_add_node_text == 'yes':
@@ -1261,13 +1262,13 @@ async def page_index_builder():
                 # Create a clean structure without unnecessary fields for description generation
                 clean_structure = create_clean_structure_for_description(structure)
                 doc_description = generate_doc_description(clean_structure, model=getattr(opt, 'summary_model', None) or opt.model)
-                structure = format_structure(structure, order=['title', 'node_id', 'start_index', 'end_index', 'summary', 'text', 'nodes'])
+                structure = format_structure(structure, order=['title', 'node_id', 'start_index', 'end_index', 'key_items', 'summary', 'text', 'nodes'])
                 return {
                     'doc_name': get_pdf_name(doc),
                     'doc_description': doc_description,
                     'structure': structure,
                 }
-        structure = format_structure(structure, order=['title', 'node_id', 'start_index', 'end_index', 'summary', 'text', 'nodes'])
+        structure = format_structure(structure, order=['title', 'node_id', 'start_index', 'end_index', 'key_items', 'summary', 'text', 'nodes'])
         return {
             'doc_name': get_pdf_name(doc),
             'structure': structure,

pageindex/tree_optimize.py

@@ -36,6 +36,13 @@
 `key_items`: the pages stay reachable by scanning the parent, but the titles
 are routing information that would otherwise be lost.
 
+merge_same_page() runs first, as a special case of the same idea. Retrieval is
+page-granular, so frontier siblings covering identical pages cannot be told apart:
+an agent routed to any of them reads the same text, and because the leaf summary
+prompt sees only that text, their summaries come back near-identical. They collapse
+into one node titled with the union of theirs, which a leaf summary call rewrites
+when the node is large enough to earn one.
+
 merge is deterministic and needs no LLM; expand proposes subsections with the
 model configured as `summary_model` (falling back to `model`) in config.yaml.
 
@@ -54,11 +61,13 @@
 import sys
 from types import SimpleNamespace
 
-from .utils import ConfigLoader, _is_openai_model, llm_acompletion
+from .utils import (ConfigLoader, _is_openai_model, llm_acompletion,
+                    strip_internal_keys)
 
 TRIGGER_PAGES = 5        # only look ahead on nodes larger than this
 ROUTING_COST = 1         # R(v), in pages
 PAGE_CHARS = 6000        # per-page text handed to the model
+TITLE_MAX_CHARS = 200    # a union title longer than this falls back to a page label
 
 EXPAND_PROMPT = """You are splitting an over-long section of a PDF into its subsections.
 
@@ -456,6 +465,70 @@ def validate(structure, page_count):
 # MERGE
 # --------------------------------------------------------------------------
 
+def page_label(node):
+    """A node's page span, for use as a title of last resort."""
+    start, end = node["start_index"], subtree_end(node)
+    return f"p.{start}" if start == end else f"p.{start}-{end}"
+
+
+def union_title(titles, node):
+    """The titles of merged same-page siblings, joined.
+
+    Falls back to a page label when the join is empty or too long to serve as a
+    title - PRML, for instance, extracts whole exercise bodies as headings, and
+    two of those joined run past a thousand characters. Titles reach the model
+    (the parent summary prompt lists them, and they survive `format_structure`),
+    so this is the field that has to stay readable; `key_items` keeps the
+    untruncated original.
+    """
+    joined = "; ".join(title for title in titles if title)
+    if not joined or len(joined) > TITLE_MAX_CHARS:
+        return page_label(node)
+    return joined
+
+
+def merge_same_page(structure, log):
+    """Collapse frontier siblings that cover exactly the same pages.
+
+    Deterministic and free. Runs before merge() because a narrower tree changes
+    its ancestors' tree_cost, and before expand() because children an expand pass
+    lands on one page are the same redundancy arriving later.
+    """
+    changed = False
+
+    def visit(nodes):
+        nonlocal changed
+        groups = {}
+        for node in nodes:
+            visit(node.get("nodes") or [])
+            if is_frontier(node):
+                groups.setdefault((node["start_index"], subtree_end(node)), []).append(node)
+
+        for span, group in groups.items():
+            if len(group) < 2:
+                continue
+            keeper, dropped = group[0], group[1:]
+            titles = []
+            for node in group:                 # document order, key_items carried forward
+                titles.append(node["title"])
+                titles.extend(node.get("key_items") or [])
+            log.append({"op": "merge_same_page", "node_id": keeper.get("node_id"),
+                        "pages": list(span), "dropped": len(dropped),
+                        "dropped_ids": [n.get("node_id") for n in dropped],
+                        "key_items": titles})
+            keeper["key_items"] = titles
+            keeper["title"] = union_title(titles, keeper)
+            # tells summarize_tree this title was synthesized and may be rewritten;
+            # stripped from the output once summaries are done
+            keeper["_same_page"] = True
+            for node in dropped:
+                nodes.remove(node)
+            changed = True
+
+    visit(structure)
+    return changed
+
+
 def merge(structure, routing, log, frozen, progress=False):
     """Collapse any subtree whose structure does not beat a linear scan.
 
@@ -477,7 +550,8 @@ def visit(node):
         checked = tree_cost_via_frontier(node, routing)
         span = S(node)
         if span <= cost:
-            removed = [c["node_id"] for c, _ in flatten(node["nodes"])]
+            # trees arrive here before ids are assigned in the main pipeline
+            removed = [c.get("node_id") for c, _ in flatten(node["nodes"])]
             # titles are routing information; keep them on the parent, in document
             # order, carrying forward anything an earlier merge already folded in
             titles = []
@@ -496,14 +570,24 @@ def visit(node):
                 node["key_items"] = titles
             frozen.add(node.get("node_id"))
             changed = True
-            note(progress, f"    merge  {node.get('node_id'):>8}  "
+            note(progress, f"    merge  {node.get('node_id') or '-':>8}  "
                            f"S={span} <= tree_cost={cost}  dropped {len(removed)} node(s)")
 
     for root in list(structure):
         visit(root)
     return changed
 
 
+def merge_tree(structure):
+    """Deterministic merge over a structure list; the no-LLM default path.
+
+    One bottom-up pass reaches the fixpoint: every decision is made after the
+    subtree below it is final.
+    """
+    merge(structure, ROUTING_COST, [], set())
+    return structure
+
+
 # --------------------------------------------------------------------------
 # EXPAND
 # --------------------------------------------------------------------------
@@ -689,11 +773,13 @@ async def optimize(structure, pages, lines, model=None, routing=ROUTING_COST,
     for round_no in range(1, max_rounds + 1):
         rounds = round_no
         note(progress, f"  round {round_no}")
+        same_page = merge_same_page(structure, log) if do_merge else False
         merged = merge(structure, routing, log, frozen, progress) if do_merge else False
         expanded = await expand(structure, pages, lines, opts, log, frozen) \
             if do_expand else False
-        log.append({"op": "round", "round": round_no, "merged": merged, "expanded": expanded})
-        if not (merged or expanded):
+        log.append({"op": "round", "round": round_no, "same_page": same_page,
+                    "merged": merged, "expanded": expanded})
+        if not (same_page or merged or expanded):
             break
 
     id_map = relabel(structure) if do_relabel else {}
@@ -703,6 +789,9 @@ async def optimize(structure, pages, lines, model=None, routing=ROUTING_COST,
     return {"structure": structure, "log": log, "rounds": rounds,
             "before": before, "after": after, "id_map": id_map,
             "merges": sum(1 for e in log if e["op"] == "merge"),
+            "same_page_merges": sum(1 for e in log if e["op"] == "merge_same_page"),
+            "same_page_dropped": sum(e["dropped"] for e in log
+                                     if e["op"] == "merge_same_page"),
             "expands": sum(1 for e in log if e.get("decision") == "expand"),
             "kept_collapsed": sum(1 for e in log if e.get("decision") == "keep_collapsed"),
             "new_issues": issues}
@@ -728,6 +817,7 @@ def optimize_tree(doc, pdf_path=None, model=None, do_expand=None, **kwargs):
     result = asyncio.run(optimize(structure, pages, lines, model=model,
                                   page_count=page_count, do_expand=do_expand,
                                   **kwargs))
+    strip_internal_keys(result["structure"])
     doc["structure"] = result["structure"]
     return result
 
@@ -835,6 +925,7 @@ async def main():
     if result["new_issues"]:
         print(f"\nnew validation issues: {result['new_issues']}")
 
+    strip_internal_keys(structure)
     refined = dict(original)
     refined["structure"] = structure
     json.dump(refined, open(out_path, "w"), indent=2, ensure_ascii=False)