diff --git a/README.md b/README.md
index 752f91dec..5ce0ca5e6 100644
--- a/README.md
+++ b/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
diff --git a/pageindex/flash/README.md b/pageindex/flash/README.md
index 675c2acec..b8b4772e3 100644
--- a/pageindex/flash/README.md
+++ b/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
}
],
diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index 804169890..3eb48d4e1 100644
--- a/pageindex/flash/api.py
+++ b/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,13 +96,16 @@ 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:
@@ -109,7 +113,7 @@ def page_index_flash(pdf, summary=True, summary_model=None,
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
diff --git a/pageindex/page_index.py b/pageindex/page_index.py
index 4731f1ed1..c0b3ea935 100644
--- a/pageindex/page_index.py
+++ b/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,
diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index 9e8dd8b1e..9277135f3 100644
--- a/pageindex/tree_optimize.py
+++ b/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,7 +570,7 @@ 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):
@@ -504,6 +578,16 @@ def visit(node):
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)
diff --git a/pageindex/utils.py b/pageindex/utils.py
index 67f9a8ddd..b41cb919e 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -656,35 +656,73 @@ def get_intro_text(node, pdf_pages, max_pages=SUMMARY_INTRO_MAX_PAGES):
return get_text_of_pdf_pages(pdf_pages, node['start_index'], end)
-def parse_summary(reply):
- """The `summary` field of a model reply, or the reply itself when there is no
- such field. Not extract_json: that rewrites `None` to `null` and collapses
- whitespace in replies that parse as written."""
+def _reply_json(reply):
+ """The JSON object in a model reply, or None when none of it parses.
+
+ Not extract_json: that rewrites `None` to `null` and collapses whitespace in
+ replies that parse as written.
+ """
if not isinstance(reply, str) or not reply.strip():
- return ""
+ return None
text = reply.strip()
if '```' in text:
text = re.sub(r'^.*?```(?:json)?\s*', '', text, flags=re.S).split('```')[0]
start, end = text.find('{'), text.rfind('}')
- if start != -1 and end > start:
- obj = text[start:end + 1]
- collapsed = ' '.join(obj.split())
- parsed = None
- # repairs, tried only once the reply fails to parse as written
- for candidate in (obj, collapsed, collapsed.replace(',]', ']').replace(',}', '}')):
- try:
- parsed = json.loads(candidate)
- break
- except json.JSONDecodeError:
- continue
- if isinstance(parsed, dict) and 'summary' in parsed:
- summary = parsed['summary']
- if isinstance(summary, list):
- summary = ' '.join(str(item).strip() for item in summary if str(item).strip())
- return str(summary).strip() if summary else ""
+ if start == -1 or end <= start:
+ return None
+ obj = text[start:end + 1]
+ collapsed = ' '.join(obj.split())
+ # repairs, tried only once the reply fails to parse as written
+ for candidate in (obj, collapsed, collapsed.replace(',]', ']').replace(',}', '}')):
+ try:
+ return json.loads(candidate)
+ except json.JSONDecodeError:
+ continue
+ return None
+
+
+def parse_summary(reply):
+ """The `summary` field of a model reply, or the reply itself when there is no
+ such field."""
+ if not isinstance(reply, str) or not reply.strip():
+ return ""
+ parsed = _reply_json(reply)
+ if isinstance(parsed, dict) and 'summary' in parsed:
+ summary = parsed['summary']
+ if isinstance(summary, list):
+ summary = ' '.join(str(item).strip() for item in summary if str(item).strip())
+ return str(summary).strip() if summary else ""
return reply.strip()
+def parse_title(reply):
+ """The `title` field of a model reply, or "" when it is absent or unusable.
+
+ Unlike parse_summary there is no falling back to the raw reply: a title that
+ did not come back as a named field is not a title, and the caller keeps the
+ deterministic one it already has.
+ """
+ parsed = _reply_json(reply)
+ if not isinstance(parsed, dict):
+ return ""
+ title = parsed.get('title')
+ if isinstance(title, list):
+ title = ' '.join(str(item).strip() for item in title if str(item).strip())
+ return ' '.join(str(title).split()) if title else ""
+
+
+def strip_internal_keys(structure):
+ """Drop the bookkeeping keys the optimize/summary passes leave behind."""
+ nodes = structure if isinstance(structure, list) else [structure]
+ for node in nodes:
+ if not isinstance(node, dict):
+ continue
+ node.pop('_same_page', None)
+ if node.get('nodes'):
+ strip_internal_keys(node['nodes'])
+ return structure
+
+
async def summarize_tree(structure, pdf_pages, model=None,
small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None):
@@ -697,28 +735,45 @@ async def summarize_tree(structure, pdf_pages, model=None,
async def ask(prompt):
async with semaphore:
- response = await llm_acompletion(model, prompt)
- return parse_summary(response)
+ return await llm_acompletion(model, prompt)
async def leaf_summary(node):
text = get_text_of_pdf_pages(pdf_pages, node['start_index'], node['end_index'])
if count_tokens(text, model="gpt-4o") < small_node_tokens:
return text.strip()
+
+ # A node merged from same-page siblings carries a title joined from theirs.
+ # This call already has the page text in front of it, so the better title
+ # costs no extra call; every other node keeps the heading the document
+ # printed, and its prompt stays byte-identical to the one without this.
+ retitle = bool(node.get('_same_page'))
+ titles = "; ".join(node.get('key_items') or [])
+ ask_title = (f"\n The text is one page holding several short sections: {titles}. "
+ f"Also return a short title, at most 12 words, naming what the "
+ f"whole page covers." if retitle else "")
+ title_field = ('\n "title": ,'
+ if retitle else "")
+
prompt = f"""You are given a text chunk from a document.
Your task is to generate a concise description of everything that is covered in the text, summarizing all its points without omitting any type of content.
- Keep the description concise and to the point, avoiding unnecessary details.
+ Keep the description concise and to the point, avoiding unnecessary details.{ask_title}
Given Text: {text}
Reply strictly in the following JSON format:
- {{
+ {{{title_field}
"points": ,
"summary":
}}
Follow strictly the above JSON return format. Do not include any other text!
"""
- return await ask(prompt)
+ reply = await ask(prompt)
+ if retitle:
+ written = parse_title(reply)
+ if written:
+ node['title'] = written
+ return parse_summary(reply)
async def parent_summary(node):
children = node['nodes']
@@ -744,7 +799,7 @@ async def parent_summary(node):
Follow strictly the above JSON return format. Do not include any other text!
"""
- return await ask(prompt)
+ return parse_summary(await ask(prompt))
async def visit(node):
children = node.get('nodes') or []
@@ -755,6 +810,7 @@ async def visit(node):
node['summary'] = await (parent_summary(node) if children else leaf_summary(node))
await asyncio.gather(*(visit(root) for root in structure))
+ strip_internal_keys(structure)
return structure
@@ -814,6 +870,7 @@ def format_structure(structure, order=None):
def page_level_thinning(structure, thinning_threshold_node_num=20, min_pages_for_large_tree=3):
+ """Legacy; superseded by tree_optimize.merge_tree."""
def count_nodes(nodes):
total = 0
for node in nodes:
diff --git a/run_pageindex.py b/run_pageindex.py
index d7cfc135f..3a8c9f8de 100644
--- a/run_pageindex.py
+++ b/run_pageindex.py
@@ -13,8 +13,9 @@
parser.add_argument('--flash', action='store_true', help='Use PageIndex Flash (with --pdf_path)')
parser.add_argument('--optimize', nargs='?', const='full', choices=['full', 'merge'],
default=None,
- help='Refine the tree for search cost: merge + LLM expand; '
- 'pass `merge` to skip the expansion pass (PDF only)')
+ help='Refine the tree with an LLM expansion pass and report search-cost '
+ 'metrics; pass `merge` to skip expansion and only report the '
+ 'deterministic merge every run performs (PDF only)')
parser.add_argument('--model', type=str, default=None, help='Model to use (overrides config.yaml)')
parser.add_argument('--summary-model', type=str, default=None,