# pyright: ignore #!/usr/bin/env python3 """ Main converter module using ProcessPoolExecutor for optimized GPU utilization. """ import tempfile from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import fitz # PyMuPDF from loguru import logger from .pipelines import init_converters def process_page_worker(args: Tuple[str, int, int, str, str]) -> Optional[Dict[str, Any]]: """ Process a single page - runs in separate process. Each process initializes its own converters for separate GPU contexts. """ page_file_str, page_num, total_pages, output_dir_str, original_pdf_name = args page_file = Path(page_file_str) output_dir = Path(output_dir_str) try: # Initialize converters in this process (separate GPU context) regular_converter, vlm_converter = init_converters() # Phase 1: Regular processing for table detection regular_result = regular_converter.convert(str(page_file)) if regular_result is None or not regular_result.document.pages: logger.error(f"Page {page_num}/{total_pages}: Phase 1 failed") return None # Analyze for tables document = regular_result.document page_tables = [ table for table in document.tables if table.prov and len(table.prov) > 0 and table.prov[0].page_no == 1 ] has_tables = len(page_tables) > 0 table_count = len(page_tables) # Log once after Phase 1 if has_tables: logger.info(f"Page {page_num}/{total_pages}: Processing") else: logger.info(f"Page {page_num}/{total_pages}: Processing (no tables)") # Phase 2: Conditional VLM processing if has_tables: vlm_result = vlm_converter.convert(str(page_file)) if vlm_result is not None: final_result = vlm_result else: logger.debug(f"VLM processing failed for page {page_num}, using regular result") final_result = regular_result else: final_result = regular_result # Generate markdown markdown_content = final_result.document.export_to_markdown() # Save output base_name = Path(original_pdf_name).stem output_filename = f"{base_name}_page_{page_num:03d}.md" output_path = output_dir / output_filename with open(output_path, "w", encoding="utf-8") as f: f.write(markdown_content) return { "page_num": page_num, "output_path": str(output_path), "has_tables": has_tables, "table_count": table_count, } except Exception as e: logger.error(f"Page {page_num}/{total_pages}: Failed - {str(e)}") return None class PDFConverter: """ High-performance PDF to Markdown converter using ProcessPoolExecutor. """ def __init__(self, output_dir: str = "output", max_workers: int = 2, debug_mode: bool = False): """ Initialize the PDF converter. Args: output_dir: Directory for output markdown files max_workers: Number of parallel process workers debug_mode: If True, saves extracted PDF pages to debug directory """ self.output_dir = Path(output_dir) self.max_workers = max_workers self.debug_mode = debug_mode self.output_dir.mkdir(parents=True, exist_ok=True) if self.debug_mode: self.debug_dir = Path(output_dir) / "debug_pages" self.debug_dir.mkdir(parents=True, exist_ok=True) logger.info(f"Debug mode enabled - pages will be saved to {self.debug_dir}") def extract_pages(self, pdf_path: Path, temp_dir: Path) -> List[Path]: """Extract individual pages from PDF to temporary files.""" doc = fitz.open(str(pdf_path)) total_pages = len(doc) page_files = [] logger.debug(f"Extracting {total_pages} pages from {pdf_path.name}") for page_num in range(total_pages): new_doc = fitz.open() new_doc.insert_pdf(doc, from_page=page_num, to_page=page_num) page_filename = f"page_{page_num + 1:03d}.pdf" page_path = temp_dir / page_filename new_doc.save(str(page_path)) # Save to debug directory if debug mode is enabled if self.debug_mode: debug_filename = f"{pdf_path.stem}_page_{page_num + 1:03d}.pdf" debug_path = self.debug_dir / debug_filename new_doc.save(str(debug_path)) logger.debug(f"Saved debug PDF: {debug_path}") new_doc.close() page_files.append(page_path) doc.close() if self.debug_mode: logger.info(f"Extracted {total_pages} pages saved to {self.debug_dir}") return page_files def convert(self, pdf_path: Path) -> Dict[str, Any]: """ Convert PDF to Markdown using ProcessPoolExecutor. Args: pdf_path: Path to input PDF file Returns: Dictionary with conversion statistics """ 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": []} with tempfile.TemporaryDirectory() as temp_dir: temp_path = Path(temp_dir) try: # Extract pages page_files = self.extract_pages(pdf_path, temp_path) total_pages = len(page_files) logger.info(f"Processing {total_pages} pages with {self.max_workers} process workers") # Prepare arguments for each page page_args = [ (str(page_file), i + 1, total_pages, str(self.output_dir), pdf_path.name) for i, page_file in enumerate(page_files) ] # Process pages using ProcessPoolExecutor successful_count = 0 failed_count = 0 output_files = [] with ProcessPoolExecutor(max_workers=self.max_workers) as executor: # Submit all tasks future_to_args = {executor.submit(process_page_worker, args): args for args in page_args} # Process as completed for future in as_completed(future_to_args): result = future.result() if result: successful_count += 1 output_files.append(result["output_path"]) else: failed_count += 1 logger.info(f"Conversion complete: {successful_count} successful, {failed_count} failed") return { "success": successful_count, "failed": failed_count, "total_pages": total_pages, "output_files": output_files, } except Exception as e: logger.error(f"Critical error during conversion: {str(e)}") return {"success": 0, "failed": 0, "total_pages": 0, "output_files": []}