# pyright: ignore import argparse import asyncio from concurrent.futures import ThreadPoolExecutor from pathlib import Path import sys from typing import Optional, Dict, Tuple from loguru import logger from docling.document_converter import DocumentConverter, PdfFormatOption from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import ( PdfPipelineOptions, VlmPipelineOptions, TableStructureOptions, TableFormerMode ) from docling.datamodel.vlm_model_specs import GRANITEDOCLING_TRANSFORMERS from docling.pipeline.vlm_pipeline import VlmPipeline from docling.pipeline.standard_pdf_pipeline import StandardPdfPipeline class TwoPhasePageConverter: def __init__(self, output_dir: str = "output", max_concurrent: int = 2): self.output_dir = Path(output_dir) self.output_dir.mkdir(exist_ok=True) self.max_concurrent = max_concurrent self.semaphore = asyncio.Semaphore(max_concurrent) # Phase 1: Regular Docling pipeline with enhanced table detection regular_pipeline_options = PdfPipelineOptions( do_table_structure=True, table_structure_options=TableStructureOptions( mode=TableFormerMode.ACCURATE, do_cell_matching=True ), images_scale=2.0, do_ocr=True ) regular_format_options = PdfFormatOption( pipeline_cls=StandardPdfPipeline, pipeline_options=regular_pipeline_options ) self.regular_converter = DocumentConverter( format_options={ InputFormat.PDF: regular_format_options } ) # Phase 2: VLM pipeline with granite_docling model vlm_pipeline_options = VlmPipelineOptions( vlm_options=GRANITEDOCLING_TRANSFORMERS ) vlm_format_options = PdfFormatOption( pipeline_cls=VlmPipeline, pipeline_options=vlm_pipeline_options ) self.vlm_converter = DocumentConverter( format_options={ InputFormat.PDF: vlm_format_options } ) logger.info("Initialized two-phase converter:") logger.info(" Phase 1: Regular pipeline with enhanced table detection") logger.info(" Phase 2: VLM pipeline with granite_docling model") logger.info(f" Output directory: {self.output_dir.absolute()}") logger.info(f" Max concurrent pages: {max_concurrent}") def _has_tables(self, page) -> Tuple[bool, int]: """Detect if a page contains tables and return table count""" table_count = 0 if hasattr(page, 'items') and page.items: for item in page.items: # Check for table items in the page item_type = type(item).__name__ if 'Table' in item_type: table_count += 1 # Alternative: check if page has table-related content if table_count == 0 and hasattr(page, 'body') and page.body: if hasattr(page.body, 'children'): for child in page.body.children: child_type = type(child).__name__ if 'Table' in child_type: table_count += 1 return table_count > 0, table_count def _generate_markdown_from_page(self, page, document, page_num: int) -> str: """Generate markdown from a single page within a document context""" try: # Create a temporary document with just this page single_page_doc = document.model_copy() single_page_doc.pages = [page] # Export to markdown markdown_content = single_page_doc.export_to_markdown() # Add page header header = f"# Page {page_num} " return header + markdown_content.strip() except Exception as e: logger.warning(f"Error generating markdown for page {page_num}: {str(e)}") return f"# Page {page_num} Error processing this page: {str(e)}" async def process_single_page(self, pdf_path: Path, page_num: int, total_pages: int) -> Optional[str]: """Process a single page through both phases if necessary""" async with self.semaphore: pdf_name = pdf_path.name logger.info(f"Starting page {page_num + 1}/{total_pages} of {pdf_name}") try: # Phase 1: Regular processing with table detection logger.debug(f"Phase 1: Regular processing for page {page_num + 1}") loop = asyncio.get_event_loop() with ThreadPoolExecutor() as executor: # Process with regular pipeline regular_result = await loop.run_in_executor( executor, self._process_with_regular_pipeline, pdf_path, page_num ) if regular_result is None: logger.error(f"Phase 1 failed for page {page_num + 1}") return None page = regular_result['page'] document = regular_result['document'] # Check for tables has_tables, table_count = self._has_tables(page) if has_tables: logger.info(f"Page {page_num + 1}: Found {table_count} table(s) - using VLM processing") # Phase 2: VLM processing for table-rich content with ThreadPoolExecutor() as executor: vlm_result = await loop.run_in_executor( executor, self._process_with_vlm_pipeline, pdf_path, page_num ) if vlm_result is not None: # Use VLM result markdown_content = self._generate_markdown_from_page( vlm_result['page'], vlm_result['document'], page_num + 1 ) processing_method = "VLM (tables detected)" else: logger.warning(f"Phase 2 VLM failed for page {page_num + 1}, falling back to regular result") # Fallback to regular result markdown_content = self._generate_markdown_from_page( page, document, page_num + 1 ) processing_method = "Regular (VLM fallback)" else: logger.info(f"Page {page_num + 1}: No tables detected - using regular processing") # Use regular result markdown_content = self._generate_markdown_from_page( page, document, page_num + 1 ) processing_method = "Regular (no tables)" # Generate output file base_name = pdf_path.stem output_filename = f"{base_name}_page_{page_num + 1:03d}.md" output_path = self.output_dir / output_filename # Write markdown content with open(output_path, "w", encoding="utf-8") as f: f.write(markdown_content) logger.success(f"Completed page {page_num + 1}/{total_pages} ({processing_method}): {output_path}") return str(output_path) except Exception as e: logger.error(f"Failed to process page {page_num + 1}/{total_pages} of {pdf_name}: {str(e)}") return None def _process_with_regular_pipeline(self, pdf_path: Path, page_num: int) -> Optional[Dict]: """Process single page with regular pipeline""" try: # For now, we process the entire document and extract the page # This could be optimized to process only the specific page result = self.regular_converter.convert(str(pdf_path)) if page_num < len(result.document.pages): return { 'page': result.document.pages[page_num], 'document': result.document } else: raise IndexError(f"Page {page_num + 1} not found in document with {len(result.document.pages)} pages") except Exception as e: logger.error(f"Regular pipeline processing failed: {str(e)}") return None def _process_with_vlm_pipeline(self, pdf_path: Path, page_num: int) -> Optional[Dict]: """Process single page with VLM pipeline""" try: result = self.vlm_converter.convert(str(pdf_path)) if page_num < len(result.document.pages): return { 'page': result.document.pages[page_num], 'document': result.document } else: raise IndexError(f"Page {page_num + 1} not found in document with {len(result.document.pages)} pages") except Exception as e: logger.error(f"VLM pipeline processing failed: {str(e)}") return None async def convert_document(self, pdf_path: Path) -> dict: """Convert all pages of a PDF document using two-phase approach""" logger.info(f"Starting two-phase conversion of {pdf_path.name}") if not pdf_path.exists(): logger.error(f"File not found: {pdf_path}") return {"success": 0, "failed": 0, "files": []} try: # First, get the total number of pages using regular pipeline logger.info("Analyzing document to determine page count...") temp_result = self.regular_converter.convert(str(pdf_path)) total_pages = len(temp_result.document.pages) logger.info(f"Document has {total_pages} pages") # Create tasks for processing each page tasks = [] for page_num in range(total_pages): task = asyncio.create_task( self.process_single_page(pdf_path, page_num, total_pages) ) tasks.append(task) # Wait for all pages to complete logger.info(f"Processing {total_pages} pages with {self.max_concurrent} concurrent workers...") results = await asyncio.gather(*tasks, return_exceptions=True) # Count successes and failures successful_files = [r for r in results if r is not None and not isinstance(r, Exception)] failed_count = total_pages - len(successful_files) logger.info(f"Conversion complete: {len(successful_files)} successful, {failed_count} failed") return { "success": len(successful_files), "failed": failed_count, "files": successful_files } except Exception as e: logger.error(f"Critical error during document conversion: {str(e)}") return {"success": 0, "failed": 0, "files": []} async def main(): parser = argparse.ArgumentParser( description="Two-phase PDF to Markdown converter: Regular processing for text, VLM for tables" ) parser.add_argument("pdf_path", help="Path to the PDF file to convert") parser.add_argument( "--output-dir", "-o", default="output", help="Output directory for markdown files (default: output)" ) parser.add_argument( "--concurrent", "-c", type=int, default=2, help="Number of pages to process concurrently (default: 2)" ) parser.add_argument( "--verbose", "-v", action="store_true", help="Enable verbose logging" ) args = parser.parse_args() # Configure logging logger.remove() # Remove default handler if args.verbose: logger.add(sys.stderr, level="DEBUG", format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}") else: logger.add(sys.stderr, level="INFO", format="{time:HH:mm:ss} | {level} | {message}") pdf_path = Path(args.pdf_path) # Create converter and process document converter = TwoPhasePageConverter( output_dir=args.output_dir, max_concurrent=args.concurrent ) results = await converter.convert_document(pdf_path) # Print summary if results["success"] > 0: logger.success(f"Successfully converted {results['success']} pages") if args.verbose: for file_path in results["files"]: logger.info(f" -> {file_path}") if results["failed"] > 0: logger.warning(f"{results['failed']} pages failed to convert") return 0 if results["success"] > 0 else 1 if __name__ == "__main__": sys.exit(asyncio.run(main()))