kreuzberg 4.0.0.rc2 → 4.0.1

data/vendor/kreuzberg/tests/pdf_hierarchy_quality.rs

@@ -0,0 +1,589 @@
+//! PDF hierarchy quality assessment tests.
+//!
+//! This module tests PDF text hierarchy extraction quality by comparing against ground truth annotations.
+//! Measures precision, recall, F1 score, and level accuracy to ensure the hierarchy detection
+//! algorithm works well on real document structures.
+//!
+//! Test philosophy:
+//! - Define ground truth hierarchies for representative PDF documents
+//! - Measure how well extracted hierarchies match ground truth
+//! - Assert minimum quality thresholds for precision/recall/F1
+//! - Verify correct hierarchy level assignments
+
+#![cfg(feature = "pdf")]
+
+use kreuzberg::pdf::hierarchy::{
+    BoundingBox, HierarchyLevel, KMeansResult, TextBlock, assign_hierarchy_levels,
+    assign_hierarchy_levels_from_clusters, cluster_font_sizes,
+};
+use serde::{Deserialize, Serialize};
+use std::fs;
+use std::path::Path;
+
+/// A bounding box annotation from ground truth.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct GroundTruthBBox {
+    left: f32,
+    top: f32,
+    right: f32,
+    bottom: f32,
+}
+
+impl GroundTruthBBox {
+    /// Convert to kreuzberg BoundingBox
+    fn to_bbox(&self) -> BoundingBox {
+        BoundingBox {
+            left: self.left,
+            top: self.top,
+            right: self.right,
+            bottom: self.bottom,
+        }
+    }
+}
+
+/// A ground truth text block with hierarchy level annotation.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct GroundTruthBlock {
+    text: String,
+    level: String,
+    bbox: GroundTruthBBox,
+}
+
+/// A page of ground truth annotations.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct GroundTruthPage {
+    page_number: u32,
+    blocks: Vec<GroundTruthBlock>,
+}
+
+/// A document with ground truth hierarchy annotations.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct GroundTruthDocument {
+    pdf_file: String,
+    pages: Vec<GroundTruthPage>,
+}
+
+/// Root structure for ground truth JSON file.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct GroundTruthFile {
+    documents: Vec<GroundTruthDocument>,
+}
+
+/// Quality metrics for hierarchy extraction.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct QualityMetrics {
+    /// Precision: (correctly identified hierarchies) / (total extracted hierarchies)
+    pub precision: f64,
+    /// Recall: (correctly identified hierarchies) / (total ground truth hierarchies)
+    pub recall: f64,
+    /// F1 Score: harmonic mean of precision and recall
+    pub f1_score: f64,
+    /// Level accuracy: percentage of blocks assigned to correct hierarchy level
+    pub level_accuracy: f64,
+    /// Number of correctly identified hierarchy blocks
+    pub true_positives: usize,
+    /// Number of incorrectly extracted hierarchy blocks
+    pub false_positives: usize,
+    /// Number of missed ground truth hierarchy blocks
+    pub false_negatives: usize,
+    /// Number of blocks with correct hierarchy level
+    pub correct_levels: usize,
+    /// Total number of blocks evaluated
+    pub total_blocks: usize,
+}
+
+impl QualityMetrics {
+    /// Create new quality metrics from test results.
+    fn new(
+        true_positives: usize,
+        false_positives: usize,
+        false_negatives: usize,
+        correct_levels: usize,
+        total_blocks: usize,
+    ) -> Self {
+        let precision = if true_positives + false_positives > 0 {
+            true_positives as f64 / (true_positives + false_positives) as f64
+        } else {
+            0.0
+        };
+
+        let recall = if true_positives + false_negatives > 0 {
+            true_positives as f64 / (true_positives + false_negatives) as f64
+        } else {
+            0.0
+        };
+
+        let f1_score = if precision + recall > 0.0 {
+            2.0 * precision * recall / (precision + recall)
+        } else {
+            0.0
+        };
+
+        let level_accuracy = if total_blocks > 0 {
+            correct_levels as f64 / total_blocks as f64
+        } else {
+            0.0
+        };
+
+        Self {
+            precision,
+            recall,
+            f1_score,
+            level_accuracy,
+            true_positives,
+            false_positives,
+            false_negatives,
+            correct_levels,
+            total_blocks,
+        }
+    }
+}
+
+/// Convert hierarchy level string to HierarchyLevel enum.
+fn parse_level(level: &str) -> HierarchyLevel {
+    match level {
+        "H1" => HierarchyLevel::H1,
+        "H2" => HierarchyLevel::H2,
+        "H3" => HierarchyLevel::H3,
+        "H4" => HierarchyLevel::H4,
+        "H5" => HierarchyLevel::H5,
+        "H6" => HierarchyLevel::H6,
+        _ => HierarchyLevel::Body,
+    }
+}
+
+/// Load ground truth annotations from JSON file.
+///
+/// Reads the hierarchy_ground_truth.json file and parses document annotations.
+///
+/// # Arguments
+///
+/// * `path` - Path to the ground truth JSON file
+///
+/// # Returns
+///
+/// Result containing the parsed GroundTruthFile or error message
+fn load_ground_truth<P: AsRef<Path>>(path: P) -> Result<GroundTruthFile, String> {
+    let content = fs::read_to_string(path).map_err(|e| format!("Failed to read file: {}", e))?;
+    serde_json::from_str(&content).map_err(|e| format!("Failed to parse JSON: {}", e))
+}
+
+/// Calculate quality metrics by comparing extracted hierarchies to ground truth.
+///
+/// Compares extracted text blocks with their hierarchy assignments to ground truth annotations.
+/// Measures:
+/// - Precision: correctly identified hierarchies / total extracted
+/// - Recall: correctly identified hierarchies / total ground truth
+/// - F1 Score: harmonic mean of precision and recall
+/// - Level Accuracy: percentage of blocks with correct hierarchy level
+///
+/// # Arguments
+///
+/// * `extracted_blocks` - Vector of extracted HierarchyBlock objects
+/// * `ground_truth_blocks` - Vector of ground truth blocks
+///
+/// # Returns
+///
+/// QualityMetrics struct with calculated precision, recall, F1, and level accuracy
+fn calculate_quality_metrics(
+    extracted_blocks: &[kreuzberg::pdf::hierarchy::HierarchyBlock],
+    ground_truth_blocks: &[GroundTruthBlock],
+) -> QualityMetrics {
+    let mut true_positives = 0;
+    let mut false_positives = 0;
+    let mut correct_levels = 0;
+
+    // For matching blocks, we use bounding box overlap and text similarity
+    let mut matched_gt_indices: Vec<bool> = vec![false; ground_truth_blocks.len()];
+
+    for extracted in extracted_blocks {
+        let mut best_match_idx: Option<usize> = None;
+        let mut best_overlap = 0.0;
+
+        // Find the best matching ground truth block by bounding box overlap
+        for (gt_idx, gt_block) in ground_truth_blocks.iter().enumerate() {
+            if matched_gt_indices[gt_idx] {
+                continue; // Already matched
+            }
+
+            let gt_bbox = gt_block.bbox.to_bbox();
+            let overlap = extracted.bbox.iou(&gt_bbox);
+
+            if overlap > best_overlap && overlap > 0.3 {
+                best_overlap = overlap;
+                best_match_idx = Some(gt_idx);
+            }
+        }
+
+        if let Some(gt_idx) = best_match_idx {
+            matched_gt_indices[gt_idx] = true;
+            true_positives += 1;
+
+            // Check if the hierarchy level matches
+            let gt_level = parse_level(&ground_truth_blocks[gt_idx].level);
+            if extracted.hierarchy_level == gt_level {
+                correct_levels += 1;
+            }
+        } else {
+            false_positives += 1;
+        }
+    }
+
+    // Count unmatched ground truth blocks as false negatives
+    let false_negatives = matched_gt_indices.iter().filter(|&&m| !m).count();
+
+    let total_blocks = extracted_blocks.len().max(ground_truth_blocks.len());
+
+    QualityMetrics::new(
+        true_positives,
+        false_positives,
+        false_negatives,
+        correct_levels,
+        total_blocks,
+    )
+}
+
+/// Create test text blocks from ground truth.
+fn create_text_blocks_from_ground_truth(gt_blocks: &[GroundTruthBlock]) -> Vec<TextBlock> {
+    gt_blocks
+        .iter()
+        .enumerate()
+        .map(|(idx, gt_block)| {
+            // Estimate font size from bbox height
+            let bbox = gt_block.bbox.to_bbox();
+            let font_size = match gt_block.level.as_str() {
+                "H1" => 28.0,
+                "H2" => 24.0,
+                "H3" => 20.0,
+                "H4" => 16.0,
+                "H5" => 14.0,
+                "H6" => 12.0,
+                _ => 10.0, // Body
+            };
+
+            TextBlock {
+                text: if gt_block.text.len() > 50 {
+                    format!("{} (Block {})", gt_block.text.chars().take(50).collect::<String>(), idx)
+                } else {
+                    gt_block.text.clone()
+                },
+                bbox,
+                font_size,
+            }
+        })
+        .collect()
+}
+
+#[test]
+fn test_hierarchy_quality_on_ground_truth() {
+    // Load ground truth data
+    let ground_truth_path = "tests/data/hierarchy_ground_truth.json";
+    let ground_truth_file = load_ground_truth(ground_truth_path).expect("Failed to load ground truth file");
+
+    println!(
+        "\nLoaded {} documents from ground truth",
+        ground_truth_file.documents.len()
+    );
+
+    let mut all_metrics: Vec<QualityMetrics> = Vec::new();
+
+    // Process each document
+    for doc in &ground_truth_file.documents {
+        println!("\nProcessing document: {}", doc.pdf_file);
+
+        for page in &doc.pages {
+            println!("  Page {}: {} blocks", page.page_number, page.blocks.len());
+
+            // Create text blocks from ground truth
+            let text_blocks = create_text_blocks_from_ground_truth(&page.blocks);
+
+            // Cluster by font size
+            let k = (text_blocks.len() / 3).clamp(1, 6); // Estimate k clusters
+            let clusters = cluster_font_sizes(&text_blocks, k).expect("Failed to cluster font sizes");
+
+            println!(
+                "    Created {} clusters from {} blocks",
+                clusters.len(),
+                text_blocks.len()
+            );
+
+            // Assign hierarchy levels from clusters
+            let hierarchy_assignments = assign_hierarchy_levels_from_clusters(&text_blocks, &clusters);
+
+            // Convert to HierarchyBlock format
+            let extracted_blocks: Vec<kreuzberg::pdf::hierarchy::HierarchyBlock> = hierarchy_assignments
+                .iter()
+                .map(|(block, level)| kreuzberg::pdf::hierarchy::HierarchyBlock {
+                    text: block.text.clone(),
+                    bbox: block.bbox,
+                    font_size: block.font_size,
+                    hierarchy_level: *level,
+                })
+                .collect();
+
+            // Calculate quality metrics
+            let metrics = calculate_quality_metrics(&extracted_blocks, &page.blocks);
+            all_metrics.push(metrics.clone());
+
+            println!("    Precision:      {:.4}", metrics.precision);
+            println!("    Recall:         {:.4}", metrics.recall);
+            println!("    F1 Score:       {:.4}", metrics.f1_score);
+            println!("    Level Accuracy: {:.4}", metrics.level_accuracy);
+        }
+    }
+
+    // Calculate average metrics
+    if !all_metrics.is_empty() {
+        let avg_precision = all_metrics.iter().map(|m| m.precision).sum::<f64>() / all_metrics.len() as f64;
+        let avg_recall = all_metrics.iter().map(|m| m.recall).sum::<f64>() / all_metrics.len() as f64;
+        let avg_f1 = all_metrics.iter().map(|m| m.f1_score).sum::<f64>() / all_metrics.len() as f64;
+        let avg_level_acc = all_metrics.iter().map(|m| m.level_accuracy).sum::<f64>() / all_metrics.len() as f64;
+
+        println!("\n=== AVERAGE METRICS ACROSS ALL PAGES ===");
+        println!("Average Precision:      {:.4}", avg_precision);
+        println!("Average Recall:         {:.4}", avg_recall);
+        println!("Average F1 Score:       {:.4}", avg_f1);
+        println!("Average Level Accuracy: {:.4}", avg_level_acc);
+
+        // Assert minimum F1 threshold
+        assert!(
+            avg_f1 > 0.85,
+            "F1 score ({:.4}) must be greater than 0.85. Metrics: precision={:.4}, recall={:.4}, level_accuracy={:.4}",
+            avg_f1,
+            avg_precision,
+            avg_recall,
+            avg_level_acc
+        );
+    }
+}
+
+#[test]
+fn test_hierarchy_clustering_consistency() {
+    // Arrange: Create a simple document with clear hierarchy
+    let blocks = vec![
+        TextBlock {
+            text: "Title".to_string(),
+            bbox: BoundingBox {
+                left: 0.0,
+                top: 0.0,
+                right: 100.0,
+                bottom: 28.0,
+            },
+            font_size: 28.0,
+        },
+        TextBlock {
+            text: "Subtitle".to_string(),
+            bbox: BoundingBox {
+                left: 0.0,
+                top: 30.0,
+                right: 100.0,
+                bottom: 54.0,
+            },
+            font_size: 24.0,
+        },
+        TextBlock {
+            text: "Section".to_string(),
+            bbox: BoundingBox {
+                left: 0.0,
+                top: 60.0,
+                right: 100.0,
+                bottom: 80.0,
+            },
+            font_size: 20.0,
+        },
+        TextBlock {
+            text: "Body paragraph".to_string(),
+            bbox: BoundingBox {
+                left: 0.0,
+                top: 90.0,
+                right: 100.0,
+                bottom: 102.0,
+            },
+            font_size: 10.0,
+        },
+    ];
+
+    // Act: Cluster and assign hierarchies
+    let clusters = cluster_font_sizes(&blocks, 4).expect("Clustering failed");
+    let assignments = assign_hierarchy_levels_from_clusters(&blocks, &clusters);
+
+    // Assert: Verify hierarchy levels are correct
+    assert_eq!(assignments.len(), 4);
+    assert_eq!(assignments[0].1, HierarchyLevel::H1, "Largest text should be H1");
+    assert_eq!(assignments[1].1, HierarchyLevel::H2, "Second largest should be H2");
+    assert_eq!(assignments[2].1, HierarchyLevel::H3, "Third largest should be H3");
+    assert_eq!(assignments[3].1, HierarchyLevel::Body, "Smallest text should be Body");
+
+    // Assert: F1 score should be perfect for this simple case
+    let quality_metrics = calculate_quality_metrics(
+        &assignments
+            .iter()
+            .map(|(b, l)| kreuzberg::pdf::hierarchy::HierarchyBlock {
+                text: b.text.clone(),
+                bbox: b.bbox,
+                font_size: b.font_size,
+                hierarchy_level: *l,
+            })
+            .collect::<Vec<_>>(),
+        &[
+            GroundTruthBlock {
+                text: "Title".to_string(),
+                level: "H1".to_string(),
+                bbox: GroundTruthBBox {
+                    left: 0.0,
+                    top: 0.0,
+                    right: 100.0,
+                    bottom: 28.0,
+                },
+            },
+            GroundTruthBlock {
+                text: "Subtitle".to_string(),
+                level: "H2".to_string(),
+                bbox: GroundTruthBBox {
+                    left: 0.0,
+                    top: 30.0,
+                    right: 100.0,
+                    bottom: 54.0,
+                },
+            },
+            GroundTruthBlock {
+                text: "Section".to_string(),
+                level: "H3".to_string(),
+                bbox: GroundTruthBBox {
+                    left: 0.0,
+                    top: 60.0,
+                    right: 100.0,
+                    bottom: 80.0,
+                },
+            },
+            GroundTruthBlock {
+                text: "Body paragraph".to_string(),
+                level: "Body".to_string(),
+                bbox: GroundTruthBBox {
+                    left: 0.0,
+                    top: 90.0,
+                    right: 100.0,
+                    bottom: 102.0,
+                },
+            },
+        ],
+    );
+
+    println!("Consistency Test - F1 Score: {:.4}", quality_metrics.f1_score);
+    assert!(
+        quality_metrics.f1_score >= 0.8,
+        "F1 score for simple hierarchy should be >= 0.8"
+    );
+}
+
+#[test]
+fn test_hierarchy_level_assignment() {
+    // Arrange: Create blocks and KMeans result
+    let blocks = vec![
+        TextBlock {
+            text: "Main Title".to_string(),
+            bbox: BoundingBox {
+                left: 50.0,
+                top: 50.0,
+                right: 150.0,
+                bottom: 100.0,
+            },
+            font_size: 28.0,
+        },
+        TextBlock {
+            text: "Section Title".to_string(),
+            bbox: BoundingBox {
+                left: 50.0,
+                top: 120.0,
+                right: 150.0,
+                bottom: 160.0,
+            },
+            font_size: 20.0,
+        },
+        TextBlock {
+            text: "Regular body text".to_string(),
+            bbox: BoundingBox {
+                left: 50.0,
+                top: 180.0,
+                right: 200.0,
+                bottom: 200.0,
+            },
+            font_size: 12.0,
+        },
+    ];
+
+    let kmeans_result = KMeansResult { labels: vec![0, 1, 2] };
+
+    // Act: Assign hierarchy levels using KMeans result
+    let result = assign_hierarchy_levels(&blocks, &kmeans_result);
+
+    // Assert: Verify correct level assignments
+    assert_eq!(result.len(), 3);
+    assert_eq!(result[0].hierarchy_level, HierarchyLevel::H1);
+    assert_eq!(result[1].hierarchy_level, HierarchyLevel::H2);
+    assert_eq!(result[2].hierarchy_level, HierarchyLevel::H3);
+}
+
+#[test]
+fn test_quality_metrics_calculation() {
+    // Arrange: Create extracted blocks and ground truth
+    let extracted = vec![
+        kreuzberg::pdf::hierarchy::HierarchyBlock {
+            text: "Title".to_string(),
+            bbox: BoundingBox {
+                left: 0.0,
+                top: 0.0,
+                right: 100.0,
+                bottom: 20.0,
+            },
+            font_size: 28.0,
+            hierarchy_level: HierarchyLevel::H1,
+        },
+        kreuzberg::pdf::hierarchy::HierarchyBlock {
+            text: "Body".to_string(),
+            bbox: BoundingBox {
+                left: 0.0,
+                top: 30.0,
+                right: 100.0,
+                bottom: 50.0,
+            },
+            font_size: 12.0,
+            hierarchy_level: HierarchyLevel::Body,
+        },
+    ];
+
+    let ground_truth = vec![
+        GroundTruthBlock {
+            text: "Title".to_string(),
+            level: "H1".to_string(),
+            bbox: GroundTruthBBox {
+                left: 0.0,
+                top: 0.0,
+                right: 100.0,
+                bottom: 20.0,
+            },
+        },
+        GroundTruthBlock {
+            text: "Body".to_string(),
+            level: "Body".to_string(),
+            bbox: GroundTruthBBox {
+                left: 0.0,
+                top: 30.0,
+                right: 100.0,
+                bottom: 50.0,
+            },
+        },
+    ];
+
+    // Act: Calculate metrics
+    let metrics = calculate_quality_metrics(&extracted, &ground_truth);
+
+    // Assert: Verify metrics
+    assert_eq!(metrics.true_positives, 2);
+    assert_eq!(metrics.false_positives, 0);
+    assert_eq!(metrics.false_negatives, 0);
+    assert_eq!(metrics.correct_levels, 2);
+    assert!(metrics.precision > 0.99);
+    assert!(metrics.recall > 0.99);
+    assert!(metrics.f1_score > 0.99);
+}

data/vendor/kreuzberg/tests/pdf_integration.rs

@@ -1,43 +1,45 @@
-//! PDF integration tests that remain specific to the Rust core.
-//!
-//! Positive-path scenarios live in the shared fixtures that back the
-//! multi-language E2E generator. This module keeps only the cases that
-//! exercise Rust-specific failure handling or error propagation.
-
-mod helpers;
-
-use helpers::*;
-use kreuzberg::core::config::ExtractionConfig;
-use kreuzberg::extract_file_sync;
-
-#[test]
-fn test_pdf_password_protected_fails_gracefully() {
-    if skip_if_missing("pdfs/copy_protected.pdf") {
-        return;
-    }
-
-    let file_path = get_test_file_path("pdfs/copy_protected.pdf");
-    let result = extract_file_sync(&file_path, None, &ExtractionConfig::default());
-
-    match result {
-        Ok(extraction_result) => {
-            assert_mime_type(&extraction_result, "application/pdf");
-            assert!(
-                extraction_result.chunks.is_none(),
-                "Chunks should be None without chunking config"
-            );
-            assert!(
-                extraction_result.detected_languages.is_none(),
-                "Language detection not enabled"
-            );
-        }
-        Err(e) => {
-            let error_msg = e.to_string().to_lowercase();
-            assert!(
-                error_msg.contains("password") || error_msg.contains("protected") || error_msg.contains("encrypted"),
-                "Error message should indicate password/protection issue, got: {}",
-                e
-            );
-        }
-    }
-}
+//! PDF integration tests that remain specific to the Rust core.
+//!
+//! Positive-path scenarios live in the shared fixtures that back the
+//! multi-language E2E generator. This module keeps only the cases that
+//! exercise Rust-specific failure handling or error propagation.
+
+#![cfg(feature = "pdf")]
+
+mod helpers;
+
+use helpers::*;
+use kreuzberg::core::config::ExtractionConfig;
+use kreuzberg::extract_file_sync;
+
+#[test]
+fn test_pdf_password_protected_fails_gracefully() {
+    if skip_if_missing("pdfs/copy_protected.pdf") {
+        return;
+    }
+
+    let file_path = get_test_file_path("pdfs/copy_protected.pdf");
+    let result = extract_file_sync(&file_path, None, &ExtractionConfig::default());
+
+    match result {
+        Ok(extraction_result) => {
+            assert_mime_type(&extraction_result, "application/pdf");
+            assert!(
+                extraction_result.chunks.is_none(),
+                "Chunks should be None without chunking config"
+            );
+            assert!(
+                extraction_result.detected_languages.is_none(),
+                "Language detection not enabled"
+            );
+        }
+        Err(e) => {
+            let error_msg = e.to_string().to_lowercase();
+            assert!(
+                error_msg.contains("password") || error_msg.contains("protected") || error_msg.contains("encrypted"),
+                "Error message should indicate password/protection issue, got: {}",
+                e
+            );
+        }
+    }
+}