natural-pdf 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

natural_pdf/classification/mixin.py

@@ -0,0 +1,149 @@
+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,62 @@
+# natural_pdf/classification/results.py
+from typing import List, Optional, Dict, Any
+from datetime import datetime
+import logging
+
+logger = logging.getLogger(__name__)
+
+class CategoryScore:
+    """Represents the score for a single category."""
+    label: str
+    confidence: float # Score between 0.0 and 1.0
+
+    def __init__(self, label: str, confidence: float):
+        # Basic validation
+        if not isinstance(label, str) or not label:
+             logger.warning(f"Initializing CategoryScore with invalid label: {label}")
+             # Fallback or raise? For now, allow but log.
+             # raise ValueError("Category label must be a non-empty string.")
+        if not isinstance(confidence, (float, int)) or not (0.0 <= confidence <= 1.0):
+             logger.warning(f"Initializing CategoryScore with invalid confidence: {confidence} for label '{label}'. Clamping to [0, 1].")
+             confidence = max(0.0, min(1.0, float(confidence)))
+             # raise ValueError("Category confidence must be a float between 0.0 and 1.0.")
+        
+        self.label = str(label)
+        self.confidence = float(confidence)
+
+    def __repr__(self):
+        return f"<CategoryScore label='{self.label}' confidence={self.confidence:.3f}>"
+
+class ClassificationResult:
+    """Holds the structured results of a classification task."""
+    model_id: str
+    using: str # Renamed from engine_type ('text' or 'vision')
+    timestamp: datetime
+    parameters: Dict[str, Any] # e.g., {'categories': [...], 'min_confidence': 0.1}
+    scores: List[CategoryScore] # List of scores above threshold, sorted by confidence
+
+    def __init__(self, model_id: str, using: str, timestamp: datetime, parameters: Dict[str, Any], scores: List[CategoryScore]):
+        if not isinstance(scores, list) or not all(isinstance(s, CategoryScore) for s in scores):
+             raise TypeError("Scores must be a list of CategoryScore objects.")
+             
+        self.model_id = str(model_id)
+        self.using = str(using) # Renamed from engine_type
+        self.timestamp = timestamp
+        self.parameters = parameters if parameters is not None else {}
+        # Ensure scores are sorted descending by confidence
+        self.scores = sorted(scores, key=lambda s: s.confidence, reverse=True)
+
+    @property
+    def top_category(self) -> Optional[str]:
+        """Returns the label of the category with the highest confidence."""
+        return self.scores[0].label if self.scores else None
+
+    @property
+    def top_confidence(self) -> Optional[float]:
+        """Returns the confidence score of the top category."""
+        return self.scores[0].confidence if self.scores else None
+
+    def __repr__(self):
+        top_cat = f" top='{self.top_category}' ({self.top_confidence:.2f})" if self.scores else ""
+        num_scores = len(self.scores)
+        return f"<ClassificationResult model='{self.model_id}' using='{self.using}' scores={num_scores}{top_cat}>" 

natural_pdf/collections/mixins.py

@@ -0,0 +1,63 @@
+import logging
+from typing import Callable, Iterable, Any, TypeVar
+from tqdm.auto import tqdm
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T") # Generic type for items in the collection
+
+class ApplyMixin:
+    """
+    Mixin class providing an `.apply()` method for collections.
+
+    Assumes the inheriting class implements `__iter__` and `__len__` appropriately
+    for the items to be processed by `apply`.
+    """
+    def _get_items_for_apply(self) -> Iterable[Any]:
+        """
+        Returns the iterable of items to apply the function to.
+        Defaults to iterating over `self`. Subclasses should override this
+        if the default iteration is not suitable for the apply operation.
+        """
+        # Default to standard iteration over the collection itself
+        return iter(self)
+
+    def apply(self: Any, func: Callable[[Any, ...], Any], *args, **kwargs) -> None:
+        """
+        Applies a function to each item in the collection.
+
+        Args:
+            func: The function to apply to each item. The item itself
+                  will be passed as the first argument to the function.
+            *args: Additional positional arguments to pass to func.
+            **kwargs: Additional keyword arguments to pass to func.
+                      A special keyword argument 'show_progress' (bool, default=False)
+                      can be used to display a progress bar.
+        """
+        show_progress = kwargs.pop('show_progress', False)
+        # Derive unit name from class name
+        unit_name = self.__class__.__name__.lower()
+        items_iterable = self._get_items_for_apply()
+
+        # Need total count for tqdm, assumes __len__ is implemented by the inheriting class
+        total_items = 0
+        try:
+             total_items = len(self)
+        except TypeError: # Handle cases where __len__ might not be defined on self
+             logger.warning(f"Could not determine collection length for progress bar.")
+
+        if show_progress and total_items > 0:
+            items_iterable = tqdm(items_iterable, total=total_items, desc=f"Applying {func.__name__}", unit=unit_name)
+        elif show_progress:
+             logger.info(f"Applying {func.__name__} (progress bar disabled for zero/unknown length).")
+
+        for item in items_iterable:
+            try:
+                # Apply the function with the item and any extra args/kwargs
+                func(item, *args, **kwargs)
+            except Exception as e:
+                # Log and continue for batch operations
+                logger.error(f"Error applying {func.__name__} to {item}: {e}", exc_info=True)
+                # Optionally add a mechanism to collect errors
+
+        # Returns None, primarily used for side effects.

natural_pdf/collections/pdf_collection.py

@@ -4,12 +4,24 @@ import logging
 import os
 import re  # Added for safe path generation
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Type, Union, Callable
+import concurrent.futures # Import concurrent.futures
+import time # Import time for logging timestamps
+import threading # Import threading for logging thread information
 
 from PIL import Image
 from tqdm import tqdm
+from tqdm.auto import tqdm as auto_tqdm
+from tqdm.notebook import tqdm as notebook_tqdm
+
+from natural_pdf.utils.tqdm_utils import get_tqdm
+
+# Get the appropriate tqdm class once
+tqdm = get_tqdm()
 
 # Set up logger early
+# Configure logging to include thread information
+# logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(threadName)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 from natural_pdf.core.pdf import PDF
@@ -36,9 +48,11 @@ except ImportError as e:
     SearchServiceProtocol, SearchOptions, Indexable = object, object, object
 
 from natural_pdf.search.searchable_mixin import SearchableMixin  # Import the new mixin
+# Import the ApplyMixin
+from natural_pdf.collections.mixins import ApplyMixin
 
 
-class PDFCollection(SearchableMixin):  # Inherit from the mixin
+class PDFCollection(SearchableMixin, ApplyMixin):  # Inherit from ApplyMixin
     def __init__(
         self,
         source: Union[str, Iterable[Union[str, "PDF"]]],
@@ -237,30 +251,214 @@ class PDFCollection(SearchableMixin):  # Inherit from the mixin
 
     def __repr__(self) -> str:
         # Removed search status
-        return f"<PDFCollection(count={len(self)})>"
+        return f"<PDFCollection(count={len(self._pdfs)})>"
 
     @property
     def pdfs(self) -> List["PDF"]:
         """Returns the list of PDF objects held by the collection."""
         return self._pdfs
 
-    def apply_ocr(self, *args, **kwargs):
-        PDF = self._get_pdf_class()
-        # Delegate to individual PDF objects
-        logger.info("Applying OCR to relevant PDFs in collection...")
-        results = []
+    def find_all(
+        self, 
+        selector: str, 
+        apply_exclusions: bool = True,  # Added explicit parameter
+        regex: bool = False,            # Added explicit parameter
+        case: bool = True,             # Added explicit parameter
+        **kwargs
+    ) -> "ElementCollection":
+        """
+        Find all elements matching the selector across all PDFs in the collection.
+        
+        This creates an ElementCollection that can span multiple PDFs. Note that
+        some ElementCollection methods have limitations when spanning PDFs.
+        
+        Args:
+            selector: CSS-like selector string to query elements
+            apply_exclusions: Whether to exclude elements in exclusion regions (default: True)
+            regex: Whether to use regex for text search in :contains (default: False)
+            case: Whether to do case-sensitive text search (default: True)
+            **kwargs: Additional keyword arguments passed to the find_all method of each PDF
+            
+        Returns:
+            ElementCollection containing all matching elements across all PDFs
+        """
+        from natural_pdf.elements.collections import ElementCollection
+        
+        # Collect elements from all PDFs
+        all_elements = []
         for pdf in self._pdfs:
-            # We need to figure out which pages belong to which PDF if batching here
-            # For now, simpler to call on each PDF
             try:
-                # Assume apply_ocr exists on PDF and accepts similar args
-                pdf.apply_ocr(*args, **kwargs)
+                # Explicitly pass the relevant arguments down
+                elements = pdf.find_all(
+                    selector,
+                    apply_exclusions=apply_exclusions,
+                    regex=regex,
+                    case=case,
+                    **kwargs
+                )
+                all_elements.extend(elements.elements)
             except Exception as e:
-                logger.error(f"Failed applying OCR to {pdf.path}: {e}", exc_info=True)
+                logger.error(f"Error finding elements in {pdf.path}: {e}", exc_info=True)
+                
+        return ElementCollection(all_elements)
+
+    def apply_ocr(
+        self,
+        engine: Optional[str] = None,
+        languages: Optional[List[str]] = None,
+        min_confidence: Optional[float] = None,
+        device: Optional[str] = None,
+        resolution: Optional[int] = None,
+        apply_exclusions: bool = True,
+        detect_only: bool = False,
+        replace: bool = True,
+        options: Optional[Any] = None,
+        pages: Optional[Union[slice, List[int]]] = None,
+        max_workers: Optional[int] = None,
+    ) -> "PDFCollection":
+        """
+        Apply OCR to all PDFs in the collection, potentially in parallel.
+
+        Args:
+            engine: OCR engine to use (e.g., 'easyocr', 'paddleocr', 'surya')
+            languages: List of language codes for OCR
+            min_confidence: Minimum confidence threshold for text detection
+            device: Device to use for OCR (e.g., 'cpu', 'cuda')
+            resolution: DPI resolution for page rendering
+            apply_exclusions: Whether to apply exclusion regions
+            detect_only: If True, only detect text regions without extracting text
+            replace: If True, replace existing OCR elements
+            options: Engine-specific options
+            pages: Specific pages to process (None for all pages)
+            max_workers: Maximum number of threads to process PDFs concurrently. 
+                         If None or 1, processing is sequential. (default: None)
+
+        Returns:
+            Self for method chaining
+        """
+        PDF = self._get_pdf_class()
+        logger.info(f"Applying OCR to {len(self._pdfs)} PDFs in collection (max_workers={max_workers})...")
+
+        # Worker function takes PDF object again
+        def _process_pdf(pdf: PDF):
+            """Helper function to apply OCR to a single PDF, handling errors."""
+            thread_id = threading.current_thread().name # Get thread name for logging
+            pdf_path = pdf.path # Get path for logging
+            logger.debug(f"[{thread_id}] Starting OCR process for: {pdf_path}")
+            start_time = time.monotonic()
+            try:
+                pdf.apply_ocr( # Call apply_ocr on the original PDF object
+                    pages=pages,
+                    engine=engine,
+                    languages=languages,
+                    min_confidence=min_confidence,
+                    device=device,
+                    resolution=resolution,
+                    apply_exclusions=apply_exclusions,
+                    detect_only=detect_only,
+                    replace=replace,
+                    options=options,
+                    # Note: We might want a max_workers here too for page rendering?
+                    # For now, PDF.apply_ocr doesn't have it.
+                )
+                end_time = time.monotonic()
+                logger.debug(f"[{thread_id}] Finished OCR process for: {pdf_path} (Duration: {end_time - start_time:.2f}s)")
+                return pdf_path, None
+            except Exception as e:
+                end_time = time.monotonic()
+                logger.error(f"[{thread_id}] Failed OCR process for {pdf_path} after {end_time - start_time:.2f}s: {e}", exc_info=False)
+                return pdf_path, e # Return path and error
+
+        # Use ThreadPoolExecutor for parallel processing if max_workers > 1
+        if max_workers is not None and max_workers > 1:
+            futures = []
+            with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="OCRWorker") as executor:
+                for pdf in self._pdfs:
+                    # Submit the PDF object to the worker function
+                    futures.append(executor.submit(_process_pdf, pdf))
+
+            # Use the selected tqdm class with as_completed for progress tracking
+            progress_bar = tqdm(
+                concurrent.futures.as_completed(futures),
+                total=len(self._pdfs),
+                desc="Applying OCR (Parallel)",
+                unit="pdf"
+            )
+            
+            for future in progress_bar:
+                pdf_path, error = future.result() # Get result (or exception)
+                if error:
+                    progress_bar.set_postfix_str(f"Error: {pdf_path}", refresh=True)
+                # Progress is updated automatically by tqdm
+
+        else: # Sequential processing (max_workers is None or 1)
+            logger.info("Applying OCR sequentially...")
+            # Use the selected tqdm class for sequential too for consistency
+            # Iterate over PDF objects directly for sequential
+            for pdf in tqdm(self._pdfs, desc="Applying OCR (Sequential)", unit="pdf"):
+                _process_pdf(pdf) # Call helper directly with PDF object
+        
+        logger.info("Finished applying OCR across the collection.")
         return self
 
-    # --- Advanced Method Placeholders ---
-    # Placeholder for categorize removed as find_relevant is now implemented
+    def correct_ocr(
+        self,
+        correction_callback: Callable[[Any], Optional[str]],
+        max_workers: Optional[int] = None,
+        progress_callback: Optional[Callable[[], None]] = None,
+    ) -> "PDFCollection":
+        """
+        Apply OCR correction to all relevant elements across all pages and PDFs
+        in the collection using a single progress bar.
+
+        Args:
+            correction_callback: Function to apply to each OCR element.
+                                 It receives the element and should return
+                                 the corrected text (str) or None.
+            max_workers: Max threads to use for parallel execution within each page.
+            progress_callback: Optional callback function to call after processing each element.
+
+        Returns:
+            Self for method chaining.
+        """
+        PDF = self._get_pdf_class() # Ensure PDF class is available
+        if not callable(correction_callback):
+            raise TypeError("`correction_callback` must be a callable function.")
+
+        logger.info(f"Gathering OCR elements from {len(self._pdfs)} PDFs for correction...")
+
+        # 1. Gather all target elements using the collection's find_all
+        #    Crucially, set apply_exclusions=False to include elements in headers/footers etc.
+        all_ocr_elements = self.find_all("text[source=ocr]", apply_exclusions=False).elements
+
+        if not all_ocr_elements:
+            logger.info("No OCR elements found in the collection to correct.")
+            return self
+
+        total_elements = len(all_ocr_elements)
+        logger.info(f"Found {total_elements} OCR elements across the collection. Starting correction process...")
+
+        # 2. Initialize the progress bar
+        progress_bar = tqdm(total=total_elements, desc="Correcting OCR Elements", unit="element")
+
+        # 3. Iterate through PDFs and delegate to PDF.correct_ocr
+        #    PDF.correct_ocr handles page iteration and passing the progress callback down.
+        for pdf in self._pdfs:
+            if not pdf.pages:
+                continue
+            try:
+                pdf.correct_ocr(
+                    correction_callback=correction_callback,
+                    max_workers=max_workers,
+                    progress_callback=progress_bar.update # Pass the bar's update method
+                )
+            except Exception as e:
+                 logger.error(f"Error occurred during correction process for PDF {pdf.path}: {e}", exc_info=True)
+                 # Decide if we should stop or continue? For now, continue.
+
+        progress_bar.close()
+
+        return self
 
     def categorize(self, categories: List[str], **kwargs):
         """Categorizes PDFs in the collection based on content or features."""
@@ -279,14 +477,17 @@ class PDFCollection(SearchableMixin):  # Inherit from the mixin
         """
         try:
             from natural_pdf.utils.packaging import create_correction_task_package
+
             # Pass the collection itself (self) as the source
             create_correction_task_package(source=self, output_zip_path=output_zip_path, **kwargs)
         except ImportError:
-            logger.error("Failed to import 'create_correction_task_package'. Packaging utility might be missing.")
+            logger.error(
+                "Failed to import 'create_correction_task_package'. Packaging utility might be missing."
+            )
             # Or raise
         except Exception as e:
             logger.error(f"Failed to export correction task for collection: {e}", exc_info=True)
-            raise # Re-raise the exception from the utility function
+            raise  # Re-raise the exception from the utility function
 
     # --- Mixin Required Implementation ---
     def get_indexable_items(self) -> Iterable[Indexable]:
@@ -306,3 +507,111 @@ class PDFCollection(SearchableMixin):  # Inherit from the mixin
                 #     logger.debug(f"Skipping empty page {page.page_number} from PDF '{pdf.path}'.")
                 #     continue
                 yield page
+
+    # --- Classification Method --- #
+    def classify_all(
+        self,
+        categories: List[str],
+        model: str = "text",
+        max_workers: Optional[int] = None,
+        **kwargs,
+    ) -> "PDFCollection":
+        """
+        Classify all pages across all PDFs in the collection, potentially in parallel.
+
+        This method uses the unified `classify_all` approach, delegating page
+        classification to each PDF's `classify_pages` method.
+        It displays a progress bar tracking individual pages.
+
+        Args:
+            categories: A list of string category names.
+            model: Model identifier ('text', 'vision', or specific HF ID).
+            max_workers: Maximum number of threads to process PDFs concurrently.
+                         If None or 1, processing is sequential.
+            **kwargs: Additional arguments passed down to `pdf.classify_pages` and
+                      subsequently to `page.classify` (e.g., device,
+                      confidence_threshold, resolution).
+
+        Returns:
+            Self for method chaining.
+
+        Raises:
+            ValueError: If categories list is empty.
+            ClassificationError: If classification fails for any page (will stop processing).
+            ImportError: If classification dependencies are missing.
+        """
+        PDF = self._get_pdf_class()
+        if not categories:
+            raise ValueError("Categories list cannot be empty.")
+
+        logger.info(f"Starting classification for {len(self._pdfs)} PDFs in collection (model: '{model}')...")
+
+        # Calculate total pages for the progress bar
+        total_pages = sum(len(pdf.pages) for pdf in self._pdfs if pdf.pages)
+        if total_pages == 0:
+            logger.warning("No pages found in the PDF collection to classify.")
+            return self
+
+        progress_bar = tqdm(
+            total=total_pages,
+            desc=f"Classifying Pages (model: {model})",
+            unit="page"
+        )
+
+        # Worker function
+        def _process_pdf_classification(pdf: PDF):
+            thread_id = threading.current_thread().name
+            pdf_path = pdf.path
+            logger.debug(f"[{thread_id}] Starting classification process for: {pdf_path}")
+            start_time = time.monotonic()
+            try:
+                # Call classify_pages on the PDF, passing the progress callback
+                pdf.classify_pages(
+                    categories=categories,
+                    model=model,
+                    progress_callback=progress_bar.update,
+                    **kwargs
+                )
+                end_time = time.monotonic()
+                logger.debug(f"[{thread_id}] Finished classification for: {pdf_path} (Duration: {end_time - start_time:.2f}s)")
+                return pdf_path, None # Return path and no error
+            except Exception as e:
+                end_time = time.monotonic()
+                # Error is logged within classify_pages, but log summary here
+                logger.error(f"[{thread_id}] Failed classification process for {pdf_path} after {end_time - start_time:.2f}s: {e}", exc_info=False)
+                # Close progress bar immediately on error to avoid hanging
+                progress_bar.close()
+                # Re-raise the exception to stop the entire collection processing
+                raise
+
+        # Use ThreadPoolExecutor for parallel processing if max_workers > 1
+        try:
+            if max_workers is not None and max_workers > 1:
+                logger.info(f"Classifying PDFs in parallel with {max_workers} workers.")
+                futures = []
+                with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="ClassifyWorker") as executor:
+                    for pdf in self._pdfs:
+                        futures.append(executor.submit(_process_pdf_classification, pdf))
+
+                    # Wait for all futures to complete (progress updated by callback)
+                    # Exceptions are raised by future.result() if worker failed
+                    for future in concurrent.futures.as_completed(futures):
+                         future.result() # Raise exception if worker failed
+
+            else: # Sequential processing
+                logger.info("Classifying PDFs sequentially.")
+                for pdf in self._pdfs:
+                    _process_pdf_classification(pdf)
+
+            logger.info("Finished classification across the collection.")
+
+        finally:
+             # Ensure progress bar is closed even if errors occurred elsewhere
+             if not progress_bar.disable and progress_bar.n < progress_bar.total:
+                 progress_bar.close()
+             elif progress_bar.disable is False:
+                  progress_bar.close()
+
+        return self
+
+    # --- End Classification Method --- #