mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-05-13 16:52:37 +02:00
feat: add PageIndex SDK with local/cloud dual-mode support (#207)
This commit is contained in:
Kylin
GitHub
parent
8f1ed7783b
commit
b63fd9779a
45 changed files with 4225 additions and 274 deletions
102
run_pageindex.py
102
run_pageindex.py
|
|
@ -1,9 +1,9 @@
|
|||
import argparse
|
||||
import os
|
||||
import json
|
||||
from pageindex import *
|
||||
from pageindex.page_index_md import md_to_tree
|
||||
from pageindex.utils import ConfigLoader
|
||||
from pageindex.index.page_index import *
|
||||
from pageindex.index.page_index_md import md_to_tree
|
||||
from pageindex.config import IndexConfig
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Set up argument parser
|
||||
|
|
@ -11,7 +11,7 @@ if __name__ == "__main__":
|
|||
parser.add_argument('--pdf_path', type=str, help='Path to the PDF file')
|
||||
parser.add_argument('--md_path', type=str, help='Path to the Markdown file')
|
||||
|
||||
parser.add_argument('--model', type=str, default=None, help='Model to use (overrides config.yaml)')
|
||||
parser.add_argument('--model', type=str, default=None, help='Model to use')
|
||||
|
||||
parser.add_argument('--toc-check-pages', type=int, default=None,
|
||||
help='Number of pages to check for table of contents (PDF only)')
|
||||
|
|
@ -20,15 +20,15 @@ if __name__ == "__main__":
|
|||
parser.add_argument('--max-tokens-per-node', type=int, default=None,
|
||||
help='Maximum number of tokens per node (PDF only)')
|
||||
|
||||
parser.add_argument('--if-add-node-id', type=str, default=None,
|
||||
help='Whether to add node id to the node')
|
||||
parser.add_argument('--if-add-node-summary', type=str, default=None,
|
||||
help='Whether to add summary to the node')
|
||||
parser.add_argument('--if-add-doc-description', type=str, default=None,
|
||||
help='Whether to add doc description to the doc')
|
||||
parser.add_argument('--if-add-node-text', type=str, default=None,
|
||||
help='Whether to add text to the node')
|
||||
|
||||
parser.add_argument('--if-add-node-id', action='store_true', default=None,
|
||||
help='Add node id to the node')
|
||||
parser.add_argument('--if-add-node-summary', action='store_true', default=None,
|
||||
help='Add summary to the node')
|
||||
parser.add_argument('--if-add-doc-description', action='store_true', default=None,
|
||||
help='Add doc description to the doc')
|
||||
parser.add_argument('--if-add-node-text', action='store_true', default=None,
|
||||
help='Add text to the node')
|
||||
|
||||
# Markdown specific arguments
|
||||
parser.add_argument('--if-thinning', type=str, default='no',
|
||||
help='Whether to apply tree thinning for markdown (markdown only)')
|
||||
|
|
@ -37,77 +37,61 @@ if __name__ == "__main__":
|
|||
parser.add_argument('--summary-token-threshold', type=int, default=200,
|
||||
help='Token threshold for generating summaries (markdown only)')
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
# Validate that exactly one file type is specified
|
||||
if not args.pdf_path and not args.md_path:
|
||||
raise ValueError("Either --pdf_path or --md_path must be specified")
|
||||
if args.pdf_path and args.md_path:
|
||||
raise ValueError("Only one of --pdf_path or --md_path can be specified")
|
||||
|
||||
|
||||
# Build IndexConfig from CLI args (None values use defaults)
|
||||
config_overrides = {
|
||||
k: v for k, v in {
|
||||
"model": args.model,
|
||||
"toc_check_page_num": args.toc_check_pages,
|
||||
"max_page_num_each_node": args.max_pages_per_node,
|
||||
"max_token_num_each_node": args.max_tokens_per_node,
|
||||
"if_add_node_id": args.if_add_node_id,
|
||||
"if_add_node_summary": args.if_add_node_summary,
|
||||
"if_add_doc_description": args.if_add_doc_description,
|
||||
"if_add_node_text": args.if_add_node_text,
|
||||
}.items() if v is not None
|
||||
}
|
||||
opt = IndexConfig(**config_overrides)
|
||||
|
||||
if args.pdf_path:
|
||||
# Validate PDF file
|
||||
if not args.pdf_path.lower().endswith('.pdf'):
|
||||
raise ValueError("PDF file must have .pdf extension")
|
||||
if not os.path.isfile(args.pdf_path):
|
||||
raise ValueError(f"PDF file not found: {args.pdf_path}")
|
||||
|
||||
# Process PDF file
|
||||
user_opt = {
|
||||
'model': args.model,
|
||||
'toc_check_page_num': args.toc_check_pages,
|
||||
'max_page_num_each_node': args.max_pages_per_node,
|
||||
'max_token_num_each_node': args.max_tokens_per_node,
|
||||
'if_add_node_id': args.if_add_node_id,
|
||||
'if_add_node_summary': args.if_add_node_summary,
|
||||
'if_add_doc_description': args.if_add_doc_description,
|
||||
'if_add_node_text': args.if_add_node_text,
|
||||
}
|
||||
opt = ConfigLoader().load({k: v for k, v in user_opt.items() if v is not None})
|
||||
|
||||
# Process the PDF
|
||||
toc_with_page_number = page_index_main(args.pdf_path, opt)
|
||||
print('Parsing done, saving to file...')
|
||||
|
||||
|
||||
# Save results
|
||||
pdf_name = os.path.splitext(os.path.basename(args.pdf_path))[0]
|
||||
pdf_name = os.path.splitext(os.path.basename(args.pdf_path))[0]
|
||||
output_dir = './results'
|
||||
output_file = f'{output_dir}/{pdf_name}_structure.json'
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
|
||||
with open(output_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(toc_with_page_number, f, indent=2)
|
||||
|
||||
|
||||
print(f'Tree structure saved to: {output_file}')
|
||||
|
||||
|
||||
elif args.md_path:
|
||||
# Validate Markdown file
|
||||
if not args.md_path.lower().endswith(('.md', '.markdown')):
|
||||
raise ValueError("Markdown file must have .md or .markdown extension")
|
||||
if not os.path.isfile(args.md_path):
|
||||
raise ValueError(f"Markdown file not found: {args.md_path}")
|
||||
|
||||
|
||||
# Process markdown file
|
||||
print('Processing markdown file...')
|
||||
|
||||
# Process the markdown
|
||||
import asyncio
|
||||
|
||||
# Use ConfigLoader to get consistent defaults (matching PDF behavior)
|
||||
from pageindex.utils import ConfigLoader
|
||||
config_loader = ConfigLoader()
|
||||
|
||||
# Create options dict with user args
|
||||
user_opt = {
|
||||
'model': args.model,
|
||||
'if_add_node_summary': args.if_add_node_summary,
|
||||
'if_add_doc_description': args.if_add_doc_description,
|
||||
'if_add_node_text': args.if_add_node_text,
|
||||
'if_add_node_id': args.if_add_node_id
|
||||
}
|
||||
|
||||
# Load config with defaults from config.yaml
|
||||
opt = config_loader.load(user_opt)
|
||||
|
||||
|
||||
toc_with_page_number = asyncio.run(md_to_tree(
|
||||
md_path=args.md_path,
|
||||
if_thinning=args.if_thinning.lower() == 'yes',
|
||||
|
|
@ -119,16 +103,16 @@ if __name__ == "__main__":
|
|||
if_add_node_text=opt.if_add_node_text,
|
||||
if_add_node_id=opt.if_add_node_id
|
||||
))
|
||||
|
||||
|
||||
print('Parsing done, saving to file...')
|
||||
|
||||
|
||||
# Save results
|
||||
md_name = os.path.splitext(os.path.basename(args.md_path))[0]
|
||||
md_name = os.path.splitext(os.path.basename(args.md_path))[0]
|
||||
output_dir = './results'
|
||||
output_file = f'{output_dir}/{md_name}_structure.json'
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
|
||||
with open(output_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(toc_with_page_number, f, indent=2, ensure_ascii=False)
|
||||
|
||||
print(f'Tree structure saved to: {output_file}')
|
||||
|
||||
print(f'Tree structure saved to: {output_file}')
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue