From 4a116daa2fefdbf6f6dd330c75c22ef2ab364caf Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 2 Aug 2026 17:56:37 +0800 Subject: [PATCH 1/3] Replace thinning with merge --- README.md | 2 +- pageindex/flash/README.md | 1 + pageindex/flash/api.py | 17 +++++++++-------- pageindex/page_index.py | 7 ++++--- pageindex/tree_optimize.py | 15 +++++++++++++-- pageindex/utils.py | 1 + run_pageindex.py | 5 +++-- 7 files changed, 32 insertions(+), 16 deletions(-) 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..1088c6fcf 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 @@ -101,7 +102,7 @@ def _optimize(structure, page_texts, do_expand, model): 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 +110,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 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..64ebdd345 100644 --- a/pageindex/tree_optimize.py +++ b/pageindex/tree_optimize.py @@ -477,7 +477,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 +497,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 +505,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 # -------------------------------------------------------------------------- diff --git a/pageindex/utils.py b/pageindex/utils.py index 67f9a8ddd..f715e40f8 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -814,6 +814,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, From 278cc7f28a8f39d0e00e2fbcdf9ed87e3fac123a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E9=B8=A3=E5=A4=A9?= Date: Sun, 2 Aug 2026 18:11:56 +0800 Subject: [PATCH 2/3] improve tree optimization Carried from bench/flash-cost; benchmark artifacts left there. --- pageindex/flash/api.py | 6 ++ pageindex/tree_optimize.py | 81 ++++++++++++++++++++++++++- pageindex/utils.py | 110 ++++++++++++++++++++++++++++--------- 3 files changed, 168 insertions(+), 29 deletions(-) diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index 1088c6fcf..3eb48d4e1 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -96,6 +96,9 @@ 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"]} @@ -123,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/tree_optimize.py b/pageindex/tree_optimize.py index 64ebdd345..24923492d 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. @@ -59,6 +66,7 @@ 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 +464,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. @@ -700,11 +772,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 {} @@ -714,6 +788,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} diff --git a/pageindex/utils.py b/pageindex/utils.py index f715e40f8..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 From e17dac205378d0055b255a4d401fe857577da59a Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 2 Aug 2026 18:23:23 +0800 Subject: [PATCH 3/3] Strip internal marker from optimize output --- pageindex/tree_optimize.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py index 24923492d..9277135f3 100644 --- a/pageindex/tree_optimize.py +++ b/pageindex/tree_optimize.py @@ -61,7 +61,8 @@ 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 @@ -816,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 @@ -923,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)