rnsr 0.1.0__py3-none-any.whl

rnsr/ingestion/header_classifier.py

@@ -0,0 +1,595 @@
+"""
+Header Classifier - Detect and Classify Headers by Level
+
+This module classifies text spans into:
+- Headers (H1, H2, H3) based on font size magnitude
+- Body text (most frequent font size)
+- Captions/footnotes (smaller than body)
+
+Header Level Mapping (adaptive):
+- Defaults: >= 24pt: H1, 18-23pt: H2, 14-17pt: H3
+- Learns from document types (legal, academic, marketing, etc.)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from datetime import datetime
+from pathlib import Path
+from threading import Lock
+from typing import Any, Literal
+
+import numpy as np
+import structlog
+from sklearn.cluster import KMeans
+
+from rnsr.models import ClassifiedSpan, FontAnalysis, SpanInfo
+
+logger = structlog.get_logger(__name__)
+
+
+# =============================================================================
+# Learned Header Thresholds Registry
+# =============================================================================
+
+DEFAULT_HEADER_THRESHOLDS_PATH = Path.home() / ".rnsr" / "learned_header_thresholds.json"
+
+
+class LearnedHeaderThresholds:
+    """
+    Registry for learning document-type-specific header thresholds.
+    
+    Different document types use different conventions:
+    - Legal briefs: 12pt bold = header
+    - Academic papers: 11pt = everything
+    - Marketing: 36pt+ = titles
+    
+    This class learns optimal thresholds from document analysis.
+    """
+    
+    # Default thresholds
+    DEFAULT_H1_MIN = 24.0
+    DEFAULT_H2_MIN = 18.0
+    DEFAULT_H3_MIN = 14.0
+    
+    def __init__(
+        self,
+        storage_path: Path | str | None = None,
+        auto_save: bool = True,
+    ):
+        """
+        Initialize the header thresholds registry.
+        
+        Args:
+            storage_path: Path to JSON file for persistence.
+            auto_save: Whether to save after changes.
+        """
+        self.storage_path = Path(storage_path) if storage_path else DEFAULT_HEADER_THRESHOLDS_PATH
+        self.auto_save = auto_save
+        
+        self._lock = Lock()
+        self._document_types: dict[str, dict[str, Any]] = {}
+        self._dirty = False
+        
+        self._load()
+    
+    def _load(self) -> None:
+        """Load learned thresholds from storage."""
+        if not self.storage_path.exists():
+            return
+        
+        try:
+            with open(self.storage_path, "r") as f:
+                data = json.load(f)
+            
+            self._document_types = data.get("document_types", {})
+            
+            logger.info(
+                "header_thresholds_loaded",
+                document_types=len(self._document_types),
+            )
+            
+        except Exception as e:
+            logger.warning("failed_to_load_header_thresholds", error=str(e))
+    
+    def _save(self) -> None:
+        """Save to storage."""
+        if not self._dirty:
+            return
+        
+        try:
+            self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            data = {
+                "version": "1.0",
+                "updated_at": datetime.utcnow().isoformat(),
+                "document_types": self._document_types,
+            }
+            
+            with open(self.storage_path, "w") as f:
+                json.dump(data, f, indent=2)
+            
+            self._dirty = False
+            
+        except Exception as e:
+            logger.warning("failed_to_save_header_thresholds", error=str(e))
+    
+    def record_thresholds(
+        self,
+        document_type: str,
+        h1_min: float,
+        h2_min: float,
+        h3_min: float,
+        body_size: float = 12.0,
+        document_example: str = "",
+    ) -> None:
+        """
+        Record observed header thresholds for a document type.
+        
+        The system averages thresholds across multiple documents
+        of the same type to learn optimal values.
+        
+        Args:
+            document_type: Type of document (legal, academic, marketing, etc.)
+            h1_min: Observed H1 minimum size.
+            h2_min: Observed H2 minimum size.
+            h3_min: Observed H3 minimum size.
+            body_size: Observed body text size.
+            document_example: Example document filename.
+        """
+        document_type = document_type.lower().strip()
+        
+        if not document_type:
+            return
+        
+        with self._lock:
+            now = datetime.utcnow().isoformat()
+            
+            if document_type not in self._document_types:
+                self._document_types[document_type] = {
+                    "count": 0,
+                    "h1_min_sum": 0.0,
+                    "h2_min_sum": 0.0,
+                    "h3_min_sum": 0.0,
+                    "body_size_sum": 0.0,
+                    "first_seen": now,
+                    "last_seen": now,
+                    "examples": [],
+                }
+                logger.info("new_document_type_learned", document_type=document_type)
+            
+            dt = self._document_types[document_type]
+            dt["count"] += 1
+            dt["h1_min_sum"] += h1_min
+            dt["h2_min_sum"] += h2_min
+            dt["h3_min_sum"] += h3_min
+            dt["body_size_sum"] += body_size
+            dt["last_seen"] = now
+            
+            if document_example and len(dt["examples"]) < 3:
+                dt["examples"].append(document_example)
+            
+            self._dirty = True
+            
+            if self.auto_save:
+                self._save()
+    
+    def get_thresholds(
+        self,
+        document_type: str | None = None,
+    ) -> dict[str, float]:
+        """
+        Get header thresholds for a document type.
+        
+        Args:
+            document_type: Type of document. If None, returns defaults.
+            
+        Returns:
+            Dict with h1_min, h2_min, h3_min, body_size.
+        """
+        if not document_type:
+            return {
+                "h1_min": self.DEFAULT_H1_MIN,
+                "h2_min": self.DEFAULT_H2_MIN,
+                "h3_min": self.DEFAULT_H3_MIN,
+                "body_size": 12.0,
+            }
+        
+        document_type = document_type.lower().strip()
+        
+        with self._lock:
+            if document_type in self._document_types:
+                dt = self._document_types[document_type]
+                count = dt["count"]
+                
+                if count > 0:
+                    return {
+                        "h1_min": dt["h1_min_sum"] / count,
+                        "h2_min": dt["h2_min_sum"] / count,
+                        "h3_min": dt["h3_min_sum"] / count,
+                        "body_size": dt["body_size_sum"] / count,
+                    }
+        
+        # Return defaults if not found
+        return self.get_thresholds(None)
+    
+    def detect_document_type(
+        self,
+        body_size: float,
+        max_font_size: float,
+        has_legal_terms: bool = False,
+        has_academic_structure: bool = False,
+    ) -> str:
+        """
+        Attempt to detect document type from characteristics.
+        
+        Args:
+            body_size: Most common font size.
+            max_font_size: Largest font in document.
+            has_legal_terms: Whether document contains legal terminology.
+            has_academic_structure: Whether document has academic structure.
+            
+        Returns:
+            Detected document type string.
+        """
+        # Simple heuristics
+        if has_legal_terms:
+            if body_size <= 12 and max_font_size <= 16:
+                return "legal_brief"
+            return "legal_general"
+        
+        if has_academic_structure:
+            return "academic"
+        
+        if max_font_size >= 36:
+            return "marketing"
+        
+        if body_size >= 14:
+            return "presentation"
+        
+        return "general"
+    
+    def get_known_document_types(self) -> list[str]:
+        """Get list of document types we have learned."""
+        with self._lock:
+            return list(self._document_types.keys())
+    
+    def get_stats(self) -> dict[str, Any]:
+        """Get statistics about learned thresholds."""
+        with self._lock:
+            return {
+                "document_types_count": len(self._document_types),
+                "document_types": list(self._document_types.keys()),
+                "total_documents_analyzed": sum(
+                    dt["count"] for dt in self._document_types.values()
+                ),
+            }
+
+
+# Global header thresholds registry
+_global_header_thresholds: LearnedHeaderThresholds | None = None
+
+
+def get_learned_header_thresholds() -> LearnedHeaderThresholds:
+    """Get the global learned header thresholds registry."""
+    global _global_header_thresholds
+    
+    if _global_header_thresholds is None:
+        custom_path = os.getenv("RNSR_HEADER_THRESHOLDS_PATH")
+        _global_header_thresholds = LearnedHeaderThresholds(
+            storage_path=custom_path if custom_path else None
+        )
+    
+    return _global_header_thresholds
+
+
+class HeaderClassifier:
+    """
+    Classifies text spans into headers and body text.
+    
+    Uses font size analysis and optional k-means clustering
+    to determine header levels. Supports adaptive thresholds
+    that learn from document types.
+    """
+
+    # Default thresholds for header levels (in points)
+    H1_MIN_SIZE = 24.0
+    H2_MIN_SIZE = 18.0
+    H3_MIN_SIZE = 14.0
+
+    def __init__(
+        self,
+        use_clustering: bool = True,
+        n_header_levels: int = 3,
+        document_type: str | None = None,
+        enable_threshold_learning: bool = True,
+    ):
+        """
+        Initialize the Header Classifier.
+        
+        Args:
+            use_clustering: Whether to use k-means clustering for header levels.
+            n_header_levels: Number of header levels to detect (default: 3).
+            document_type: Optional document type for adaptive thresholds.
+            enable_threshold_learning: Whether to learn thresholds from documents.
+        """
+        self.use_clustering = use_clustering
+        self.n_header_levels = n_header_levels
+        self.document_type = document_type
+        self.enable_threshold_learning = enable_threshold_learning
+        
+        # Get learned thresholds registry
+        self._threshold_registry = get_learned_header_thresholds() if enable_threshold_learning else None
+        
+        # Set thresholds based on document type
+        self._update_thresholds(document_type)
+    
+    def _update_thresholds(self, document_type: str | None) -> None:
+        """Update thresholds based on document type."""
+        if self._threshold_registry and document_type:
+            thresholds = self._threshold_registry.get_thresholds(document_type)
+            self.H1_MIN_SIZE = thresholds["h1_min"]
+            self.H2_MIN_SIZE = thresholds["h2_min"]
+            self.H3_MIN_SIZE = thresholds["h3_min"]
+        else:
+            # Use class defaults
+            self.H1_MIN_SIZE = HeaderClassifier.H1_MIN_SIZE
+            self.H2_MIN_SIZE = HeaderClassifier.H2_MIN_SIZE
+            self.H3_MIN_SIZE = HeaderClassifier.H3_MIN_SIZE
+    
+    def set_document_type(self, document_type: str) -> None:
+        """Set document type and update thresholds."""
+        self.document_type = document_type
+        self._update_thresholds(document_type)
+    
+    def learn_from_analysis(
+        self,
+        analysis: FontAnalysis,
+        detected_h1_size: float | None = None,
+        detected_h2_size: float | None = None,
+        detected_h3_size: float | None = None,
+        document_name: str = "",
+    ) -> None:
+        """
+        Learn thresholds from document analysis.
+        
+        Call this after processing a document to record observed values.
+        """
+        if not self._threshold_registry or not self.document_type:
+            return
+        
+        # Use detected sizes or infer from analysis
+        h1_size = detected_h1_size or analysis.header_threshold + 6
+        h2_size = detected_h2_size or analysis.header_threshold + 3
+        h3_size = detected_h3_size or analysis.header_threshold
+        
+        self._threshold_registry.record_thresholds(
+            document_type=self.document_type,
+            h1_min=h1_size,
+            h2_min=h2_size,
+            h3_min=h3_size,
+            body_size=analysis.body_size,
+            document_example=document_name,
+        )
+
+    def classify_spans(
+        self,
+        spans: list[SpanInfo],
+        analysis: FontAnalysis,
+    ) -> list[ClassifiedSpan]:
+        """
+        Classify all spans into headers or body text.
+        
+        Args:
+            spans: List of SpanInfo from font analysis.
+            analysis: FontAnalysis with body size and threshold.
+            
+        Returns:
+            List of ClassifiedSpan with role and header_level assigned.
+        """
+        if not spans:
+            return []
+        
+        # First pass: identify potential headers
+        potential_headers: list[SpanInfo] = []
+        for span in spans:
+            if self._is_header_candidate(span, analysis):
+                potential_headers.append(span)
+        
+        # Determine header levels
+        if potential_headers and self.use_clustering:
+            level_mapping = self._cluster_header_levels(potential_headers)
+        else:
+            level_mapping = {}
+        
+        # Classify all spans
+        classified: list[ClassifiedSpan] = []
+        for span in spans:
+            role, level = self._classify_single_span(span, analysis, level_mapping)
+            
+            classified_span = ClassifiedSpan(
+                text=span.text,
+                font_size=span.font_size,
+                font_name=span.font_name,
+                is_bold=span.is_bold,
+                is_italic=span.is_italic,
+                bbox=span.bbox,
+                page_num=span.page_num,
+                role=role,
+                header_level=level,
+            )
+            classified.append(classified_span)
+        
+        # Log classification stats
+        header_count = sum(1 for s in classified if s.role == "header")
+        logger.info(
+            "spans_classified",
+            total=len(classified),
+            headers=header_count,
+            body=len(classified) - header_count,
+        )
+        
+        return classified
+
+    def _is_header_candidate(
+        self, 
+        span: SpanInfo, 
+        analysis: FontAnalysis,
+    ) -> bool:
+        """
+        Determine if a span is a header candidate.
+        
+        A span is a header candidate if:
+        1. Font size > header_threshold, OR
+        2. Bold text at or above body size
+        """
+        # Size-based detection
+        if span.font_size > analysis.header_threshold:
+            return True
+        
+        # Bold text at body size or larger
+        if span.is_bold and span.font_size >= analysis.body_size:
+            return True
+        
+        return False
+
+    def _classify_single_span(
+        self,
+        span: SpanInfo,
+        analysis: FontAnalysis,
+        level_mapping: dict[float, int],
+    ) -> tuple[Literal["header", "body", "caption", "footnote"], int]:
+        """
+        Classify a single span.
+        
+        Returns:
+            Tuple of (role, header_level).
+        """
+        # Check if it's a header
+        if self._is_header_candidate(span, analysis):
+            # Use clustering-based level if available
+            if span.font_size in level_mapping:
+                level = level_mapping[span.font_size]
+            else:
+                # Fall back to absolute thresholds
+                level = self._get_level_by_size(span.font_size)
+            
+            return ("header", level)
+        
+        # Check for captions/footnotes (smaller than body)
+        caption_threshold = analysis.body_size - (analysis.body_size * 0.2)
+        if span.font_size < caption_threshold:
+            return ("caption", 0)
+        
+        # Default to body text
+        return ("body", 0)
+
+    def _get_level_by_size(self, font_size: float) -> int:
+        """
+        Get header level based on absolute font size thresholds.
+        
+        Args:
+            font_size: The font size in points.
+            
+        Returns:
+            Header level (1, 2, or 3).
+        """
+        if font_size >= self.H1_MIN_SIZE:
+            return 1
+        elif font_size >= self.H2_MIN_SIZE:
+            return 2
+        else:
+            return 3
+
+    def _cluster_header_levels(
+        self, 
+        headers: list[SpanInfo],
+    ) -> dict[float, int]:
+        """
+        Use k-means clustering to determine header levels from actual data.
+        
+        This adapts to documents with non-standard font sizes.
+        
+        Args:
+            headers: List of spans identified as headers.
+            
+        Returns:
+            Dict mapping font_size to header_level (1, 2, or 3).
+        """
+        if len(headers) < self.n_header_levels:
+            # Not enough headers to cluster
+            return {}
+        
+        # Get unique font sizes
+        unique_sizes = list(set(h.font_size for h in headers))
+        
+        if len(unique_sizes) < self.n_header_levels:
+            # Fewer unique sizes than levels - assign directly
+            sorted_sizes = sorted(unique_sizes, reverse=True)
+            return {size: i + 1 for i, size in enumerate(sorted_sizes)}
+        
+        # Perform k-means clustering
+        try:
+            X = np.array(unique_sizes).reshape(-1, 1)
+            n_clusters = min(self.n_header_levels, len(unique_sizes))
+            
+            kmeans = KMeans(
+                n_clusters=n_clusters, 
+                random_state=42,
+                n_init=10,
+            ).fit(X)
+            
+            # Map clusters to levels by size (largest = H1)
+            cluster_centers = [
+                (i, kmeans.cluster_centers_[i][0]) 
+                for i in range(n_clusters)
+            ]
+            cluster_centers.sort(key=lambda x: -x[1])  # Descending by size
+            
+            cluster_to_level = {
+                cluster: level + 1 
+                for level, (cluster, _) in enumerate(cluster_centers)
+            }
+            
+            # Map each unique size to its level
+            size_to_level = {}
+            for size in unique_sizes:
+                cluster = kmeans.predict([[size]])[0]
+                size_to_level[size] = cluster_to_level[cluster]
+            
+            logger.debug(
+                "header_levels_clustered",
+                mapping=size_to_level,
+            )
+            
+            return size_to_level
+            
+        except Exception as e:
+            logger.warning("clustering_failed", error=str(e))
+            return {}
+
+
+def classify_headers(
+    spans: list[SpanInfo],
+    analysis: FontAnalysis,
+    use_clustering: bool = True,
+) -> list[ClassifiedSpan]:
+    """
+    Convenience function to classify spans into headers and body text.
+    
+    Args:
+        spans: List of SpanInfo from font analysis.
+        analysis: FontAnalysis with body size and threshold.
+        use_clustering: Whether to use k-means for header levels.
+        
+    Returns:
+        List of ClassifiedSpan with roles assigned.
+        
+    Example:
+        analysis, spans = analyze_font_histogram("doc.pdf")
+        classified = classify_headers(spans, analysis)
+        headers = [s for s in classified if s.role == "header"]
+    """
+    classifier = HeaderClassifier(use_clustering=use_clustering)
+    return classifier.classify_spans(spans, analysis)