natural-pdf 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

natural_pdf/classification/manager.py

@@ -0,0 +1,422 @@
+import logging
+import time
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+from PIL import Image
+
+# Use try-except for robustness if dependencies are missing
+try:
+    import torch
+    from transformers import (
+        AutoModelForSequenceClassification,
+        AutoModelForZeroShotImageClassification,
+        AutoTokenizer,
+        pipeline,
+    )
+
+    _CLASSIFICATION_AVAILABLE = True
+except ImportError:
+    _CLASSIFICATION_AVAILABLE = False
+    # Define dummy types for type hinting if imports fail
+    pipeline = object
+    AutoTokenizer = object
+    AutoModelForZeroShotImageClassification = object
+    AutoModelForSequenceClassification = object
+    torch = None
+
+from natural_pdf.utils.tqdm_utils import get_tqdm
+
+# Import result classes
+from .results import CategoryScore, ClassificationResult
+
+if TYPE_CHECKING:
+    from transformers import Pipeline
+
+
+logger = logging.getLogger(__name__)
+
+# Global cache for models/pipelines
+_PIPELINE_CACHE: Dict[str, "Pipeline"] = {}
+_TOKENIZER_CACHE: Dict[str, Any] = {}
+_MODEL_CACHE: Dict[str, Any] = {}
+
+
+class ClassificationError(Exception):
+    """Custom exception for classification errors."""
+
+    pass
+
+
+class ClassificationManager:
+    """Manages classification models and execution."""
+
+    DEFAULT_TEXT_MODEL = "facebook/bart-large-mnli"
+    DEFAULT_VISION_MODEL = "openai/clip-vit-base-patch16"
+
+    def __init__(
+        self,
+        model_mapping: Optional[Dict[str, str]] = None,
+        default_device: Optional[str] = None,
+    ):
+        """
+        Initialize the ClassificationManager.
+
+        Args:
+            model_mapping: Optional dictionary mapping aliases ('text', 'vision') to model IDs.
+            default_device: Default device ('cpu', 'cuda') if not specified in classify calls.
+        """
+        if not _CLASSIFICATION_AVAILABLE:
+            raise ImportError(
+                "Classification dependencies missing. "
+                'Install with: pip install "natural-pdf[classification]"'
+            )
+
+        self.pipelines: Dict[Tuple[str, str], "Pipeline"] = (
+            {}
+        )  # Cache: (model_id, device) -> pipeline
+
+        self.device = default_device
+        logger.info(f"ClassificationManager initialized on device: {self.device}")
+
+    def is_available(self) -> bool:
+        """Check if required dependencies are installed."""
+        return _CLASSIFICATION_AVAILABLE
+
+    def _get_pipeline(self, model_id: str, using: str) -> "Pipeline":
+        """Get or create a classification pipeline."""
+        cache_key = f"{model_id}_{using}_{self.device}"
+        if cache_key not in _PIPELINE_CACHE:
+            logger.info(
+                f"Loading {using} classification pipeline for model '{model_id}' on device '{self.device}'..."
+            )
+            start_time = time.time()
+            try:
+                task = (
+                    "zero-shot-classification"
+                    if using == "text"
+                    else "zero-shot-image-classification"
+                )
+                _PIPELINE_CACHE[cache_key] = pipeline(task, model=model_id, device=self.device)
+                end_time = time.time()
+                logger.info(
+                    f"Pipeline for '{model_id}' loaded in {end_time - start_time:.2f} seconds."
+                )
+            except Exception as e:
+                logger.error(
+                    f"Failed to load pipeline for model '{model_id}' (using: {using}): {e}",
+                    exc_info=True,
+                )
+                raise ClassificationError(
+                    f"Failed to load pipeline for model '{model_id}'. Ensure the model ID is correct and supports the {task} task."
+                ) from e
+        return _PIPELINE_CACHE[cache_key]
+
+    def infer_using(self, model_id: str, using: Optional[str] = None) -> str:
+        """Infers processing mode ('text' or 'vision') if not provided."""
+        if using in ["text", "vision"]:
+            return using
+
+        # Simple inference based on common model names
+        normalized_model_id = model_id.lower()
+        if (
+            "clip" in normalized_model_id
+            or "vit" in normalized_model_id
+            or "siglip" in normalized_model_id
+        ):
+            logger.debug(f"Inferred using='vision' for model '{model_id}'")
+            return "vision"
+        if (
+            "bart" in normalized_model_id
+            or "bert" in normalized_model_id
+            or "mnli" in normalized_model_id
+            or "xnli" in normalized_model_id
+            or "deberta" in normalized_model_id
+        ):
+            logger.debug(f"Inferred using='text' for model '{model_id}'")
+            return "text"
+
+        # Fallback or raise error? Let's try loading text first, then vision.
+        logger.warning(
+            f"Could not reliably infer mode for '{model_id}'. Trying text, then vision pipeline loading."
+        )
+        try:
+            self._get_pipeline(model_id, "text")
+            logger.info(f"Successfully loaded '{model_id}' as a text model.")
+            return "text"
+        except Exception:
+            logger.warning(f"Failed to load '{model_id}' as text model. Trying vision.")
+            try:
+                self._get_pipeline(model_id, "vision")
+                logger.info(f"Successfully loaded '{model_id}' as a vision model.")
+                return "vision"
+            except Exception as e_vision:
+                logger.error(
+                    f"Failed to load '{model_id}' as either text or vision model.", exc_info=True
+                )
+                raise ClassificationError(
+                    f"Cannot determine mode for model '{model_id}'. Please specify `using='text'` or `using='vision'`. Error: {e_vision}"
+                )
+
+    def classify_item(
+        self,
+        item_content: Union[str, Image.Image],
+        categories: List[str],
+        model_id: Optional[str] = None,
+        using: Optional[str] = None,
+        min_confidence: float = 0.0,
+        multi_label: bool = False,
+        **kwargs,
+    ) -> ClassificationResult:  # Return ClassificationResult
+        """Classifies a single item (text or image)."""
+
+        # Determine model and engine type
+        effective_using = using
+        if model_id is None:
+            # Try inferring based on content type
+            if isinstance(item_content, str):
+                effective_using = "text"
+                model_id = self.DEFAULT_TEXT_MODEL
+            elif isinstance(item_content, Image.Image):
+                effective_using = "vision"
+                model_id = self.DEFAULT_VISION_MODEL
+            else:
+                raise TypeError(f"Unsupported item_content type: {type(item_content)}")
+        else:
+            # Infer engine type if not given
+            effective_using = self.infer_using(model_id, using)
+            # Set default model if needed (though should usually be provided if engine known)
+            if model_id is None:
+                model_id = (
+                    self.DEFAULT_TEXT_MODEL
+                    if effective_using == "text"
+                    else self.DEFAULT_VISION_MODEL
+                )
+
+        if not categories:
+            raise ValueError("Categories list cannot be empty.")
+
+        pipeline_instance = self._get_pipeline(model_id, effective_using)
+        timestamp = datetime.now()
+        parameters = {  # Store parameters used for this run
+            "categories": categories,
+            "model_id": model_id,
+            "using": effective_using,
+            "min_confidence": min_confidence,
+            "multi_label": multi_label,
+            **kwargs,
+        }
+
+        logger.debug(
+            f"Classifying content (type: {type(item_content).__name__}) with model '{model_id}'"
+        )
+        try:
+            # Handle potential kwargs for specific pipelines if needed
+            # The zero-shot pipelines expect `candidate_labels`
+            result_raw = pipeline_instance(
+                item_content, candidate_labels=categories, multi_label=multi_label, **kwargs
+            )
+            logger.debug(f"Raw pipeline result: {result_raw}")
+
+            # --- Process raw result into ClassificationResult --- #
+            scores_list: List[CategoryScore] = []
+
+            # Handle text pipeline format (dict with 'labels' and 'scores')
+            if isinstance(result_raw, dict) and "labels" in result_raw and "scores" in result_raw:
+                for label, score_val in zip(result_raw["labels"], result_raw["scores"]):
+                    if score_val >= min_confidence:
+                        try:
+                            scores_list.append(CategoryScore(label=label, confidence=score_val))
+                        except (ValueError, TypeError) as score_err:
+                            logger.warning(
+                                f"Skipping invalid score from text pipeline: label='{label}', score={score_val}. Error: {score_err}"
+                            )
+            # Handle vision pipeline format (list of dicts with 'label' and 'score')
+            elif isinstance(result_raw, list) and all(
+                isinstance(item, dict) and "label" in item and "score" in item
+                for item in result_raw
+            ):
+                for item in result_raw:
+                    score_val = item["score"]
+                    label = item["label"]
+                    if score_val >= min_confidence:
+                        try:
+                            scores_list.append(CategoryScore(label=label, confidence=score_val))
+                        except (ValueError, TypeError) as score_err:
+                            logger.warning(
+                                f"Skipping invalid score from vision pipeline: label='{label}', score={score_val}. Error: {score_err}"
+                            )
+            else:
+                logger.warning(
+                    f"Unexpected raw result format from pipeline for model '{model_id}': {type(result_raw)}. Cannot extract scores."
+                )
+                # Return empty result?
+                # scores_list = []
+
+            return ClassificationResult(
+                model_id=model_id,
+                using=effective_using,
+                timestamp=timestamp,
+                parameters=parameters,
+                scores=scores_list,
+            )
+            # --- End Processing --- #
+
+        except Exception as e:
+            logger.error(f"Classification failed for model '{model_id}': {e}", exc_info=True)
+            # Return an empty result object on failure?
+            # return ClassificationResult(model_id=model_id, engine_type=engine_type, timestamp=timestamp, parameters=parameters, scores=[])
+            raise ClassificationError(
+                f"Classification failed using model '{model_id}'. Error: {e}"
+            ) from e
+
+    def classify_batch(
+        self,
+        item_contents: List[Union[str, Image.Image]],
+        categories: List[str],
+        model_id: Optional[str] = None,
+        using: Optional[str] = None,
+        min_confidence: float = 0.0,
+        multi_label: bool = False,
+        batch_size: int = 8,
+        progress_bar: bool = True,
+        **kwargs,
+    ) -> List[ClassificationResult]:  # Return list of ClassificationResult
+        """Classifies a batch of items (text or image) using the pipeline's batching."""
+        if not item_contents:
+            return []
+
+        # Determine model and engine type (assuming uniform type in batch)
+        first_item = item_contents[0]
+        effective_using = using
+        if model_id is None:
+            if isinstance(first_item, str):
+                effective_using = "text"
+                model_id = self.DEFAULT_TEXT_MODEL
+            elif isinstance(first_item, Image.Image):
+                effective_using = "vision"
+                model_id = self.DEFAULT_VISION_MODEL
+            else:
+                raise TypeError(f"Unsupported item_content type in batch: {type(first_item)}")
+        else:
+            effective_using = self.infer_using(model_id, using)
+            if model_id is None:
+                model_id = (
+                    self.DEFAULT_TEXT_MODEL
+                    if effective_using == "text"
+                    else self.DEFAULT_VISION_MODEL
+                )
+
+        if not categories:
+            raise ValueError("Categories list cannot be empty.")
+
+        pipeline_instance = self._get_pipeline(model_id, effective_using)
+        timestamp = datetime.now()  # Single timestamp for the batch run
+        parameters = {  # Parameters for the whole batch
+            "categories": categories,
+            "model_id": model_id,
+            "using": effective_using,
+            "min_confidence": min_confidence,
+            "multi_label": multi_label,
+            "batch_size": batch_size,
+            **kwargs,
+        }
+
+        logger.info(
+            f"Classifying batch of {len(item_contents)} items with model '{model_id}' (batch size: {batch_size})"
+        )
+        batch_results_list: List[ClassificationResult] = []
+
+        try:
+            # Use pipeline directly for batching
+            results_iterator = pipeline_instance(
+                item_contents,
+                candidate_labels=categories,
+                multi_label=multi_label,
+                batch_size=batch_size,
+                **kwargs,
+            )
+
+            # Wrap with tqdm for progress if requested
+            total_items = len(item_contents)
+            if progress_bar:
+                # Get the appropriate tqdm class
+                tqdm_class = get_tqdm()
+                results_iterator = tqdm_class(
+                    results_iterator,
+                    total=total_items,
+                    desc=f"Classifying batch ({model_id})",
+                    leave=False,  # Don't leave progress bar hanging
+                )
+
+            for raw_result in results_iterator:
+                # --- Process each raw result (which corresponds to ONE input item) --- #
+                scores_list: List[CategoryScore] = []
+                try:
+                    # Check for text format (dict with 'labels' and 'scores')
+                    if (
+                        isinstance(raw_result, dict)
+                        and "labels" in raw_result
+                        and "scores" in raw_result
+                    ):
+                        for label, score_val in zip(raw_result["labels"], raw_result["scores"]):
+                            if score_val >= min_confidence:
+                                try:
+                                    scores_list.append(
+                                        CategoryScore(label=label, confidence=score_val)
+                                    )
+                                except (ValueError, TypeError) as score_err:
+                                    logger.warning(
+                                        f"Skipping invalid score from text pipeline batch: label='{label}', score={score_val}. Error: {score_err}"
+                                    )
+                    # Check for vision format (list of dicts with 'label' and 'score')
+                    elif isinstance(raw_result, list):
+                        for item in raw_result:
+                            try:
+                                score_val = item["score"]
+                                label = item["label"]
+                                if score_val >= min_confidence:
+                                    scores_list.append(
+                                        CategoryScore(label=label, confidence=score_val)
+                                    )
+                            except (KeyError, ValueError, TypeError) as item_err:
+                                logger.warning(
+                                    f"Skipping invalid item in vision result list from batch: {item}. Error: {item_err}"
+                                )
+                    else:
+                        logger.warning(
+                            f"Unexpected raw result format in batch item from model '{model_id}': {type(raw_result)}. Cannot extract scores."
+                        )
+
+                except Exception as proc_err:
+                    logger.error(
+                        f"Error processing result item in batch: {proc_err}", exc_info=True
+                    )
+                    # scores_list remains empty for this item
+
+                # Append result object for this item
+                batch_results_list.append(
+                    ClassificationResult(
+                        model_id=model_id,
+                        using=effective_using,
+                        timestamp=timestamp,  # Use same timestamp for batch
+                        parameters=parameters,  # Use same params for batch
+                        scores=scores_list,
+                    )
+                )
+                # --- End Processing --- #
+
+            if len(batch_results_list) != total_items:
+                logger.warning(
+                    f"Batch classification returned {len(batch_results_list)} results, but expected {total_items}. Results might be incomplete or misaligned."
+                )
+
+            return batch_results_list
+
+        except Exception as e:
+            logger.error(f"Batch classification failed for model '{model_id}': {e}", exc_info=True)
+            # Return list of empty results?
+            # return [ClassificationResult(model_id=model_id, s=engine_type, timestamp=timestamp, parameters=parameters, scores=[]) for _ in item_contents]
+            raise ClassificationError(
+                f"Batch classification failed using model '{model_id}'. Error: {e}"
+            ) from e

natural_pdf/classification/mixin.py

@@ -0,0 +1,163 @@
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+
+# Assuming PIL is installed as it's needed for vision
+try:
+    from PIL import Image
+except ImportError:
+    Image = None  # type: ignore
+
+# Import result classes
+from .results import ClassificationResult  # Assuming results.py is in the same dir
+
+if TYPE_CHECKING:
+    # Avoid runtime import cycle
+    from natural_pdf.core.page import Page
+    from natural_pdf.elements.region import Region
+
+    from .manager import ClassificationManager
+
+logger = logging.getLogger(__name__)
+
+
+class ClassificationMixin:
+    """
+    Mixin class providing classification capabilities to Page and Region objects.
+    Relies on a ClassificationManager being accessible, typically via the parent PDF.
+    """
+
+    # --- Abstract methods/properties required by the host class --- #
+    # These must be implemented by classes using this mixin (Page, Region)
+
+    def _get_classification_manager(self) -> "ClassificationManager":
+        """Should return the ClassificationManager instance."""
+        raise NotImplementedError
+
+    def _get_classification_content(self, model_type: str, **kwargs) -> Union[str, "Image"]:
+        """Should return the text content (str) or image (PIL.Image) for classification."""
+        raise NotImplementedError
+
+    # Host class needs 'analyses' attribute initialized as Dict[str, Any]
+    # analyses: Dict[str, Any]
+
+    # --- End Abstract --- #
+
+    def classify(
+        self,
+        categories: List[str],
+        model: Optional[str] = None,  # Default handled by manager
+        using: Optional[str] = None,  # Renamed parameter
+        min_confidence: float = 0.0,
+        analysis_key: str = "classification",  # Default key
+        multi_label: bool = False,
+        **kwargs,
+    ) -> "ClassificationMixin":  # Return self for chaining
+        """
+        Classifies this item (Page or Region) using the configured manager.
+
+        Stores the result in self.analyses[analysis_key]. If analysis_key is not
+        provided, it defaults to 'classification' and overwrites any previous
+        result under that key.
+
+        Args:
+            categories: A list of string category names.
+            model: Model identifier (e.g., 'text', 'vision', HF ID). Defaults handled by manager.
+            using: Optional processing mode ('text' or 'vision'). If None, inferred by manager.
+            min_confidence: Minimum confidence threshold for results (0.0-1.0).
+            analysis_key: Key under which to store the result in `self.analyses`.
+                          Defaults to 'classification'.
+            multi_label: Whether to allow multiple labels (passed to HF pipeline).
+            **kwargs: Additional arguments passed to the ClassificationManager.
+
+        Returns:
+            Self for method chaining.
+        """
+        # Ensure analyses dict exists
+        if not hasattr(self, "analyses") or self.analyses is None:
+            logger.warning("'analyses' attribute not found or is None. Initializing as empty dict.")
+            self.analyses = {}
+
+        try:
+            manager = self._get_classification_manager()
+
+            # Determine the effective model ID and engine type
+            effective_model_id = model
+            inferred_using = manager.infer_using(
+                model if model else manager.DEFAULT_TEXT_MODEL, using
+            )
+
+            # If model was not provided, use the manager's default for the inferred engine type
+            if effective_model_id is None:
+                effective_model_id = (
+                    manager.DEFAULT_TEXT_MODEL
+                    if inferred_using == "text"
+                    else manager.DEFAULT_VISION_MODEL
+                )
+                logger.debug(
+                    f"No model provided, using default for mode '{inferred_using}': '{effective_model_id}'"
+                )
+
+            # Get content based on the *final* determined engine type
+            content = self._get_classification_content(model_type=inferred_using, **kwargs)
+
+            # Manager now returns a ClassificationResult object
+            result_obj: ClassificationResult = manager.classify_item(
+                item_content=content,
+                categories=categories,
+                model_id=effective_model_id,  # Pass the resolved model ID
+                using=inferred_using,  # Pass renamed argument
+                min_confidence=min_confidence,
+                multi_label=multi_label,
+                **kwargs,
+            )
+
+            # Store the structured result object under the specified key
+            self.analyses[analysis_key] = result_obj
+            logger.debug(f"Stored classification result under key '{analysis_key}': {result_obj}")
+
+        except NotImplementedError as nie:
+            logger.error(f"Classification cannot proceed: {nie}")
+            raise
+        except Exception as e:
+            logger.error(f"Classification failed: {e}", exc_info=True)
+            # Optionally re-raise or just log and return self
+            # raise
+
+        return self
+
+    @property
+    def classification_results(self) -> Optional[ClassificationResult]:
+        """Returns the ClassificationResult from the *default* ('classification') key, or None."""
+        if not hasattr(self, "analyses") or self.analyses is None:
+            return None
+        # Return the result object directly from the default key
+        return self.analyses.get("classification")
+
+    @property
+    def category(self) -> Optional[str]:
+        """Returns the top category label from the *default* ('classification') key, or None."""
+        result_obj = self.classification_results  # Uses the property above
+        # Access the property on the result object
+        return result_obj.top_category if result_obj else None
+
+    @property
+    def category_confidence(self) -> Optional[float]:
+        """Returns the top category confidence from the *default* ('classification') key, or None."""
+        result_obj = self.classification_results  # Uses the property above
+        # Access the property on the result object
+        return result_obj.top_confidence if result_obj else None
+
+    # Maybe add a helper to get results by specific key?
+    def get_classification_result(
+        self, analysis_key: str = "classification"
+    ) -> Optional[ClassificationResult]:
+        """Gets a classification result object stored under a specific key."""
+        if not hasattr(self, "analyses") or self.analyses is None:
+            return None
+        result = self.analyses.get(analysis_key)
+        if result is not None and not isinstance(result, ClassificationResult):
+            logger.warning(
+                f"Item found under key '{analysis_key}' is not a ClassificationResult (type: {type(result)}). Returning None."
+            )
+            return None
+        return result

natural_pdf/classification/results.py

@@ -0,0 +1,80 @@
+# natural_pdf/classification/results.py
+import logging
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CategoryScore:
+    """Represents a category and its confidence score from classification."""
+
+    category: str
+    score: float
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {"category": self.category, "score": self.score}
+
+
+@dataclass
+class ClassificationResult:
+    """Results from a classification operation."""
+
+    category: str
+    score: float
+    scores: List[CategoryScore]
+    model_id: str
+    timestamp: datetime
+    using: str  # 'text' or 'vision'
+    parameters: Optional[Dict[str, Any]] = None
+
+    def __init__(
+        self,
+        category: str,
+        score: float,
+        scores: List[CategoryScore],
+        model_id: str,
+        using: str,
+        parameters: Optional[Dict[str, Any]] = None,
+        timestamp: Optional[datetime] = None,
+    ):
+        self.category = category
+        self.score = score
+        self.scores = scores
+        self.model_id = model_id
+        self.using = using
+        self.parameters = parameters or {}
+        self.timestamp = timestamp or datetime.now()
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert the classification result to a dictionary for serialization.
+
+        Returns:
+            Dictionary representation of the classification result
+        """
+        return {
+            "category": self.category,
+            "score": self.score,
+            "scores": [s.to_dict() for s in self.scores],
+            "model_id": self.model_id,
+            "using": self.using,
+            "parameters": self.parameters,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+    @property
+    def top_category(self) -> str:
+        """Returns the category with the highest score."""
+        return self.category
+
+    @property
+    def top_confidence(self) -> float:
+        """Returns the confidence score of the top category."""
+        return self.score
+
+    def __repr__(self) -> str:
+        return f"<ClassificationResult category='{self.category}' score={self.score:.3f} model='{self.model_id}'>"