# pyright: ignore import argparse import asyncio from concurrent.futures import ThreadPoolExecutor from pathlib import Path import sys from typing import Optional from loguru import logger from docling.document_converter import DocumentConverter, PdfFormatOption from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import VlmPipelineOptions from docling.datamodel.vlm_model_specs import GRANITEDOCLING_TRANSFORMERS from docling.pipeline.vlm_pipeline import VlmPipeline class PageByPageConverter: 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) # Configure VLM pipeline with granite_docling model vlm_pipeline_options = VlmPipelineOptions( vlm_options=GRANITEDOCLING_TRANSFORMERS ) # Configure PDF format to use VLM pipeline pdf_format_options = PdfFormatOption( pipeline_cls=VlmPipeline, pipeline_options=vlm_pipeline_options ) self.converter = DocumentConverter( format_options={ InputFormat.PDF: pdf_format_options } ) logger.info("Initialized converter with VLM pipeline using granite_docling model") logger.info(f"Output directory: {self.output_dir.absolute()}") logger.info(f"Max concurrent pages: {max_concurrent}") async def process_page(self, page_data: dict, page_num: int, total_pages: int, pdf_name: str) -> Optional[str]: """Process a single page and return the output file path if successful""" async with self.semaphore: logger.info(f"Starting page {page_num + 1}/{total_pages} of {pdf_name}") try: # Run the markdown generation in a thread pool to avoid blocking loop = asyncio.get_event_loop() with ThreadPoolExecutor() as executor: markdown_content = await loop.run_in_executor( executor, self._generate_page_markdown, page_data ) # Generate output filename base_name = Path(pdf_name).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}: {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 _generate_page_markdown(self, page_data: dict) -> str: """Generate markdown from a single page data - runs in thread pool""" try: page = page_data['page'] page_num = page_data.get('page_num', 'Unknown') header = f"# Page {page_num} " content_parts = [] # Extract content from the page's body items if hasattr(page, 'items') and page.items: for item in page.items: # Handle different types of content items if hasattr(item, 'text') and item.text: # Text items content_parts.append(item.text) elif hasattr(item, 'content') and item.content: # Content-based items content_parts.append(item.content) # If no items found, try to get content from the page's main text if not content_parts and hasattr(page, 'text') and page.text: content_parts.append(page.text) # If still no content, extract from body if available if not content_parts and hasattr(page, 'body') and page.body: if hasattr(page.body, 'text') and page.body.text: content_parts.append(page.body.text) content = " ".join(content_parts) if content_parts else "No content extracted from this page." # Clean up the content content = content.strip() return header + content except Exception as e: logger.warning(f"Error generating markdown for page {page_data.get('page_num', 'Unknown')}: {str(e)}") page_num = page_data.get('page_num', 'Unknown') return f"# Page {page_num} Error processing this page: {str(e)}" async def convert_document(self, pdf_path: Path) -> dict: """Convert all pages of a PDF document concurrently""" logger.info(f"Starting 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, convert the entire document once logger.info("Converting document with VLM pipeline...") result = self.converter.convert(str(pdf_path)) total_pages = len(result.document.pages) logger.info(f"Document has {total_pages} pages") # Create tasks for processing each page's markdown tasks = [] for page_num, page in enumerate(result.document.pages): page_data = { 'page': page, 'page_num': page_num + 1, 'document': result.document } task = asyncio.create_task( self.process_page(page_data, page_num, total_pages, pdf_path.name) ) 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="Convert PDF pages to Markdown using Docling VLM pipeline with granite_docling model" ) 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 = PageByPageConverter( 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()))