diff --git a/.gitignore b/.gitignore index 5193735ca..434466009 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ __pycache__ logs/ .pageindex/ dist/ +results/ diff --git a/examples/documents/test_doc.txt b/examples/documents/test_doc.txt new file mode 100644 index 000000000..24fd7527e --- /dev/null +++ b/examples/documents/test_doc.txt @@ -0,0 +1,12 @@ +# Test Document + +This is a test document to verify the text conversion feature. + +## Section 1 +Here is some text in section 1. + +### Subsection 1.1 +More text here. + +## Section 2 +And this is section 2. diff --git a/examples/documents/test_doc2.txt b/examples/documents/test_doc2.txt new file mode 100644 index 000000000..9de52edb2 --- /dev/null +++ b/examples/documents/test_doc2.txt @@ -0,0 +1,9 @@ +# Second Document + +This is another test document to test merging. + +## Another Section +Here is some other text. + +### Another Subsection +More text here for doc 2. diff --git a/pageindex/corpus_builder.py b/pageindex/corpus_builder.py new file mode 100644 index 000000000..c6cc5aa1b --- /dev/null +++ b/pageindex/corpus_builder.py @@ -0,0 +1,61 @@ +import json +import os +try: + from .utils import write_node_id +except ImportError: + from utils import write_node_id + +def build_corpus(corpus_files: list, corpus_name: str = "Corpus") -> dict: + """ + Takes a list of file paths to existing '*_structure.json' files, + merges them into a single corpus tree under a root node, and renumbers + the node_ids sequentially. + """ + corpus_root = { + "title": corpus_name, + "nodes": [] + } + + total_line_count = 0 + + for path in corpus_files: + if not os.path.exists(path): + print(f"Warning: Corpus file not found: {path}") + continue + + with open(path, 'r', encoding='utf-8') as f: + try: + doc_data = json.load(f) + except json.JSONDecodeError: + print(f"Warning: Failed to parse JSON from {path}") + continue + + # Get the document title and structure + doc_title = doc_data.get("doc_name", os.path.basename(path)) + structure = doc_data.get("structure", []) + total_line_count += doc_data.get("line_count", 0) + + doc_node = { + "title": doc_title, + "nodes": structure + } + + corpus_root["nodes"].append(doc_node) + + # Renumber the entire tree sequentially starting from 0001 + # write_node_id expects a list (or a dict) and modifies in place, + # and we want it to start from 1 (the default in write_node_id is 0 but it increments before stringifying, wait, let's check). + # In utils.py: + # def write_node_id(data, node_id=0): + # if isinstance(data, dict): + # data['node_id'] = str(node_id).zfill(4) + # node_id += 1 ... + # So if we pass node_id=1, the root gets 0001. + + write_node_id([corpus_root], node_id=1) + + return { + "doc_name": corpus_name, + "line_count": total_line_count, + "structure": [corpus_root] + } diff --git a/pageindex/format_converter.py b/pageindex/format_converter.py new file mode 100644 index 000000000..1682cd72a --- /dev/null +++ b/pageindex/format_converter.py @@ -0,0 +1,95 @@ +import os + +def convert_docx_to_md(docx_path: str) -> str: + """ + Reads a .docx file and converts it to markdown, preserving heading levels. + Requires python-docx. + """ + try: + import docx + except ImportError: + raise ImportError("The 'python-docx' package is required to parse .docx files. Please install it.") + + doc = docx.Document(docx_path) + md_lines = [] + + for para in doc.paragraphs: + text = para.text.strip() + if not text: + continue + + style_name = para.style.name if para.style else "" + + # Check if the style is a heading + if style_name.startswith('Heading'): + try: + # e.g., "Heading 1" -> level 1 + level_str = style_name.replace('Heading', '').strip() + level = int(level_str) + # Cap the level at 6 for markdown + level = min(level, 6) + md_lines.append(f"{'#' * level} {text}\n") + except ValueError: + # Fallback if parsing fails + md_lines.append(f"{text}\n") + elif style_name == 'Title': + md_lines.append(f"# {text}\n") + else: + md_lines.append(f"{text}\n") + + return "\n".join(md_lines) + + +def convert_html_to_md(html_path: str) -> str: + """ + Reads an .html file and converts it to markdown using BeautifulSoup and markdownify. + """ + try: + from bs4 import BeautifulSoup + import markdownify + except ImportError: + raise ImportError("The 'beautifulsoup4' and 'markdownify' packages are required to parse .html files.") + + with open(html_path, 'r', encoding='utf-8') as f: + html_content = f.read() + + soup = BeautifulSoup(html_content, 'html.parser') + + # Remove script and style elements + for script in soup(["script", "style", "nav", "footer", "header", "aside"]): + script.extract() + + # Convert to markdown with ATX headings (### instead of underlines) + md_text = markdownify.markdownify(str(soup), heading_style="ATX") + + # Clean up excessive newlines + import re + md_text = re.sub(r'\n{3,}', '\n\n', md_text).strip() + return md_text + + +def convert_txt_to_md(txt_path: str) -> str: + """ + Reads a plain .txt file. It doesn't need actual conversion, but we wrap it here for consistency. + """ + with open(txt_path, 'r', encoding='utf-8') as f: + return f.read() + + +def convert_to_markdown(file_path: str) -> str: + """ + Determines the file type and returns its markdown representation. + """ + ext = os.path.splitext(file_path)[1].lower() + + if ext == '.docx': + return convert_docx_to_md(file_path) + elif ext in ['.html', '.htm']: + return convert_html_to_md(file_path) + elif ext == '.txt': + return convert_txt_to_md(file_path) + elif ext in ['.md', '.markdown']: + with open(file_path, 'r', encoding='utf-8') as f: + return f.read() + else: + raise ValueError(f"Unsupported file format for conversion: {ext}") diff --git a/requirements.txt b/requirements.txt index 2406eef11..16f64797c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,6 @@ python-dotenv==1.2.2 pyyaml==6.0.2 regex>=2024.0.0 sortedcontainers==2.4.0 +python-docx>=1.1.0 +beautifulsoup4>=4.12.0 +markdownify>=0.13.1 diff --git a/run_pageindex.py b/run_pageindex.py index f2642b8a6..bfb9723ce 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -14,6 +14,14 @@ 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('--docx_path', type=str, help='Path to the DOCX file') + parser.add_argument('--html_path', type=str, help='Path to the HTML file') + parser.add_argument('--txt_path', type=str, help='Path to the TXT file') + + # Corpus specific arguments + parser.add_argument('--corpus_name', type=str, default='Corpus', help='Name of the unified corpus root node (used with --corpus_files)') + parser.add_argument('--corpus_files', nargs='+', help='List of *_structure.json files to merge into a single corpus') + parser.add_argument('--mode', choices=['flash', 'standard'], default='flash', help='Processing mode (default: flash)') parser.add_argument('--flash', action='store_true', default=False, @@ -58,14 +66,62 @@ parser.add_argument('--summary-token-threshold', type=int, default=200, help='Token threshold for generating summaries (markdown only)') args = parser.parse_args() + + if args.corpus_files: + from pageindex.corpus_builder import build_corpus + print(f"Building corpus '{args.corpus_name}' from {len(args.corpus_files)} files...") + corpus_result = build_corpus(args.corpus_files, args.corpus_name) + + output_dir = './results' + os.makedirs(output_dir, exist_ok=True) + safe_name = args.corpus_name.replace(" ", "_").replace("/", "_") + output_file = f'{output_dir}/{safe_name}_structure.json' + + with open(output_file, 'w', encoding='utf-8') as f: + json.dump(corpus_result, f, indent=2, ensure_ascii=False) + + print(f"Corpus structure saved to: {output_file}") + import sys + sys.exit(0) + if args.flash: args.mode = 'flash' # 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") + paths_provided = sum(x is not None for x in [args.pdf_path, args.md_path, args.docx_path, args.html_path, args.txt_path]) + if paths_provided == 0: + raise ValueError("One of --pdf_path, --md_path, --docx_path, --html_path, or --txt_path must be specified") + if paths_provided > 1: + raise ValueError("Only one file path can be specified") + + if args.docx_path or args.html_path or args.txt_path: + from pageindex.format_converter import convert_to_markdown + source_path = args.docx_path or args.html_path or args.txt_path + + if args.docx_path and not args.docx_path.lower().endswith('.docx'): + raise ValueError("DOCX file must have .docx extension") + if args.html_path and not args.html_path.lower().endswith(('.html', '.htm')): + raise ValueError("HTML file must have .html or .htm extension") + if args.txt_path and not args.txt_path.lower().endswith('.txt'): + raise ValueError("TXT file must have .txt extension") + if not os.path.isfile(source_path): + raise ValueError(f"File not found: {source_path}") + + print(f"Converting {source_path} to Markdown...") + md_content = convert_to_markdown(source_path) + + output_dir = './results' + os.makedirs(output_dir, exist_ok=True) + filename = os.path.splitext(os.path.basename(source_path))[0] + tmp_md_path = os.path.join(output_dir, f"{filename}_converted.md") + + with open(tmp_md_path, 'w', encoding='utf-8') as f: + f.write(md_content) + + print(f"Converted markdown saved to: {tmp_md_path}") + # Route to the markdown logic + args.md_path = tmp_md_path + if args.optimize is not None and not (args.pdf_path and args.mode == 'flash'): raise ValueError("--optimize requires Flash mode with --pdf_path") if args.optimize is None: