apunkt/PageIndex
1
0
Fork 0
mirror of https://github.com/VectifyAI/PageIndex.git synced 2026-05-02 03:22:37 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

fix params

Browse source
This commit is contained in:
Ray 2025-08-28 12:45:39 +08:00
parent 480f7583f7
commit 6d1b505541
6 changed files with 126 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

82
run_pageindex.py
View file

@ -7,24 +7,27 @@ from pageindex.page_index_md import md_to_tree
if __name__ == "__main__":
# Set up argument parser
parser = argparse.ArgumentParser(description='Process PDF or Markdown document and generate structure')
parser.add_argument('--file_path', type=str, help='Path to the PDF or Markdown file')
parser.add_argument('--file_type', type=str, choices=['pdf', 'markdown', 'md'], default='pdf',
help='Type of file to process (pdf, markdown, or md)')
parser.add_argument('--model', type=str, default='gpt-4o-2024-11-20', help='Model to use')
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-4.1', 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='no',
help='Whether to add summary to the node')
parser.add_argument('--if-add-doc-description', type=str, default='yes',
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)')
@ -34,14 +37,19 @@ if __name__ == "__main__":
help='Token threshold for generating summaries (markdown only)')
args = parser.parse_args()
# Determine file type from extension if not specified
if args.file_type == 'pdf' and args.file_path:
if args.file_path.lower().endswith(('.md', '.markdown')):
args.file_type = 'markdown'
elif not args.file_path.lower().endswith('.pdf'):
raise ValueError("File must be a PDF or Markdown file")
# 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.file_type == 'pdf':
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(
@ -56,42 +64,62 @@ if __name__ == "__main__":
)
# Process the PDF
toc_with_page_number = page_index_main(args.file_path, opt)
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.file_path))[0]
pdf_name = os.path.splitext(os.path.basename(args.pdf_path))[0]
os.makedirs('./results', exist_ok=True)
with open(f'./results/{pdf_name}_structure.json', 'w', encoding='utf-8') as f:
json.dump(toc_with_page_number, f, indent=2)
elif args.file_type in ['markdown', 'md']:
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...')
# Configure markdown options
if_thinning = args.if_thinning.lower() == 'yes'
if_summary = args.if_add_node_summary.lower() == 'yes'
# 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.file_path,
if_thinning=if_thinning,
md_path=args.md_path,
if_thinning=args.if_thinning.lower() == 'yes',
min_token_threshold=args.thinning_threshold,
if_summary=if_summary,
if_add_node_summary=opt.if_add_node_summary,
summary_token_threshold=args.summary_token_threshold,
model=args.model
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.file_path))[0]
md_name = os.path.splitext(os.path.basename(args.md_path))[0]
os.makedirs('./results', exist_ok=True)
with open(f'./results/{md_name}_structure.json', 'w', encoding='utf-8') as f:
json.dump(toc_with_page_number, f, indent=2, ensure_ascii=False)
else:
raise ValueError(f"Unsupported file type: {args.file_type}. Supported types are 'pdf', 'markdown', or 'md'")
json.dump(toc_with_page_number, f, indent=2, ensure_ascii=False)