apunkt/PageIndex
1
0
Fork 0
mirror of https://github.com/VectifyAI/PageIndex.git synced 2026-04-24 23:56:21 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

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

43
pageindex/page_index_md.py
View file

@ -1,6 +1,7 @@
import asyncio
import json
import re
import os
try:
from .utils import *
except:
@ -239,7 +240,7 @@ def clean_tree_for_output(tree_nodes):
return cleaned_nodes
async def md_to_tree(md_path, if_thinning=True, min_token_threshold=None, if_summary=True, summary_token_threshold=None, model=None):
async def md_to_tree(md_path, if_thinning=True, min_token_threshold=None, if_add_node_summary='no', summary_token_threshold=None, model=None, if_add_doc_description='yes', if_add_node_text='no', if_add_node_id='yes'):
with open(md_path, 'r', encoding='utf-8') as f:
markdown_content = f.read()
@ -257,14 +258,44 @@ async def md_to_tree(md_path, if_thinning=True, min_token_threshold=None, if_sum
print(f"Building tree from nodes...")
tree_structure = build_tree_from_nodes(nodes_with_content)
if if_summary:
print(f"Generating summaries for each node...")
tree_structure = await generate_summaries_for_structure_md(tree_structure,summary_token_threshold=summary_token_threshold, model=model)
# Add node IDs if requested (matching PDF behavior)
if if_add_node_id == 'yes':
write_node_id(tree_structure)
print(f"Formatting tree structure...")
tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'summary', 'prefix_summary', 'text', 'line_num', 'nodes'])
return tree_structure
if if_add_node_summary == 'yes':
# Always include text for summary generation
tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'summary', 'prefix_summary', 'text', 'line_num', 'nodes'])
print(f"Generating summaries for each node...")
tree_structure = await generate_summaries_for_structure_md(tree_structure, summary_token_threshold=summary_token_threshold, model=model)
if if_add_node_text == 'no':
# Remove text after summary generation if not requested
tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'summary', 'prefix_summary', 'line_num', 'nodes'])
if if_add_doc_description == 'yes':
print(f"Generating document description...")
# Create a clean structure without unnecessary fields for description generation
clean_structure = create_clean_structure_for_description(tree_structure)
doc_description = generate_doc_description(clean_structure, model=model)
return {
'doc_name': os.path.splitext(os.path.basename(md_path))[0],
'doc_description': doc_description,
'structure': tree_structure,
}
else:
# No summaries needed, format based on text preference
if if_add_node_text == 'yes':
tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'summary', 'prefix_summary', 'text', 'line_num', 'nodes'])
else:
tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'summary', 'prefix_summary', 'line_num', 'nodes'])
return {
'doc_name': os.path.splitext(os.path.basename(md_path))[0],
'structure': tree_structure,
}
if __name__ == "__main__":