granite_docling.py

"""
Granite Docling 258M Implementation

This module provides an interface to the IBM Granite Docling 258M model
for document processing and conversion tasks.
"""

import os
import logging
import time
from pathlib import Path
from typing import Union, Optional, Dict, Any

from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import (
    PdfPipelineOptions,
    VlmPipelineOptions,
    ResponseFormat,
    AcceleratorDevice,
    vlm_model_specs
)
from docling.pipeline.vlm_pipeline import VlmPipeline

# Additional imports for fast document analysis
try:
    import fitz  # PyMuPDF for fast PDF metadata extraction
    PYMUPDF_AVAILABLE = True
except ImportError:
    PYMUPDF_AVAILABLE = False

try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class GraniteDocling:
    """
    A wrapper class for the IBM Granite Docling 258M model.

    This class provides an easy-to-use interface for document processing
    using the Granite Docling model through the Docling framework.
    """

    def __init__(
        self,
        model_type: str = "transformers",
        artifacts_path: Optional[str] = None
    ):
        """
        Initialize the Granite Docling processor.

        Args:
            model_type: Model type - "transformers" or "mlx"
            artifacts_path: Path to cached model artifacts
        """
        self.model_type = model_type.lower()
        self.artifacts_path = artifacts_path

        # Choose the appropriate model configuration
        if self.model_type == "mlx":
            self.vlm_model = vlm_model_specs.GRANITEDOCLING_MLX
        else:
            self.vlm_model = vlm_model_specs.GRANITEDOCLING_TRANSFORMERS

        # Initialize the document converter
        self._setup_converter()

    def _setup_converter(self):
        """Set up the document converter with Granite Docling configuration."""

        # Set up VLM pipeline options using the pre-configured Granite Docling model
        pipeline_options = VlmPipelineOptions(vlm_options=self.vlm_model)

        # Configure PDF processing options
        pdf_options = PdfFormatOption(
            pipeline_cls=VlmPipeline,
            pipeline_options=pipeline_options,
        )

        # If artifacts path is specified, add it to PDF pipeline options
        if self.artifacts_path:
            pdf_pipeline_options = PdfPipelineOptions(artifacts_path=self.artifacts_path)
            pdf_options.pipeline_options = pdf_pipeline_options

        # Initialize the document converter
        self.converter = DocumentConverter(
            format_options={
                InputFormat.PDF: pdf_options,
            }
        )

        logger.info(f"Initialized Granite Docling with model type: {self.model_type}")

    def analyze_document_structure(
        self,
        source: Union[str, Path],
        sample_pages: int = 3,
        max_sample_chars: int = 2000
    ) -> Dict[str, Any]:
        """
        Fast document structure analysis without full conversion.

        This method provides lightweight document insights including:
        - Basic metadata (pages, size, type)
        - Structure detection (headers, tables, images)
        - Content sampling from first few pages
        - Performance optimized for large documents

        Args:
            source: Path to the document
            sample_pages: Number of pages to sample for content analysis
            max_sample_chars: Maximum characters to extract for preview

        Returns:
            Dictionary containing document analysis and structure information
        """
        start_time = time.time()

        try:
            source_path = Path(source)
            logger.info(f"Analyzing document structure: {source}")

            # Initialize analysis result
            analysis_result = {
                "source": str(source),
                "file_name": source_path.name,
                "file_size_mb": round(source_path.stat().st_size / (1024 * 1024), 2),
                "analysis_time_seconds": 0,
                "document_type": source_path.suffix.lower(),
                "structure_detected": {},
                "content_preview": "",
                "metadata_extraction": {},
                "processing_approach": "fast_analysis"
            }

            # PDF-specific fast analysis
            if source_path.suffix.lower() == '.pdf' and PYMUPDF_AVAILABLE:
                analysis_result.update(self._analyze_pdf_structure(source, sample_pages, max_sample_chars))

            # Image file analysis
            elif source_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff'] and PIL_AVAILABLE:
                analysis_result.update(self._analyze_image_structure(source))

            # For other formats, use docling but with limited sampling
            else:
                analysis_result.update(self._analyze_other_format_structure(source, sample_pages, max_sample_chars))

            analysis_result["analysis_time_seconds"] = round(time.time() - start_time, 2)

            logger.info(f"Document analysis completed in {analysis_result['analysis_time_seconds']} seconds")
            return analysis_result

        except Exception as e:
            logger.error(f"Error analyzing document structure {source}: {str(e)}")
            return {
                "source": str(source),
                "error": str(e),
                "analysis_time_seconds": round(time.time() - start_time, 2),
                "processing_approach": "fast_analysis_failed"
            }

    def _analyze_pdf_structure(self, source: Union[str, Path], sample_pages: int, max_sample_chars: int) -> Dict[str, Any]:
        """Fast PDF structure analysis using PyMuPDF."""
        try:
            doc = fitz.open(str(source))
            total_pages = doc.page_count

            # Extract metadata
            metadata = doc.metadata

            # Sample pages for structure analysis
            pages_to_sample = min(sample_pages, total_pages)
            sample_text = ""
            headers_found = []
            tables_detected = 0
            images_detected = 0
            text_density_avg = 0

            for page_num in range(pages_to_sample):
                page = doc[page_num]

                # Get text content
                page_text = page.get_text()
                sample_text += page_text[:max_sample_chars // pages_to_sample] + "\n"

                # Detect structure elements
                text_dict = page.get_text("dict")

                # Count images
                images_detected += len(page.get_images())

                # Estimate text density
                text_density_avg += len(page_text.strip()) / max(1, page.rect.width * page.rect.height) * 10000

                # Simple header detection (large/bold text)
                for block in text_dict.get("blocks", []):
                    if "lines" in block:
                        for line in block["lines"]:
                            for span in line.get("spans", []):
                                text = span.get("text", "").strip()
                                if text and len(text) < 100:  # Potential header
                                    font_size = span.get("size", 12)
                                    font_flags = span.get("flags", 0)

                                    # Check if text looks like a header (large font or bold)
                                    if font_size > 14 or (font_flags & 2**4):  # Bold flag
                                        headers_found.append(text)

                # Simple table detection (look for aligned text patterns)
                tables_detected += self._estimate_tables_in_page_text(page_text)

            doc.close()

            text_density_avg = round(text_density_avg / pages_to_sample, 2) if pages_to_sample > 0 else 0

            return {
                "total_pages": total_pages,
                "pages_analyzed": pages_to_sample,
                "metadata_extraction": {
                    "title": metadata.get("title", ""),
                    "author": metadata.get("author", ""),
                    "creation_date": metadata.get("creationDate", ""),
                    "modification_date": metadata.get("modDate", "")
                },
                "structure_detected": {
                    "headers_found": len(set(headers_found)),
                    "sample_headers": list(set(headers_found))[:5],
                    "estimated_tables": tables_detected,
                    "images_detected": images_detected,
                    "text_density": text_density_avg,
                    "has_text": len(sample_text.strip()) > 50
                },
                "content_preview": sample_text[:max_sample_chars].strip()
            }

        except Exception as e:
            logger.warning(f"PyMuPDF analysis failed, falling back: {e}")
            return self._analyze_other_format_structure(source, sample_pages, max_sample_chars)

    def _analyze_image_structure(self, source: Union[str, Path]) -> Dict[str, Any]:
        """Fast image file analysis."""
        try:
            with Image.open(source) as img:
                return {
                    "total_pages": 1,
                    "pages_analyzed": 1,
                    "metadata_extraction": {
                        "format": img.format,
                        "mode": img.mode,
                        "size": f"{img.size[0]}x{img.size[1]}",
                        "has_exif": bool(getattr(img, '_getexif', lambda: None)())
                    },
                    "structure_detected": {
                        "content_type": "image",
                        "requires_ocr": True,
                        "estimated_text_content": "unknown_until_ocr"
                    },
                    "content_preview": f"Image file: {img.format} format, {img.size[0]}x{img.size[1]} pixels"
                }
        except Exception as e:
            logger.warning(f"Image analysis failed: {e}")
            return {
                "total_pages": 1,
                "structure_detected": {"content_type": "image", "analysis_failed": str(e)},
                "content_preview": "Image analysis failed"
            }

    def _analyze_other_format_structure(self, source: Union[str, Path], sample_pages: int, max_sample_chars: int) -> Dict[str, Any]:
        """Lightweight analysis for other formats using minimal docling processing."""
        try:
            # Use docling but process minimally - just get basic structure
            result = self.converter.convert(source=str(source))
            document = result.document

            # Get basic info without full markdown conversion
            total_pages = len(document.pages) if hasattr(document, 'pages') else 1

            # Sample first few pages only
            pages_to_analyze = min(sample_pages, total_pages)
            sample_content = ""

            if hasattr(document, 'pages'):
                for i in range(pages_to_analyze):
                    if i < len(document.pages):
                        page = document.pages[i]
                        # Get text content from page without full markdown processing
                        if hasattr(page, 'text'):
                            sample_content += str(page.text)[:max_sample_chars // pages_to_analyze] + "\n"

            # If we still don't have content, do a quick markdown export of first portion
            if not sample_content:
                full_content = document.export_to_markdown()
                sample_content = full_content[:max_sample_chars]

            # Quick structure analysis
            headers_found = [line.strip() for line in sample_content.split('\n') if line.strip().startswith('#')]
            table_lines = [line for line in sample_content.split('\n') if '|' in line and line.strip()]

            return {
                "total_pages": total_pages,
                "pages_analyzed": pages_to_analyze,
                "structure_detected": {
                    "headers_found": len(headers_found),
                    "sample_headers": headers_found[:5],
                    "estimated_tables": len([line for line in table_lines if line.count('|') > 1]),
                    "has_markdown_structure": len(headers_found) > 0 or len(table_lines) > 0
                },
                "content_preview": sample_content.strip()
            }

        except Exception as e:
            logger.warning(f"Docling lightweight analysis failed: {e}")
            return {
                "total_pages": 1,
                "structure_detected": {"analysis_method": "file_info_only"},
                "content_preview": "Unable to analyze document structure"
            }

    def _estimate_tables_in_page_text(self, text: str) -> int:
        """Estimate number of tables in text by looking for aligned patterns."""
        lines = text.split('\n')
        potential_table_lines = 0

        for line in lines:
            # Look for lines with multiple whitespace-separated columns
            parts = line.strip().split()
            if len(parts) >= 3:  # At least 3 columns
                # Check if parts look like tabular data (numbers, short text)
                if any(part.replace('.', '').replace(',', '').isdigit() for part in parts):
                    potential_table_lines += 1

        # Rough estimate: every 5+ aligned lines might be a table
        return potential_table_lines // 5

    def convert_document(
        self,
        source: Union[str, Path],
        output_format: str = "markdown"
    ) -> Dict[str, Any]:
        """
        Convert a document using the Granite Docling model.

        Args:
            source: Path to the document or URL
            output_format: Output format (currently supports 'markdown')

        Returns:
            Dictionary containing the conversion result and metadata
        """
        try:
            logger.info(f"Converting document: {source}")

            # Convert the document
            result = self.converter.convert(source=str(source))
            document = result.document

            # Extract the converted content
            if output_format.lower() == "markdown":
                content = document.export_to_markdown()
            else:
                content = str(document)

            # Prepare result dictionary
            conversion_result = {
                "content": content,
                "source": str(source),
                "format": output_format,
                "pages": len(document.pages) if hasattr(document, 'pages') else 1,
                "metadata": {
                    "model_type": self.model_type,
                    "model_config": str(self.vlm_model.__class__.__name__)
                }
            }

            logger.info(f"Successfully converted document with {conversion_result['pages']} pages")
            return conversion_result

        except Exception as e:
            logger.error(f"Error converting document {source}: {str(e)}")
            raise

    def convert_to_file(
        self,
        source: Union[str, Path],
        output_path: Union[str, Path],
        output_format: str = "markdown"
    ) -> Dict[str, Any]:
        """
        Convert a document and save the result to a file.

        Args:
            source: Path to the input document or URL
            output_path: Path where the converted document will be saved
            output_format: Output format (currently supports 'markdown')

        Returns:
            Dictionary containing the conversion result and metadata
        """
        # Convert the document
        result = self.convert_document(source, output_format)

        # Save to file
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(result["content"])

        result["output_path"] = str(output_path)
        logger.info(f"Saved converted document to: {output_path}")

        return result

    def batch_convert(
        self,
        sources: list,
        output_dir: Union[str, Path],
        output_format: str = "markdown"
    ) -> list:
        """
        Convert multiple documents in batch.

        Args:
            sources: List of document paths or URLs
            output_dir: Directory to save converted documents
            output_format: Output format for all documents

        Returns:
            List of conversion results
        """
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        results = []

        for source in sources:
            try:
                # Generate output filename
                source_path = Path(source)
                if output_format.lower() == "markdown":
                    output_filename = source_path.stem + ".md"
                else:
                    output_filename = source_path.stem + f".{output_format}"

                output_path = output_dir / output_filename

                # Convert and save
                result = self.convert_to_file(source, output_path, output_format)
                results.append(result)

            except Exception as e:
                logger.error(f"Failed to convert {source}: {str(e)}")
                results.append({
                    "source": str(source),
                    "error": str(e),
                    "success": False
                })

        return results


def download_models():
    """Download the required Granite Docling models."""
    try:
        import subprocess
        logger.info("Downloading Granite Docling models...")
        subprocess.run([
            "docling-tools", "models", "download-hf-repo",
            "ibm-granite/granite-docling-258M"
        ], check=True)
        logger.info("Models downloaded successfully!")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to download models: {e}")
        raise
    except FileNotFoundError:
        logger.error("docling-tools not found. Please install docling first.")
        raise


if __name__ == "__main__":
    # Example usage
    granite = GraniteDocling()

    # Example conversion (replace with actual document path)
    # result = granite.convert_document("path/to/document.pdf")
    # print(result["content"])