From cf04a580b2b25d62e09d3f600fcf0d882ce3c977 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 2 Aug 2026 21:29:47 +0800 Subject: [PATCH] Keep merge inside --optimize --- pageindex/flash/README.md | 5 +++-- pageindex/flash/api.py | 11 +---------- run_pageindex.py | 5 ++--- 3 files changed, 6 insertions(+), 15 deletions(-) diff --git a/pageindex/flash/README.md b/pageindex/flash/README.md index 8336132c1..85b4b5ede 100644 --- a/pageindex/flash/README.md +++ b/pageindex/flash/README.md @@ -30,14 +30,15 @@ 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 + "key_items": [str], # with --optimize: titles of merged-away subsections "nodes": [...], # absent on leaf nodes } ], } ``` -Page indexes are 1-based. `nodes` nests the same shape recursively. +Page indexes are 1-based. `nodes` nests the same shape recursively. Without +`--optimize` the extracted tree is returned as-is. ## Benchmark diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index 3eb48d4e1..365664f58 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -68,13 +68,6 @@ def _validate_pdf(pdf): return pdf -def _merge(structure): - from ..tree_optimize import merge_tree - from ..utils import write_node_id - merge_tree(structure) - write_node_id(structure) - - async def _summarize(structure, page_list, model, concurrency=None): from ..utils import summarize_tree await summarize_tree(structure, page_list, model=model, concurrency=concurrency) @@ -105,15 +98,13 @@ 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, 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. """ + """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 before summaries: a deterministic merge collapses subtrees whose structure does not beat a linear scan, keeping the removed titles on the parent as ``key_items``, then an LLM pass expands oversized sections. Without it the extracted tree is returned unchanged. optimize_expand: if False, run the merge but skip the LLM expansion. 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: - _merge(structure) if summary and structure: import asyncio from ..utils import ConfigLoader diff --git a/run_pageindex.py b/run_pageindex.py index 3a8c9f8de..fab01fd36 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -13,9 +13,8 @@ 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 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)') + help='Refine the tree for search cost: a deterministic merge, then an ' + 'LLM expansion pass; pass `merge` to run the merge alone (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,