Add pyproject.toml to make package installable

.gitignore

@@ -13,3 +13,6 @@ log/*
 logs/
 parts/*
 json_results/*
+*.egg-info/
+dist/
+build/

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}')

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"]

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}')
+    main()