mirage-benchmark 1.0.4__py3-none-any.whl

mirage/pipeline/pdf_processor.py

@@ -0,0 +1,598 @@
+"""
+Docling-Based Document Processing and Table Image Extraction
+
+Supports: single PDF/HTML file, folder of documents, or zip file containing documents.
+Supported formats: PDF, HTML, XHTML (via Docling library)
+Configuration via config.yaml under pdf_processing section.
+"""
+
+import logging
+import time
+from pathlib import Path
+import requests
+import base64
+from io import BytesIO
+from PIL import Image
+import matplotlib.pyplot as plt
+import os
+import json
+import zipfile
+import tempfile
+import shutil
+import yaml
+
+import pandas as pd
+import torch
+
+from docling_core.types.doc import ImageRefMode, PictureItem, TableItem
+from docling_core.types.doc.document import DescriptionAnnotation
+from docling.datamodel.base_models import InputFormat
+from docling.datamodel.pipeline_options import (
+    PdfPipelineOptions,
+    AcceleratorDevice,
+    AcceleratorOptions,
+    PictureDescriptionApiOptions,
+    EasyOcrOptions
+)
+from docling.document_converter import DocumentConverter, PdfFormatOption, HTMLFormatOption
+from docling.utils.export import generate_multimodal_pages
+
+
+def load_config(config_path="config.yaml"):
+    """Load configuration from YAML file."""
+    with open(config_path, 'r') as f:
+        return yaml.safe_load(f)
+
+# Load config
+CONFIG = load_config()
+PDF_CONFIG = CONFIG.get("pdf_processing", {})
+BACKEND_CONFIG = CONFIG.get("backend", {})
+
+# --- Configuration from config.yaml ---
+IMAGE_RESOLUTION_SCALE = PDF_CONFIG.get("image_resolution_scale", 2.0)
+INPUT_PATH = PDF_CONFIG.get("input_path", "data/documents")
+OUTPUT_DIR = Path(PDF_CONFIG.get("output_dir", "trials/pdf2md/output"))
+MODEL_NAME = PDF_CONFIG.get("model_name", "qwen2.5vl:32b")
+NUM_THREADS = PDF_CONFIG.get("num_threads", 14)
+CUDA_DEVICE_ID = PDF_CONFIG.get("cuda_device_id", 1)
+
+# API Configuration from backend settings
+# API configuration from backend config or environment
+API_KEY_FILE = os.environ.get("GEMINI_API_KEY_PATH", os.path.expanduser("~/.config/gemini/api_key.txt"))
+API_URL = os.environ.get("LLM_API_URL", "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent")
+
+from prompt import PROMPTS_DESC
+
+logging.basicConfig(level=logging.INFO)
+_log = logging.getLogger(__name__)
+
+
+# Supported document extensions
+SUPPORTED_EXTENSIONS = {'.pdf', '.html', '.htm', '.xhtml'}
+
+def get_input_format(file_path: Path) -> InputFormat:
+    """Get the InputFormat enum for a given file path."""
+    ext = file_path.suffix.lower()
+    if ext == '.pdf':
+        return InputFormat.PDF
+    elif ext in {'.html', '.htm', '.xhtml'}:
+        return InputFormat.HTML
+    else:
+        raise ValueError(f"Unsupported file extension: {ext}")
+
+def collect_input_files(input_path):
+    """
+    Collect document files (PDF, HTML) from input path.
+    
+    Args:
+        input_path: str - path to a single document, folder of documents, or zip file
+        
+    Returns:
+        tuple: (list of document paths, temp_dir or None)
+               temp_dir is returned if zip was extracted (caller should clean up)
+    """
+    input_path = Path(input_path)
+    temp_dir = None
+    doc_files = []
+    
+    if not input_path.exists():
+        raise FileNotFoundError(f"Input path does not exist: {input_path}")
+    
+    if input_path.is_file():
+        if input_path.suffix.lower() in SUPPORTED_EXTENSIONS:
+            # Single document file
+            doc_files = [input_path]
+        elif input_path.suffix.lower() == '.zip':
+            # Zip file - extract to temp directory
+            temp_dir = tempfile.mkdtemp(prefix="doc_extract_")
+            _log.info(f"Extracting zip file to: {temp_dir}")
+            with zipfile.ZipFile(input_path, 'r') as zf:
+                zf.extractall(temp_dir)
+            # Recursively find all supported documents in extracted content
+            for ext in SUPPORTED_EXTENSIONS:
+                doc_files.extend(Path(temp_dir).rglob(f"*{ext}"))
+        else:
+            raise ValueError(f"Unsupported file type: {input_path.suffix}. Supported: {SUPPORTED_EXTENSIONS}")
+    elif input_path.is_dir():
+        # Folder - recursively find all supported documents
+        for ext in SUPPORTED_EXTENSIONS:
+            doc_files.extend(input_path.rglob(f"*{ext}"))
+    else:
+        raise ValueError(f"Invalid input path: {input_path}")
+    
+    # Sort by file size (smallest first) for faster initial feedback
+    doc_files = sorted(doc_files, key=lambda p: p.stat().st_size)
+    
+    # Log counts by format
+    pdf_count = sum(1 for f in doc_files if f.suffix.lower() == '.pdf')
+    html_count = sum(1 for f in doc_files if f.suffix.lower() in {'.html', '.htm', '.xhtml'})
+    _log.info(f"Found {len(doc_files)} document files to process (PDF: {pdf_count}, HTML: {html_count})")
+    return doc_files, temp_dir
+
+# Backward compatibility alias
+def collect_pdf_files(input_path):
+    """Backward compatibility wrapper for collect_input_files."""
+    return collect_input_files(input_path)
+
+def motormaven_vlm_options():
+    """Configure PictureDescriptionApiOptions for motormaven endpoint"""
+    # Load API key from file
+    with open(API_KEY_FILE, 'r') as f:
+        api_key = f.read().strip()
+    
+    options = PictureDescriptionApiOptions(
+        url=API_URL,
+        params=dict(
+            model=MODEL_NAME,
+        ),
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        },
+        prompt=PROMPTS_DESC["image"],
+        timeout=120,
+        retries=10,
+    )
+    return options
+
+def check_cuda_memory(device_id=0):
+    if torch.cuda.is_available():
+        total_memory = torch.cuda.get_device_properties(device_id).total_memory
+        reserved_memory = torch.cuda.memory_reserved(device_id)
+        allocated_memory = torch.cuda.memory_allocated(device_id)
+        free_memory = reserved_memory - allocated_memory
+        unreserved_memory = total_memory - reserved_memory
+
+        print(f"Total memory:     {total_memory / 1024 ** 3:.2f} GiB")
+        print(f"Reserved memory:  {reserved_memory / 1024 ** 3:.2f} GiB")
+        print(f"Allocated memory: {allocated_memory / 1024 ** 3:.2f} GiB")
+        print(f"Free memory:      {free_memory / 1024 ** 3:.2f} GiB (within reserved)")
+        print(f"Unreserved memory: {unreserved_memory / 1024 ** 3:.2f} GiB")
+    else:
+        print("CUDA is not available.")
+
+
+def is_bbox_inside(inner_bbox, outer_bbox):
+    """Check if inner_bbox is inside outer_bbox. Both are BoundingBox objects with l, t, r, b.
+    Coordinate origin is bottom-left: b < t (bottom has lower y than top)."""
+    return (inner_bbox.l >= outer_bbox.l and 
+            inner_bbox.r <= outer_bbox.r and 
+            inner_bbox.b >= outer_bbox.b and 
+            inner_bbox.t <= outer_bbox.t)
+
+
+def get_pictures_inside_tables(conv_res):
+    """Return set of picture indices that are inside tables."""
+    pictures_to_skip = set()
+    
+    # Build table info: {page_no: [bbox1, bbox2, ...]}
+    table_bboxes_by_page = {}
+    for table in conv_res.document.tables:
+        if table.prov:
+            for prov in table.prov:
+                page_no = prov.page_no
+                bbox = prov.bbox
+                if page_no not in table_bboxes_by_page:
+                    table_bboxes_by_page[page_no] = []
+                table_bboxes_by_page[page_no].append(bbox)
+    
+    # Check each picture
+    for i, picture in enumerate(conv_res.document.pictures):
+        if picture.prov:
+            for prov in picture.prov:
+                pic_page = prov.page_no
+                pic_bbox = prov.bbox
+                # Check if this picture is inside any table on the same page
+                if pic_page in table_bboxes_by_page:
+                    for table_bbox in table_bboxes_by_page[pic_page]:
+                        if is_bbox_inside(pic_bbox, table_bbox):
+                            pictures_to_skip.add(i)
+                            print(f"Picture {i} on page {pic_page} is inside a table - will skip annotation")
+                            break
+    
+    return pictures_to_skip
+
+
+def make_api_call_with_retries(payload, api_url, headers, item_type, item_index, output_dir):
+    """Make API call with retry logic for server errors and response recording"""
+    retry_count = 0
+    while True:
+        try:
+            response = requests.post(api_url, json=payload, headers=headers, timeout=120)
+            
+            if response.status_code == 200:
+                # Response saving disabled - no longer saving individual response files
+                # responses_dir = output_dir / "api_responses"
+                # responses_dir.mkdir(parents=True, exist_ok=True)
+                # response_file = responses_dir / f"{item_type}_{item_index}_response.json"
+                # with open(response_file, 'w') as f:
+                #     json.dump(response.json(), f, indent=2)
+                return response
+            elif 400 <= response.status_code < 500:
+                # Client error - skip
+                print(f"Client error for {item_type} {item_index}: HTTP {response.status_code} - skipping")
+                return None
+            elif response.status_code >= 500:
+                # Server error - retry with exponential backoff
+                retry_count += 1
+                wait_time = min(60, 2 ** min(retry_count, 6))  # Cap at 60 seconds
+                print(f"Server error for {item_type} {item_index}: HTTP {response.status_code} - retrying in {wait_time}s (attempt {retry_count})")
+                time.sleep(wait_time)
+                continue
+            else:
+                print(f"Unexpected status for {item_type} {item_index}: HTTP {response.status_code} - skipping")
+                return None
+                
+        except Exception as e:
+            # Network/connection errors - retry
+            retry_count += 1
+            wait_time = min(60, 2 ** min(retry_count, 6))
+            print(f"Network error for {item_type} {item_index}: {str(e)} - retrying in {wait_time}s (attempt {retry_count})")
+            time.sleep(wait_time)
+            continue
+
+def annotate_items_with_images(conv_res, model_name=MODEL_NAME, api_url=API_URL, max_tokens=1000, use_batch=True):
+    """Annotate both pictures and tables that don't have annotations using their base64 image data.
+    
+    Args:
+        conv_res: Conversion result from docling
+        model_name: VLM model name
+        api_url: API endpoint URL
+        max_tokens: Max tokens for response
+        use_batch: If True, use batch processing for faster annotation
+        
+    Returns:
+        Set of picture indices that are inside tables (skipped).
+    """
+    
+    # Get pictures that are inside tables (skip these)
+    pictures_to_skip = get_pictures_inside_tables(conv_res)
+    
+    if use_batch:
+        return _annotate_items_batch(conv_res, model_name, pictures_to_skip)
+    else:
+        return _annotate_items_sequential(conv_res, model_name, api_url, pictures_to_skip)
+
+
+def _annotate_items_batch(conv_res, model_name, pictures_to_skip):
+    """Batch annotate pictures and tables using async batch processing."""
+    from call_llm import batch_call_vlm_base64
+    
+    # Collect all items that need annotation
+    batch_requests = []  # List of (prompt, base64, mime_type)
+    item_refs = []  # Track which item each request corresponds to: ('picture'|'table', index, item)
+    
+    # Collect pictures
+    for i, item in enumerate(conv_res.document.pictures):
+        if i in pictures_to_skip:
+            print(f"Picture {i}: Skipping - inside a table")
+            continue
+        if not item.annotations:
+            try:
+                base64_data = str(item.image.uri).split(',')[1] if item.image and item.image.uri else None
+                if not base64_data:
+                    print(f"Picture {i}: No base64 data available")
+                    continue
+                print(f"Picture {i}: Queued for batch (base64 len: {len(base64_data)})")
+                batch_requests.append((PROMPTS_DESC["image"], base64_data, "image/png"))
+                item_refs.append(('picture', i, item))
+            except Exception as e:
+                print(f"Error preparing picture {i}: {str(e)}")
+    
+    # Collect tables
+    for i, item in enumerate(conv_res.document.tables):
+        if not item.annotations:
+            try:
+                base64_data = str(item.image.uri.path).split(',')[1] if item.image and item.image.uri else None
+                if not base64_data:
+                    print(f"Table {i}: No base64 data available")
+                    continue
+                print(f"Table {i}: Queued for batch (base64 len: {len(base64_data)})")
+                batch_requests.append((PROMPTS_DESC["table"], base64_data, "image/png"))
+                item_refs.append(('table', i, item))
+            except Exception as e:
+                print(f"Error preparing table {i}: {str(e)}")
+    
+    if not batch_requests:
+        print("No items to annotate")
+        return pictures_to_skip
+    
+    # Execute batch VLM call
+    print(f"\n⚡ Batch annotating {len(batch_requests)} items...")
+    responses = batch_call_vlm_base64(batch_requests, show_progress=True)
+    
+    # Apply responses to items
+    success_count = 0
+    for (item_type, idx, item), response in zip(item_refs, responses):
+        if response and not response.startswith("ERROR:"):
+            annotation = DescriptionAnnotation(
+                kind='description',
+                text=response,
+                provenance=model_name
+            )
+            item.annotations.append(annotation)
+            print(f"✅ Added annotation to {item_type} {idx}")
+            success_count += 1
+        else:
+            print(f"❌ Failed to annotate {item_type} {idx}: {response[:100] if response else 'No response'}")
+    
+    print(f"\n📊 Batch annotation complete: {success_count}/{len(batch_requests)} successful")
+    return pictures_to_skip
+
+
+def _annotate_items_sequential(conv_res, model_name, api_url, pictures_to_skip):
+    """Sequential annotation (original implementation, kept as fallback)."""
+    # Load API key for authentication
+    with open(API_KEY_FILE, 'r') as f:
+        api_key = f.read().strip()
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+    
+    # Process pictures
+    for i, item in enumerate(conv_res.document.pictures):
+        if i in pictures_to_skip:
+            print(f"Picture {i}: Skipping - inside a table")
+            continue
+        if not item.annotations:
+            try:
+                base64_data = str(item.image.uri).split(',')[1] if item.image and item.image.uri else None
+                if not base64_data:
+                    print(f"Picture {i}: No base64 data available")
+                    continue
+                    
+                print(f"Picture {i}: Using base64 data, length: {len(base64_data)}")
+                
+                payload = {
+                    "model": model_name,
+                    "messages": [{
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": PROMPTS_DESC["image"]},
+                            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_data}"}}
+                        ]
+                    }]
+                }
+                
+                response = make_api_call_with_retries(payload, api_url, headers, "picture", i, OUTPUT_DIR)
+                
+                if response and response.status_code == 200:
+                    result = response.json()
+                    description_text = result.get('choices', [{}])[0].get('message', {}).get('content', 'No description available')
+                    annotation = DescriptionAnnotation(kind='description', text=description_text, provenance=model_name)
+                    item.annotations.append(annotation)
+                    print(f"Added annotation to picture {i}")
+                else:
+                    print(f"Failed to annotate picture {i}: HTTP {response.status_code if response else 'No response'}")
+            except Exception as e:
+                print(f"Error annotating picture {i}: {str(e)}")
+    
+    # Process tables
+    for i, item in enumerate(conv_res.document.tables):
+        if not item.annotations:
+            try:
+                base64_data = str(item.image.uri.path).split(',')[1] if item.image and item.image.uri else None
+                if not base64_data:
+                    print(f"Table {i}: No base64 data available")
+                    continue
+                    
+                print(f"Table {i}: Using base64 data, length: {len(base64_data)}")
+                
+                payload = {
+                    "model": model_name,
+                    "messages": [{
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": PROMPTS_DESC["table"]},
+                            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_data}"}}
+                        ]
+                    }]
+                }
+                
+                response = make_api_call_with_retries(payload, api_url, headers, "table", i, OUTPUT_DIR)
+                
+                if response and response.status_code == 200:
+                    result = response.json()
+                    description_text = result.get('choices', [{}])[0].get('message', {}).get('content', 'No description available')
+                    annotation = DescriptionAnnotation(kind='description', text=description_text, provenance=model_name)
+                    item.annotations.append(annotation)
+                    print(f"Added annotation to table {i}")
+                else:
+                    print(f"Failed to annotate table {i}: HTTP {response.status_code if response else 'No response'}")
+            except Exception as e:
+                print(f"Error annotating table {i}: {str(e)}")
+    
+    return pictures_to_skip
+
+
+def configure_pipeline_options(model_name:str="granite3.2-vision:latest", cuda_device_id:int=None):
+    pipeline_options = PdfPipelineOptions()
+    pipeline_options.enable_remote_services = True
+    pipeline_options.images_scale = IMAGE_RESOLUTION_SCALE
+    pipeline_options.generate_page_images = False # True For debugging
+    pipeline_options.generate_parsed_pages = False  # True For debugging
+    pipeline_options.generate_picture_images = True
+    pipeline_options.do_picture_classification = False  # False to Avoid CUDA OOM
+    pipeline_options.do_picture_description = True   # False Avoid CUDA OOM
+    pipeline_options.do_ocr = True
+    pipeline_options.ocr_options = EasyOcrOptions()
+    pipeline_options.do_code_enrichment = True
+    pipeline_options.do_formula_enrichment = True
+    pipeline_options.do_table_structure = True
+    pipeline_options.table_structure_options.do_cell_matching = True
+    pipeline_options.generate_table_images = True
+    
+    print(f"DEBUG: Pipeline options - generate_table_images: {pipeline_options.generate_table_images}")
+    print(f"DEBUG: Pipeline options - do_table_structure: {pipeline_options.do_table_structure}")
+    
+    # Set accelerator with specific CUDA device if provided
+    if cuda_device_id is not None:
+        print(f"DEBUG: Using CUDA device {cuda_device_id}")
+        pipeline_options.accelerator_options = AcceleratorOptions(
+            num_threads=NUM_THREADS, device=f"cuda:{cuda_device_id}"
+        )
+    else:
+        pipeline_options.accelerator_options = AcceleratorOptions(
+            num_threads=NUM_THREADS, device=AcceleratorDevice.AUTO
+        )
+
+    ### Set picture description API options for motormaven endpoint
+    pipeline_options.picture_description_options = motormaven_vlm_options()
+    
+    print("DEBUG: Using manual table annotation (no built-in table description options)")
+    
+    return pipeline_options
+
+def process_single_document(doc_path, doc_converter, output_dir, is_pdf=True):
+    """Process a single document file (PDF or HTML) and save outputs.
+    
+    Args:
+        doc_path: Path to the document file
+        doc_converter: DocumentConverter instance
+        output_dir: Output directory path
+        is_pdf: Whether the document is a PDF (affects annotation behavior)
+    """
+    doc_start = time.time()
+    
+    try:
+        conv_res = doc_converter.convert(str(doc_path))
+        doc_filename = conv_res.input.file.stem
+        
+        # Create per-document output directory
+        doc_output_dir = output_dir / doc_filename
+        doc_output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Annotate images and tables (only for PDFs with complex layouts)
+        # HTML files typically have simpler structure
+        if is_pdf:
+            pictures_to_skip = annotate_items_with_images(conv_res)
+            
+            # Create folders for images
+            tables_dir = doc_output_dir / "tables"
+            tables_dir.mkdir(parents=True, exist_ok=True)
+            
+            # Save table images
+            table_counter = 0
+            for element, _level in conv_res.document.iterate_items():
+                if isinstance(element, TableItem):
+                    table_counter += 1
+                    element_image_filename = tables_dir / f"{doc_filename}-table-{table_counter}.png"
+                    _log.info(f"Table {element.self_ref} - Caption: {element.caption_text(doc=conv_res.document)}")
+                    try:
+                        with element_image_filename.open("wb") as fp:
+                            element.get_image(conv_res.document).save(fp, "PNG")
+                    except Exception as e:
+                        _log.warning(f"Could not save table image: {e}")
+        
+        # Create artifacts directory for referenced images
+        artifacts_dir = doc_output_dir / "ref_artifacts"
+        artifacts_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Save markdown with externally referenced pictures
+        md_filename = doc_output_dir / f"{doc_filename}_ref.md"
+        conv_res.document.save_as_markdown(md_filename, image_mode=ImageRefMode.REFERENCED, artifacts_dir=artifacts_dir)
+        
+        elapsed = time.time() - doc_start
+        file_type = "PDF" if is_pdf else "HTML"
+        _log.info(f"✅ Processed {file_type} {doc_filename} in {elapsed:.1f}s")
+        return True, doc_filename, elapsed
+        
+    except Exception as e:
+        elapsed = time.time() - doc_start
+        _log.error(f"❌ Failed to process {doc_path.name}: {str(e)}")
+        return False, doc_path.name, elapsed
+
+# Backward compatibility alias
+def process_single_pdf(pdf_path, doc_converter, output_dir):
+    """Backward compatibility wrapper for process_single_document."""
+    return process_single_document(pdf_path, doc_converter, output_dir, is_pdf=True)
+
+
+def create_multi_format_converter(cuda_device_id=None):
+    """Create a DocumentConverter that supports both PDF and HTML formats."""
+    # PDF pipeline options (full processing)
+    pdf_pipeline_options = configure_pipeline_options(cuda_device_id=cuda_device_id)
+    
+    # Initialize document converter with both PDF and HTML support
+    doc_converter = DocumentConverter(
+        format_options={
+            InputFormat.PDF: PdfFormatOption(pipeline_options=pdf_pipeline_options),
+            InputFormat.HTML: HTMLFormatOption(),  # HTML uses default options
+        }
+    )
+    return doc_converter
+
+
+if __name__ == "__main__":
+    start_time = time.time()
+    
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+    check_cuda_memory(CUDA_DEVICE_ID)
+    
+    # Collect document files (PDF and HTML) from input path
+    doc_files, temp_dir = collect_input_files(INPUT_PATH)
+    
+    if not doc_files:
+        _log.error("No document files found to process")
+        exit(1)
+    
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    
+    # Initialize document converter with multi-format support
+    doc_converter = create_multi_format_converter(cuda_device_id=CUDA_DEVICE_ID)
+    
+    # Process all documents
+    results = []
+    for i, doc_path in enumerate(doc_files, 1):
+        is_pdf = doc_path.suffix.lower() == '.pdf'
+        file_type = "PDF" if is_pdf else "HTML"
+        
+        _log.info(f"\n{'='*60}")
+        _log.info(f"Processing {file_type} {i}/{len(doc_files)}: {doc_path.name}")
+        _log.info(f"{'='*60}")
+        
+        success, name, elapsed = process_single_document(doc_path, doc_converter, OUTPUT_DIR, is_pdf=is_pdf)
+        results.append((success, name, elapsed, file_type))
+    
+    # Cleanup temp directory if created from zip
+    if temp_dir:
+        _log.info(f"Cleaning up temp directory: {temp_dir}")
+        shutil.rmtree(temp_dir, ignore_errors=True)
+    
+    # Summary
+    total_time = time.time() - start_time
+    successful = sum(1 for r in results if r[0])
+    failed = len(results) - successful
+    pdf_count = sum(1 for r in results if r[3] == "PDF")
+    html_count = sum(1 for r in results if r[3] == "HTML")
+    
+    _log.info(f"\n{'='*60}")
+    _log.info(f"PROCESSING COMPLETE")
+    _log.info(f"{'='*60}")
+    _log.info(f"Total documents: {len(results)} (PDF: {pdf_count}, HTML: {html_count})")
+    _log.info(f"Successful: {successful}")
+    _log.info(f"Failed: {failed}")
+    _log.info(f"Total time: {total_time:.1f}s")
+    _log.info(f"Output directory: {OUTPUT_DIR}")