rnsr 0.1.0__py3-none-any.whl

rnsr/benchmarks/finance_bench.py

@@ -0,0 +1,147 @@
+"""
+FinanceBench Dataset Loader for RNSR Evaluation
+
+FinanceBench is a benchmark for financial question answering (QA) using large language models (LLMs).
+It assesses the ability of LLMs to answer questions about financial documents, requiring retrieval
+from complex PDFs (tables, charts, etc.).
+
+Repository: https://huggingface.co/datasets/PatronusAI/financebench
+"""
+
+import os
+import requests
+import hashlib
+from pathlib import Path
+from dataclasses import dataclass
+from typing import List, Dict, Any, Optional
+
+import structlog
+from datasets import load_dataset  # type: ignore
+
+from rnsr.benchmarks.standard_benchmarks import BenchmarkDataset, BenchmarkQuestion
+
+logger = structlog.get_logger(__name__)
+
+CACHE_DIR = Path("rnsr/benchmarks/data/financebench")
+
+class FinanceBenchLoader:
+    """Loader for the FinanceBench dataset."""
+    
+    @staticmethod
+    def _download_pdf(url: str, doc_name: str) -> Optional[Path]:
+        """
+        Download PDF from URL and cache it locally.
+        Returns the path to the cached PDF.
+        """
+        if not url:
+            return None
+            
+        # Create a safe filename (hash + original name sanitized)
+        url_hash = hashlib.md5(url.encode()).hexdigest()[:8]
+        safe_name = "".join(c for c in doc_name if c.isalnum() or c in (' ', '.', '_', '-')).strip()
+        safe_name = safe_name.replace(" ", "_")
+        if not safe_name.lower().endswith(".pdf"):
+            safe_name += ".pdf"
+            
+        file_path = CACHE_DIR / f"{url_hash}_{safe_name}"
+        
+        if file_path.exists():
+            return file_path
+        
+        try:
+            CACHE_DIR.mkdir(parents=True, exist_ok=True)
+            logger.info("Downloading PDF", url=url, path=str(file_path))
+            response = requests.get(url, timeout=30)
+            response.raise_for_status()
+            
+            with open(file_path, "wb") as f:
+                f.write(response.content)
+                
+            return file_path
+        except Exception as e:
+            logger.error("Failed to download PDF", url=url, error=str(e))
+            return None
+
+    @staticmethod
+    def load(
+        split: str = "train",
+        max_samples: Optional[int] = None,
+        download_pdfs: bool = True
+    ) -> BenchmarkDataset:
+        """
+        Load the FinanceBench dataset.
+        
+        Args:
+            split: Dataset split to load (usually 'train' as test is hidden or same)
+            max_samples: Max number of questions to load
+            download_pdfs: Whether to download the referenced PDFs
+            
+        Returns:
+            BenchmarkDataset containing FinanceBench questions
+        """
+        try:
+            dataset = load_dataset("PatronusAI/financebench", split=split)
+        except Exception as e:
+            logger.error("Failed to load FinanceBench dataset", error=str(e))
+            return BenchmarkDataset(
+                name="FinanceBench",
+                description="Financial QA (Failed to load)",
+                questions=[],
+                metrics=[],
+                source_url=""
+            )
+        
+        questions: List[BenchmarkQuestion] = []
+        
+        # FinanceBench structure:
+        # question, answer, evidence_text, doc_name, doc_link, id, etc.
+        
+        count = 0
+        for item in dataset:
+            if not isinstance(item, dict):
+                continue
+
+            if max_samples and count >= max_samples:
+                break
+                
+            doc_link = item.get("doc_link")
+            doc_name = item.get("doc_name", "unknown_doc")
+            
+            pdf_path = None
+            if download_pdfs and doc_link:
+                # Some links might be missing or broken, handle gracefully
+                pdf_path = FinanceBenchLoader._download_pdf(doc_link, doc_name)
+            
+            # If we couldn't get the PDF, we might skip or mark it
+            # But we'll include it with metadata indicating missing file for now
+            
+            # Extract answer - typically a string in FinanceBench
+            answer = item.get("answer", "")
+            
+            # Create question object
+            q = BenchmarkQuestion(
+                id=f"fb_{count}",  # FinanceBench doesn't have stable IDs in some versions
+                question=item["question"],
+                answer=str(answer),
+                supporting_facts=[item.get("evidence_text", "")],
+                context=[],  # Context is the PDF file, not text chunks
+                reasoning_type="financial-retrieval",
+                metadata={
+                    "doc_name": doc_name,
+                    "doc_link": doc_link,
+                    "pdf_path": str(pdf_path) if pdf_path else None,
+                    "page_index": item.get("page_index"), # Sometimes available
+                    "dataset": "financebench"
+                }
+            )
+            
+            questions.append(q)
+            count += 1
+            
+        return BenchmarkDataset(
+            name="FinanceBench",
+            description="Financial Question Answering on Complex PDFs",
+            questions=questions,
+            metrics=["answer_evaluation_llm"], # Will need LLM-based eval
+            source_url="https://huggingface.co/datasets/PatronusAI/financebench"
+        )

rnsr/benchmarks/pdf_merger.py

@@ -0,0 +1,178 @@
+"""
+PDF Merger Utility for "Chaos Mode" Benchmarking.
+
+This tool merges multiple random PDFs into a single "Frankenstein" document
+to test retrieving information from a specific sub-document within a larger,
+noisy context. This simulates searching through a "Binder" or a merged scan package.
+"""
+
+import random
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+import structlog
+
+# Try to import fitz (PyMuPDF)
+try:
+    import fitz  # type: ignore
+except ImportError:
+    fitz = None
+
+logger = structlog.get_logger(__name__)
+
+class PDFMerger:
+    """Helper to merge PDFs for chaos testing."""
+    
+    @staticmethod
+    def merge_pdfs(
+        target_pdf_path: Path, 
+        distractor_pdf_paths: List[Path], 
+        output_path: Path,
+        insert_position: str = "random"
+    ) -> Dict[str, Any]:
+        """
+        Merge target_pdf into a list of distractor_pdfs.
+        
+        Args:
+            target_pdf_path: The PDF containing the answer.
+            distractor_pdf_paths: List of irrelevant PDFs.
+            output_path: Where to save the merged file.
+            insert_position: 'start', 'end', or 'random'.
+            
+        Returns:
+            Metadata about the merge (page ranges).
+        """
+        if not fitz:
+            raise ImportError("PyMuPDF (fitz) is required for PDF merging. Install 'pymupdf'.")
+            
+        merged_doc = fitz.open()
+        
+        # Prepare list of docs to merge
+        # (doc_object, is_target, label)
+        docs_to_merge = []
+        
+        # Load distractors
+        for p in distractor_pdf_paths:
+            try:
+                doc = fitz.open(p)
+                docs_to_merge.append((doc, False, p.name))
+            except Exception as e:
+                logger.warning(f"Could not open distractor {p}: {e}")
+
+        # Load target
+        try:
+            target_doc = fitz.open(target_pdf_path)
+            target_entry = (target_doc, True, target_pdf_path.name)
+        except Exception as e:
+            logger.error(f"Could not open target PDF {target_pdf_path}: {e}")
+            return {}
+
+        # Insert target
+        if insert_position == "start":
+            docs_to_merge.insert(0, target_entry)
+        elif insert_position == "end":
+            docs_to_merge.append(target_entry)
+        else: # random
+            idx = random.randint(0, len(docs_to_merge))
+            docs_to_merge.insert(idx, target_entry)
+            
+        # Perform merge
+        current_page = 0
+        target_page_range = (0, 0)
+        
+        for doc, is_target, label in docs_to_merge:
+            page_count = doc.page_count
+            merged_doc.insert_pdf(doc)
+            
+            start = current_page
+            end = current_page + page_count - 1
+            
+            if is_target:
+                target_page_range = (start, end)
+                
+            current_page += page_count
+            doc.close()
+            
+        # Save
+        merged_doc.save(output_path)
+        merged_doc.close()
+        
+        return {
+            "merged_file": str(output_path),
+            "target_filename": target_pdf_path.name,
+            "target_page_range": target_page_range, # 0-indexed, inclusive
+            "total_pages": current_page
+        }
+
+    @staticmethod
+    def create_chaos_dataset(
+        base_dataset_questions, 
+        pool_of_pdfs: List[Path], 
+        output_dir: Path,
+        num_distractors: int = 3
+    ):
+        """
+        Takes a list of BenchmarkQuestions (e.g. from FinanceBench) and creates
+        a new version where each source PDF is merged with random distractors.
+        """
+        output_dir.mkdir(parents=True, exist_ok=True)
+        chaos_questions = []
+        
+        for q in base_dataset_questions:
+            try:
+                metadata = q.metadata or {}
+                original_pdf = metadata.get("pdf_path")
+                
+                if not original_pdf:
+                    continue
+                    
+                target_path = Path(original_pdf)
+                if not target_path.exists():
+                    logger.warning(f"Original PDF not found: {target_path}")
+                    continue
+                
+                # Select random distractors (excluding self)
+                candidates = [p for p in pool_of_pdfs if p.name != target_path.name]
+                if len(candidates) < num_distractors:
+                    distractors = candidates
+                else:
+                    distractors = random.sample(candidates, num_distractors)
+                
+                # Merge
+                merged_filename = f"chaos_{q.id}.pdf"
+                merged_path = output_dir / merged_filename
+                
+                merge_info = PDFMerger.merge_pdfs(
+                    target_path, distractors, merged_path
+                )
+                
+                # Clone question and update metadata
+                new_meta = metadata.copy()
+                new_meta.update({
+                    "original_pdf_path": str(target_path),
+                    "pdf_path": str(merged_path), # Point to the chaotic file
+                    "chaos_mode": True,
+                    "target_page_range": merge_info["target_page_range"],
+                    "distractors": [d.name for d in distractors]
+                })
+                
+                # Update question text slightly to hint? No, let's keep it hard.
+                # q.question = q.question # Unchanged
+                
+                # Update question object (requires creating new instance as it might be immutable-ish)
+                from rnsr.benchmarks.standard_benchmarks import BenchmarkQuestion
+                new_q = BenchmarkQuestion(
+                    id=f"{q.id}_chaos",
+                    question=q.question,
+                    answer=q.answer,
+                    supporting_facts=q.supporting_facts,
+                    context=q.context, # Usually empty for PDF benchmarks
+                    reasoning_type=q.reasoning_type,
+                    metadata=new_meta
+                )
+                chaos_questions.append(new_q)
+                
+            except Exception as e:
+                logger.error(f"Failed to process chaos for question {q.id}: {e}")
+                
+        return chaos_questions
+

rnsr/benchmarks/performance.py

@@ -0,0 +1,321 @@
+"""
+Performance Benchmarks
+
+Measures:
+- Ingestion time (PDF → Tree)
+- Indexing time (Tree → Skeleton + KV)
+- Query latency (Question → Answer)
+- Memory usage
+- Throughput (pages/sec, queries/sec)
+"""
+
+from __future__ import annotations
+
+import gc
+import time
+import tracemalloc
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+@dataclass
+class BenchmarkResult:
+    """Result from a single benchmark run."""
+    
+    name: str
+    duration_seconds: float
+    memory_peak_mb: float
+    memory_current_mb: float
+    throughput: float | None = None
+    throughput_unit: str = ""
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    @property
+    def duration_ms(self) -> float:
+        return self.duration_seconds * 1000
+    
+    def __str__(self) -> str:
+        parts = [
+            f"{self.name}:",
+            f"  Time: {self.duration_ms:.2f}ms",
+            f"  Memory Peak: {self.memory_peak_mb:.2f}MB",
+        ]
+        if self.throughput:
+            parts.append(f"  Throughput: {self.throughput:.2f} {self.throughput_unit}")
+        return "\n".join(parts)
+
+
+@dataclass
+class PerformanceBenchmark:
+    """Collection of performance benchmark results."""
+    
+    ingestion: BenchmarkResult | None = None
+    indexing: BenchmarkResult | None = None
+    query_cold: BenchmarkResult | None = None
+    query_warm: BenchmarkResult | None = None
+    total_time_seconds: float = 0.0
+    
+    def summary(self) -> dict[str, Any]:
+        """Get summary statistics."""
+        return {
+            "ingestion_ms": self.ingestion.duration_ms if self.ingestion else None,
+            "indexing_ms": self.indexing.duration_ms if self.indexing else None,
+            "query_cold_ms": self.query_cold.duration_ms if self.query_cold else None,
+            "query_warm_ms": self.query_warm.duration_ms if self.query_warm else None,
+            "total_time_seconds": self.total_time_seconds,
+            "peak_memory_mb": max(
+                r.memory_peak_mb for r in [self.ingestion, self.indexing, self.query_cold]
+                if r is not None
+            ) if any([self.ingestion, self.indexing, self.query_cold]) else 0,
+        }
+    
+    def __str__(self) -> str:
+        lines = ["=== Performance Benchmark Results ==="]
+        for result in [self.ingestion, self.indexing, self.query_cold, self.query_warm]:
+            if result:
+                lines.append(str(result))
+        lines.append(f"\nTotal Time: {self.total_time_seconds:.2f}s")
+        return "\n".join(lines)
+
+
+def _measure_execution(
+    func: Callable[[], Any],
+    name: str,
+    warmup_runs: int = 0,
+) -> tuple[Any, BenchmarkResult]:
+    """
+    Execute a function and measure time and memory.
+    
+    Args:
+        func: Function to execute
+        name: Name for the benchmark
+        warmup_runs: Number of warmup runs before measurement
+        
+    Returns:
+        Tuple of (function result, BenchmarkResult)
+    """
+    # Warmup
+    for _ in range(warmup_runs):
+        func()
+        gc.collect()
+    
+    # Force garbage collection before measurement
+    gc.collect()
+    
+    # Start memory tracking
+    tracemalloc.start()
+    
+    # Time the execution
+    start_time = time.perf_counter()
+    result = func()
+    end_time = time.perf_counter()
+    
+    # Get memory stats
+    current, peak = tracemalloc.get_traced_memory()
+    tracemalloc.stop()
+    
+    benchmark_result = BenchmarkResult(
+        name=name,
+        duration_seconds=end_time - start_time,
+        memory_peak_mb=peak / 1024 / 1024,
+        memory_current_mb=current / 1024 / 1024,
+    )
+    
+    return result, benchmark_result
+
+
+def run_ingestion_benchmark(
+    pdf_path: Path | str,
+    iterations: int = 1,
+) -> BenchmarkResult:
+    """
+    Benchmark PDF ingestion.
+    
+    Args:
+        pdf_path: Path to PDF file
+        iterations: Number of iterations to average
+        
+    Returns:
+        BenchmarkResult with timing and memory stats
+    """
+    from rnsr.ingestion import ingest_document
+    
+    pdf_path = Path(pdf_path)
+    
+    if not pdf_path.exists():
+        raise FileNotFoundError(f"PDF not found: {pdf_path}")
+    
+    # Get file size
+    file_size_mb = pdf_path.stat().st_size / 1024 / 1024
+    
+    results: list[BenchmarkResult] = []
+    final_result = None
+    
+    for i in range(iterations):
+        logger.info("ingestion_benchmark_iteration", iteration=i + 1, total=iterations)
+        
+        def run_ingestion():
+            return ingest_document(pdf_path)
+        
+        ingestion_result, benchmark = _measure_execution(
+            run_ingestion,
+            f"Ingestion (iter {i + 1})",
+        )
+        
+        # Calculate throughput (pages/sec)
+        if hasattr(ingestion_result, 'tree') and ingestion_result.tree:
+            page_count = ingestion_result.tree.total_nodes
+            benchmark.throughput = page_count / benchmark.duration_seconds
+            benchmark.throughput_unit = "nodes/sec"
+        
+        benchmark.metadata = {
+            "file": pdf_path.name,
+            "file_size_mb": file_size_mb,
+            "tier_used": ingestion_result.tier_used if hasattr(ingestion_result, 'tier_used') else None,
+        }
+        
+        results.append(benchmark)
+        final_result = ingestion_result
+    
+    # Average results
+    avg_duration = sum(r.duration_seconds for r in results) / len(results)
+    max_memory = max(r.memory_peak_mb for r in results)
+    avg_throughput = sum(r.throughput or 0 for r in results) / len(results)
+    
+    return BenchmarkResult(
+        name="Ingestion",
+        duration_seconds=avg_duration,
+        memory_peak_mb=max_memory,
+        memory_current_mb=results[-1].memory_current_mb,
+        throughput=avg_throughput if avg_throughput > 0 else None,
+        throughput_unit="nodes/sec",
+        metadata={
+            "iterations": iterations,
+            "file": pdf_path.name,
+            "file_size_mb": file_size_mb,
+        },
+    )
+
+
+def run_query_benchmark(
+    questions: list[str],
+    skeleton: dict,
+    kv_store: Any,
+    warmup: bool = True,
+) -> tuple[BenchmarkResult, BenchmarkResult]:
+    """
+    Benchmark query execution.
+    
+    Args:
+        questions: List of questions to benchmark
+        skeleton: Skeleton index
+        kv_store: KV store with content
+        warmup: Whether to do warmup run
+        
+    Returns:
+        Tuple of (cold_start_result, warm_result)
+    """
+    from rnsr.agent import run_navigator
+    
+    if not questions:
+        raise ValueError("At least one question required")
+    
+    # Cold start benchmark (first query)
+    def run_cold_query():
+        return run_navigator(
+            question=questions[0],
+            skeleton=skeleton,
+            kv_store=kv_store,
+            max_iterations=10,
+        )
+    
+    _, cold_result = _measure_execution(run_cold_query, "Query (Cold Start)")
+    cold_result.metadata = {"question": questions[0][:50]}
+    
+    # Warm benchmark (average of remaining queries)
+    warm_times: list[float] = []
+    warm_memories: list[float] = []
+    
+    for q in questions[1:] if len(questions) > 1 else questions:
+        def run_warm_query():
+            return run_navigator(
+                question=q,
+                skeleton=skeleton,
+                kv_store=kv_store,
+                max_iterations=10,
+            )
+        
+        _, warm_bench = _measure_execution(run_warm_query, "Query (Warm)")
+        warm_times.append(warm_bench.duration_seconds)
+        warm_memories.append(warm_bench.memory_peak_mb)
+    
+    warm_result = BenchmarkResult(
+        name="Query (Warm)",
+        duration_seconds=sum(warm_times) / len(warm_times) if warm_times else 0,
+        memory_peak_mb=max(warm_memories) if warm_memories else 0,
+        memory_current_mb=warm_memories[-1] if warm_memories else 0,
+        throughput=len(warm_times) / sum(warm_times) if sum(warm_times) > 0 else None,
+        throughput_unit="queries/sec",
+        metadata={"query_count": len(warm_times)},
+    )
+    
+    return cold_result, warm_result
+
+
+def run_end_to_end_benchmark(
+    pdf_path: Path | str,
+    questions: list[str],
+) -> PerformanceBenchmark:
+    """
+    Run complete end-to-end benchmark.
+    
+    Args:
+        pdf_path: Path to PDF file
+        questions: List of test questions
+        
+    Returns:
+        PerformanceBenchmark with all results
+    """
+    from rnsr.indexing import build_skeleton_index
+    from rnsr.ingestion import ingest_document
+    
+    pdf_path = Path(pdf_path)
+    total_start = time.perf_counter()
+    
+    benchmark = PerformanceBenchmark()
+    
+    # 1. Ingestion benchmark
+    logger.info("benchmark_ingestion_start", file=pdf_path.name)
+    
+    def do_ingest():
+        return ingest_document(pdf_path)
+    
+    ingestion_result, benchmark.ingestion = _measure_execution(do_ingest, "Ingestion")
+    benchmark.ingestion.metadata = {"file": pdf_path.name}
+    
+    # 2. Indexing benchmark
+    logger.info("benchmark_indexing_start")
+    
+    def do_index():
+        return build_skeleton_index(ingestion_result.tree)
+    
+    (skeleton, kv_store), benchmark.indexing = _measure_execution(do_index, "Indexing")
+    benchmark.indexing.metadata = {"node_count": len(skeleton)}
+    
+    # 3. Query benchmarks
+    if questions:
+        logger.info("benchmark_query_start", question_count=len(questions))
+        benchmark.query_cold, benchmark.query_warm = run_query_benchmark(
+            questions, skeleton, kv_store
+        )
+    
+    benchmark.total_time_seconds = time.perf_counter() - total_start
+    
+    logger.info("benchmark_complete", total_seconds=benchmark.total_time_seconds)
+    
+    return benchmark