diff --git a/.gitignore b/.gitignore index 47d38baef..5d8265eeb 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,6 @@ log/* logs/ parts/* json_results/* +*.egg-info/ +dist/ +build/ diff --git a/pageindex/cli.py b/pageindex/cli.py new file mode 100644 index 000000000..ece9619e3 --- /dev/null +++ b/pageindex/cli.py @@ -0,0 +1,119 @@ +import argparse +import asyncio +import json +import os + + +def main(): + parser = argparse.ArgumentParser(description='Process PDF or Markdown document and generate structure') + 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='gpt-4o-2024-11-20', help='Model to use') + + parser.add_argument('--toc-check-pages', type=int, default=20, + help='Number of pages to check for table of contents (PDF only)') + parser.add_argument('--max-pages-per-node', type=int, default=10, + help='Maximum number of pages per node (PDF only)') + parser.add_argument('--max-tokens-per-node', type=int, default=20000, + help='Maximum number of tokens per node (PDF only)') + + parser.add_argument('--if-add-node-id', type=str, default='yes', + help='Whether to add node id to the node') + parser.add_argument('--if-add-node-summary', type=str, default='yes', + help='Whether to add summary to the node') + parser.add_argument('--if-add-doc-description', type=str, default='no', + help='Whether to add doc description to the doc') + parser.add_argument('--if-add-node-text', type=str, default='no', + help='Whether to 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)') + parser.add_argument('--thinning-threshold', type=int, default=5000, + help='Minimum token threshold for thinning (markdown only)') + parser.add_argument('--summary-token-threshold', type=int, default=200, + help='Token threshold for generating summaries (markdown only)') + args = parser.parse_args() + + 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") + + if args.pdf_path: + from pageindex import config, page_index_main + + 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}") + + opt = config( + 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 + ) + + toc_with_page_number = page_index_main(args.pdf_path, opt) + print('Parsing done, saving to file...') + + 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: + from pageindex.page_index_md import md_to_tree + from pageindex.utils import ConfigLoader + + 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}") + + print('Processing markdown file...') + + config_loader = ConfigLoader() + 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 + } + 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', + min_token_threshold=args.thinning_threshold, + if_add_node_summary=opt.if_add_node_summary, + summary_token_threshold=args.summary_token_threshold, + model=opt.model, + if_add_doc_description=opt.if_add_doc_description, + if_add_node_text=opt.if_add_node_text, + if_add_node_id=opt.if_add_node_id + )) + + print('Parsing done, saving to file...') + + 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}') diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..aa580f6c1 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,29 @@ +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[project] +name = "pageindex" +version = "0.1.0b0" +description = "Transform long documents into hierarchical tree structures for LLM-based retrieval" +readme = "README.md" +license = { file = "LICENSE" } +requires-python = ">=3.8" +dependencies = [ + "openai>=1.0.0", + "pymupdf>=1.26.0", + "PyPDF2>=3.0.0", + "python-dotenv>=1.0.0", + "tiktoken>=0.5.0", + "pyyaml>=6.0", +] + +[project.scripts] +pageindex = "pageindex.cli:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["pageindex*"] + +[tool.setuptools.package-data] +pageindex = ["config.yaml"] diff --git a/run_pageindex.py b/run_pageindex.py index 107024505..1153613e9 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -1,133 +1,4 @@ -import argparse -import os -import json -from pageindex import * -from pageindex.page_index_md import md_to_tree +from pageindex.cli import main if __name__ == "__main__": - # Set up argument parser - parser = argparse.ArgumentParser(description='Process PDF or Markdown document and generate structure') - 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='gpt-4o-2024-11-20', help='Model to use') - - parser.add_argument('--toc-check-pages', type=int, default=20, - help='Number of pages to check for table of contents (PDF only)') - parser.add_argument('--max-pages-per-node', type=int, default=10, - help='Maximum number of pages per node (PDF only)') - parser.add_argument('--max-tokens-per-node', type=int, default=20000, - help='Maximum number of tokens per node (PDF only)') - - parser.add_argument('--if-add-node-id', type=str, default='yes', - help='Whether to add node id to the node') - parser.add_argument('--if-add-node-summary', type=str, default='yes', - help='Whether to add summary to the node') - parser.add_argument('--if-add-doc-description', type=str, default='no', - help='Whether to add doc description to the doc') - parser.add_argument('--if-add-node-text', type=str, default='no', - help='Whether to 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)') - parser.add_argument('--thinning-threshold', type=int, default=5000, - help='Minimum token threshold for thinning (markdown only)') - 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") - - 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 - # Configure options - opt = config( - 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 - ) - - # 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] - 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', - min_token_threshold=args.thinning_threshold, - if_add_node_summary=opt.if_add_node_summary, - summary_token_threshold=args.summary_token_threshold, - model=opt.model, - if_add_doc_description=opt.if_add_doc_description, - 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] - 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}') \ No newline at end of file + main()