mirage-benchmark 1.0.4__py3-none-any.whl

mirage/utils/stats.py

@@ -0,0 +1,626 @@
+#!/usr/bin/env python3
+"""
+Dataset Statistics Calculator
+Measures pages, images, and tokens for each dataset using chunks.json files.
+"""
+
+import os
+import re
+import zipfile
+import io
+import json
+from pathlib import Path
+from collections import defaultdict
+from typing import Dict, List, Any, Optional
+
+import pypdfium2 as pdfium
+from tqdm import tqdm
+
+
+def get_tokenizer():
+    """Get tokenizer - uses GPT2 if available, otherwise word-based estimate."""
+    try:
+        from transformers import GPT2TokenizerFast
+        return GPT2TokenizerFast.from_pretrained("gpt2")
+    except Exception:
+        return None
+
+
+def count_tokens(text: str, tokenizer=None) -> int:
+    """Count tokens - uses tokenizer if available, otherwise word-based estimate."""
+    if not text:
+        return 0
+    
+    if tokenizer:
+        try:
+            return len(tokenizer.encode(text))
+        except Exception:
+            pass
+    
+    # Fallback: approximate tokens as ~1.3 * words (GPT average)
+    words = len(text.split())
+    return int(words * 1.3)
+
+
+def count_pages_from_zip(zip_path: Path) -> Dict[str, int]:
+    """Count pages per document from zip file (PDFs or HTML files)."""
+    page_counts = {}
+    try:
+        with zipfile.ZipFile(zip_path, 'r') as zf:
+            for name in zf.namelist():
+                stem = Path(name).stem
+                lower_name = name.lower()
+                
+                if lower_name.endswith('.pdf'):
+                    try:
+                        pdf_bytes = zf.read(name)
+                        pdf = pdfium.PdfDocument(pdf_bytes)
+                        page_counts[stem] = len(pdf)
+                    except Exception:
+                        page_counts[stem] = 0
+                elif lower_name.endswith(('.html', '.htm')):
+                    # HTML files count as 1 "page"
+                    page_counts[stem] = 1
+    except Exception as e:
+        print(f"Error reading zip {zip_path}: {e}")
+    return page_counts
+
+
+def get_file_size_mb_from_zip(zip_path: Path) -> Dict[str, float]:
+    """Get file sizes per document from zip file."""
+    sizes = {}
+    try:
+        with zipfile.ZipFile(zip_path, 'r') as zf:
+            for info in zf.infolist():
+                lower = info.filename.lower()
+                if lower.endswith(('.pdf', '.html', '.htm')):
+                    stem = Path(info.filename).stem
+                    sizes[stem] = info.file_size / (1024 * 1024)
+    except Exception:
+        pass
+    return sizes
+
+
+def get_file_type_from_zip(zip_path: Path) -> str:
+    """Detect primary file type in zip."""
+    try:
+        with zipfile.ZipFile(zip_path, 'r') as zf:
+            names = zf.namelist()
+            pdfs = sum(1 for n in names if n.lower().endswith('.pdf'))
+            htmls = sum(1 for n in names if n.lower().endswith(('.html', '.htm')))
+            if pdfs > htmls:
+                return "pdf"
+            elif htmls > 0:
+                return "html"
+    except Exception:
+        pass
+    return "unknown"
+
+
+def analyze_chunks_json(chunks_path: Path, tokenizer=None) -> Dict[str, Any]:
+    """Analyze a chunks.json file to extract stats per PDF."""
+    pdf_stats = defaultdict(lambda: {
+        "chunks": 0,
+        "text_chunks": 0,
+        "image_chunks": 0,
+        "table_chunks": 0,
+        "tokens": 0,
+        "chars": 0
+    })
+    
+    try:
+        with open(chunks_path, 'r') as f:
+            chunks = json.load(f)
+        
+        for chunk in chunks:
+            file_name = chunk.get("file_name", "unknown")
+            chunk_type = chunk.get("chunk_type", "").lower()
+            content = chunk.get("content", "")
+            artifact = chunk.get("artifact", "None")
+            
+            stats = pdf_stats[file_name]
+            stats["chunks"] += 1
+            stats["chars"] += len(content)
+            stats["tokens"] += count_tokens(content, tokenizer)
+            
+            # Classify chunk type
+            if "image" in chunk_type or (artifact and artifact != "None"):
+                stats["image_chunks"] += 1
+            elif "table" in chunk_type:
+                stats["table_chunks"] += 1
+            else:
+                stats["text_chunks"] += 1
+                
+    except Exception as e:
+        print(f"Error reading {chunks_path}: {e}")
+    
+    return dict(pdf_stats)
+
+
+def analyze_dataset(dataset_name: str, chunks_path: Optional[Path], zip_path: Optional[Path], tokenizer=None) -> Dict[str, Any]:
+    """Analyze a single dataset using chunks.json and/or zip file."""
+    stats = {
+        "dataset_name": dataset_name,
+        "source": "chunks.json" if chunks_path else "zip",
+        "num_pdfs": 0,
+        "total_pages": 0,
+        "total_chunks": 0,
+        "total_text_chunks": 0,
+        "total_image_chunks": 0,
+        "total_table_chunks": 0,
+        "total_tokens": 0,
+        "total_chars": 0,
+        "total_size_mb": 0.0,
+        "pdfs": []
+    }
+    
+    # Get page counts and file sizes from zip
+    page_counts = {}
+    file_sizes = {}
+    file_type = "unknown"
+    if zip_path and zip_path.exists():
+        page_counts = count_pages_from_zip(zip_path)
+        file_sizes = get_file_size_mb_from_zip(zip_path)
+        file_type = get_file_type_from_zip(zip_path)
+        stats["total_size_mb"] = sum(file_sizes.values())
+        stats["file_type"] = file_type
+    
+    # Get chunk stats from chunks.json
+    chunk_stats = {}
+    if chunks_path and chunks_path.exists():
+        chunk_stats = analyze_chunks_json(chunks_path, tokenizer)
+    
+    # Merge stats - use chunk_stats keys as primary if available
+    all_pdfs = set(chunk_stats.keys()) | set(page_counts.keys())
+    stats["num_pdfs"] = len(all_pdfs)
+    
+    for pdf_name in sorted(all_pdfs):
+        pdf_info = {
+            "filename": pdf_name,
+            "pages": page_counts.get(pdf_name, 0),
+            "size_mb": round(file_sizes.get(pdf_name, 0), 2),
+            "chunks": chunk_stats.get(pdf_name, {}).get("chunks", 0),
+            "text_chunks": chunk_stats.get(pdf_name, {}).get("text_chunks", 0),
+            "image_chunks": chunk_stats.get(pdf_name, {}).get("image_chunks", 0),
+            "table_chunks": chunk_stats.get(pdf_name, {}).get("table_chunks", 0),
+            "tokens": chunk_stats.get(pdf_name, {}).get("tokens", 0),
+            "chars": chunk_stats.get(pdf_name, {}).get("chars", 0),
+        }
+        stats["pdfs"].append(pdf_info)
+        
+        stats["total_pages"] += pdf_info["pages"]
+        stats["total_chunks"] += pdf_info["chunks"]
+        stats["total_text_chunks"] += pdf_info["text_chunks"]
+        stats["total_image_chunks"] += pdf_info["image_chunks"]
+        stats["total_table_chunks"] += pdf_info["table_chunks"]
+        stats["total_tokens"] += pdf_info["tokens"]
+        stats["total_chars"] += pdf_info["chars"]
+    
+    stats["total_size_mb"] = round(stats["total_size_mb"], 2)
+    
+    # Compute averages
+    if stats["num_pdfs"] > 0:
+        stats["avg_pages_per_pdf"] = round(stats["total_pages"] / stats["num_pdfs"], 1)
+        stats["avg_chunks_per_pdf"] = round(stats["total_chunks"] / stats["num_pdfs"], 1)
+        stats["avg_images_per_pdf"] = round(stats["total_image_chunks"] / stats["num_pdfs"], 1)
+        stats["avg_tokens_per_pdf"] = int(stats["total_tokens"] / stats["num_pdfs"])
+    if stats["total_pages"] > 0:
+        stats["avg_tokens_per_page"] = int(stats["total_tokens"] / stats["total_pages"])
+    
+    return stats
+
+
+def print_summary(all_stats: List[Dict]) -> None:
+    """Print formatted summary of all datasets."""
+    print("\n" + "=" * 110)
+    print("DATASET STATISTICS SUMMARY (from chunks.json)")
+    print("=" * 110)
+    
+    print(f"\n{'Dataset':<28} {'Docs':>5} {'Pages':>6} {'Chunks':>7} {'Images':>7} {'Tables':>7} {'Tokens':>12} {'Size(MB)':>10}")
+    print("-" * 110)
+    
+    totals = defaultdict(int)
+    totals["total_size_mb"] = 0.0
+    
+    for stats in all_stats:
+        print(f"{stats['dataset_name']:<28} {stats['num_pdfs']:>5} {stats['total_pages']:>6} "
+              f"{stats['total_chunks']:>7} {stats['total_image_chunks']:>7} {stats['total_table_chunks']:>7} "
+              f"{stats['total_tokens']:>12,} {stats['total_size_mb']:>10.1f}")
+        
+        totals["num_pdfs"] += stats["num_pdfs"]
+        totals["total_pages"] += stats["total_pages"]
+        totals["total_chunks"] += stats["total_chunks"]
+        totals["total_image_chunks"] += stats["total_image_chunks"]
+        totals["total_table_chunks"] += stats["total_table_chunks"]
+        totals["total_tokens"] += stats["total_tokens"]
+        totals["total_size_mb"] += stats["total_size_mb"]
+    
+    print("-" * 110)
+    print(f"{'TOTAL':<28} {totals['num_pdfs']:>5} {totals['total_pages']:>6} "
+          f"{totals['total_chunks']:>7} {totals['total_image_chunks']:>7} {totals['total_table_chunks']:>7} "
+          f"{totals['total_tokens']:>12,} {totals['total_size_mb']:>10.1f}")
+    print("=" * 110)
+    
+    # Print averages
+    print("\nAVERAGES PER DATASET:")
+    print(f"{'Dataset':<28} {'Pg/Doc':>8} {'Chunks/Doc':>10} {'Img/Doc':>8} {'Tok/Doc':>12} {'Tok/Page':>10}")
+    print("-" * 80)
+    for stats in all_stats:
+        print(f"{stats['dataset_name']:<28} {stats.get('avg_pages_per_pdf', 0):>8.1f} "
+              f"{stats.get('avg_chunks_per_pdf', 0):>10.1f} {stats.get('avg_images_per_pdf', 0):>8.1f} "
+              f"{stats.get('avg_tokens_per_pdf', 0):>12,} {stats.get('avg_tokens_per_page', 0):>10,}")
+
+
+def find_datasets(data_dir: Path, results_dir: Path) -> List[Dict]:
+    """Find all datasets and their corresponding chunks.json files."""
+    datasets = []
+    
+    # Get all zip files in data directory
+    zip_files = {z.stem: z for z in data_dir.glob("*.zip")}
+    
+    # Get all chunks.json files in results directory
+    chunks_files = {}
+    if results_dir.exists():
+        for chunks_path in results_dir.glob("*/chunks.json"):
+            dataset_name = chunks_path.parent.name
+            chunks_files[dataset_name] = chunks_path
+    
+    # Merge: prefer chunks.json when available
+    all_datasets = set(zip_files.keys()) | set(chunks_files.keys())
+    
+    for name in sorted(all_datasets):
+        datasets.append({
+            "name": name,
+            "zip_path": zip_files.get(name),
+            "chunks_path": chunks_files.get(name)
+        })
+    
+    return datasets
+
+
+def compute_dataset_stats(
+    output_dir: str,
+    pdf_dir: str = None,
+    chunks_file: str = None,
+    tokenizer=None
+) -> Dict[str, Any]:
+    """
+    Compute comprehensive dataset statistics from a trial results directory.
+    
+    Args:
+        output_dir: Path to the output directory (e.g., output/results/my_dataset)
+        pdf_dir: Path to source PDFs directory (for page counts)
+        chunks_file: Path to chunks.json (defaults to output_dir/chunks.json)
+        tokenizer: Optional tokenizer for accurate token counting
+    
+    Returns:
+        Dict with keys: total_images, total_tables, total_pages, total_tokens,
+                       num_pdfs, per_pdf_stats, etc.
+    """
+    output_path = Path(output_dir)
+    
+    # Default chunks file path
+    if chunks_file is None:
+        chunks_file = output_path / "chunks.json"
+    else:
+        chunks_file = Path(chunks_file)
+    
+    stats = {
+        "num_pdfs": 0,
+        "total_images": 0,
+        "total_tables": 0,
+        "total_pages": 0,
+        "total_tokens": 0,
+        "total_chars": 0,
+        "total_chunks": 0,
+        "per_pdf_stats": []
+    }
+    
+    # Get PDF names from chunks.json
+    pdf_names = set()
+    chunks_data = []
+    if chunks_file.exists():
+        try:
+            with open(chunks_file, 'r') as f:
+                chunks_data = json.load(f)
+            for chunk in chunks_data:
+                file_name = chunk.get("file_name", "")
+                if file_name:
+                    pdf_names.add(file_name)
+            stats["total_chunks"] = len(chunks_data)
+        except Exception as e:
+            print(f"Warning: Could not read chunks.json: {e}")
+    
+    stats["num_pdfs"] = len(pdf_names)
+    
+    # Count images and tables from markdown directory structure
+    markdown_dir = output_path / "markdown"
+    all_image_files = set()
+    all_table_files = set()
+    per_pdf_images = defaultdict(set)
+    per_pdf_tables = defaultdict(set)
+    
+    if markdown_dir.exists():
+        for subdir in markdown_dir.iterdir():
+            if subdir.is_dir():
+                pdf_name = subdir.name
+                
+                # Count images from ref_artifacts/
+                artifact_dir = subdir / "ref_artifacts"
+                if artifact_dir.exists():
+                    for img_file in artifact_dir.glob("image_*"):
+                        if img_file.suffix.lower() in ['.png', '.jpg', '.jpeg', '.gif', '.webp']:
+                            all_image_files.add(img_file.name)
+                            per_pdf_images[pdf_name].add(img_file.name)
+                
+                # Count tables from tables/
+                table_dir = subdir / "tables"
+                if table_dir.exists():
+                    for table_file in table_dir.glob("*"):
+                        if table_file.suffix.lower() in ['.png', '.jpg', '.jpeg']:
+                            all_table_files.add(table_file.name)
+                            per_pdf_tables[pdf_name].add(table_file.name)
+    
+    stats["total_images"] = len(all_image_files)
+    stats["total_tables"] = len(all_table_files)
+    
+    # Get page counts from source PDFs
+    page_counts = {}
+    if pdf_dir:
+        pdf_path = Path(pdf_dir)
+        if pdf_path.exists():
+            for pdf_file in pdf_path.glob("*.pdf"):
+                pdf_stem = pdf_file.stem
+                # Match PDF names to chunk file names
+                for chunk_name in pdf_names:
+                    if chunk_name.lower() in pdf_stem.lower() or pdf_stem.lower() in chunk_name.lower():
+                        try:
+                            pdf = pdfium.PdfDocument(pdf_file)
+                            page_counts[chunk_name] = len(pdf)
+                        except Exception:
+                            page_counts[chunk_name] = 0
+                        break
+    
+    stats["total_pages"] = sum(page_counts.values())
+    
+    # Count tokens from chunks
+    total_tokens = 0
+    total_chars = 0
+    per_pdf_tokens = defaultdict(int)
+    per_pdf_chars = defaultdict(int)
+    
+    for chunk in chunks_data:
+        content = chunk.get("content", "")
+        file_name = chunk.get("file_name", "unknown")
+        
+        chars = len(content)
+        tokens = count_tokens(content, tokenizer)
+        
+        total_chars += chars
+        total_tokens += tokens
+        per_pdf_chars[file_name] += chars
+        per_pdf_tokens[file_name] += tokens
+    
+    stats["total_tokens"] = total_tokens
+    stats["total_chars"] = total_chars
+    
+    # Build per-PDF stats
+    for pdf_name in sorted(pdf_names):
+        pdf_stats = {
+            "filename": pdf_name,
+            "pages": page_counts.get(pdf_name, 0),
+            "images": len(per_pdf_images.get(pdf_name, set())),
+            "tables": len(per_pdf_tables.get(pdf_name, set())),
+            "tokens": per_pdf_tokens.get(pdf_name, 0),
+            "chars": per_pdf_chars.get(pdf_name, 0)
+        }
+        stats["per_pdf_stats"].append(pdf_stats)
+    
+    # Compute averages
+    if stats["num_pdfs"] > 0:
+        stats["avg_pages_per_pdf"] = round(stats["total_pages"] / stats["num_pdfs"], 1)
+        stats["avg_images_per_pdf"] = round(stats["total_images"] / stats["num_pdfs"], 1)
+        stats["avg_tables_per_pdf"] = round(stats["total_tables"] / stats["num_pdfs"], 1)
+        stats["avg_tokens_per_pdf"] = int(stats["total_tokens"] / stats["num_pdfs"])
+    
+    if stats["total_pages"] > 0:
+        stats["avg_tokens_per_page"] = int(stats["total_tokens"] / stats["total_pages"])
+    
+    return stats
+
+
+def print_dataset_stats(stats: Dict[str, Any]) -> None:
+    """Print formatted dataset statistics."""
+    print("\n" + "=" * 60)
+    print("DATASET STATISTICS")
+    print("=" * 60)
+    print(f"\nImages: {stats['total_images']}")
+    print(f"Tables: {stats['total_tables']}")
+    print(f"Pages: {stats['total_pages']}")
+    print(f"Tokens: {stats['total_tokens']:,}")
+    print(f"PDFs: {stats['num_pdfs']}")
+    print(f"Chunks: {stats['total_chunks']}")
+    print("=" * 60)
+
+
+def compute_qa_category_stats(qa_data: List[Dict]) -> Dict[str, Any]:
+    """
+    Compute QA category statistics including multihop/multimodal intersection.
+    
+    Args:
+        qa_data: List of QA pairs from deduplicated dataset
+    
+    Returns:
+        Dict with category counts and percentages
+    """
+    total = len(qa_data)
+    if total == 0:
+        return {
+            'total_qa_pairs': 0,
+            'multihop_count': 0,
+            'multimodal_count': 0,
+            'multihop_multimodal_count': 0,
+            'multihop_only_count': 0,
+            'multimodal_only_count': 0,
+            'text_only_count': 0,
+            'avg_difficulty': 0.0,
+            'avg_relevance': 0.0
+        }
+    
+    multihop = 0
+    multimodal = 0
+    both = 0
+    
+    # Track difficulty and relevance scores
+    difficulty_scores = []
+    relevance_scores = []
+    
+    for qa in qa_data:
+        # Multihop: hop_count > 1 or multiple chunks added
+        hop_count = qa.get('hop_count', 0)
+        chunks_added = qa.get('chunks_added', [])
+        is_multihop = hop_count > 1 or (isinstance(chunks_added, list) and len(chunks_added) > 1)
+        
+        # Multimodal: has image_path in context chunks OR markdown images in content
+        # Match logic from metrics_optimized.py for consistency
+        context_chunks = qa.get('context_chunks', [])
+        is_multimodal = False
+        for chunk in context_chunks:
+            if isinstance(chunk, dict):
+                # Check image_path field
+                image_path = chunk.get('image_path')
+                if image_path and image_path not in ('None', 'null', None) and str(image_path).strip():
+                    is_multimodal = True
+                    break
+                # Check content for markdown image references (e.g., ![alt](path))
+                content = chunk.get('content', '')
+                if content and re.search(r'!\[[^\]]*\]\([^)]+\)', content):
+                    is_multimodal = True
+                    break
+        
+        if is_multihop:
+            multihop += 1
+        if is_multimodal:
+            multimodal += 1
+        if is_multihop and is_multimodal:
+            both += 1
+        
+        # Extract difficulty and relevance scores
+        try:
+            diff = qa.get('difficulty_score', qa.get('difficulty', 0))
+            if diff is not None and str(diff).strip():
+                difficulty_scores.append(float(diff))
+        except (ValueError, TypeError):
+            pass
+        try:
+            rel = qa.get('relevance_score', qa.get('relevance', 0))
+            if rel is not None and str(rel).strip():
+                relevance_scores.append(float(rel))
+        except (ValueError, TypeError):
+            pass
+    
+    # Exclusive counts
+    multihop_only = multihop - both
+    multimodal_only = multimodal - both
+    text_only = total - multihop - multimodal + both
+    
+    # Compute averages
+    avg_difficulty = round(sum(difficulty_scores) / len(difficulty_scores), 2) if difficulty_scores else 0.0
+    avg_relevance = round(sum(relevance_scores) / len(relevance_scores), 2) if relevance_scores else 0.0
+    
+    return {
+        'total_qa_pairs': total,
+        'multihop_count': multihop,
+        'multimodal_count': multimodal,
+        'multihop_multimodal_count': both,
+        'multihop_only_count': multihop_only,
+        'multimodal_only_count': multimodal_only,
+        'text_only_count': text_only,
+        'multihop_pct': round(100 * multihop / total, 1),
+        'multimodal_pct': round(100 * multimodal / total, 1),
+        'multihop_multimodal_pct': round(100 * both / total, 1),
+        'multihop_only_pct': round(100 * multihop_only / total, 1),
+        'multimodal_only_pct': round(100 * multimodal_only / total, 1),
+        'text_only_pct': round(100 * text_only / total, 1),
+        'avg_difficulty': avg_difficulty,
+        'avg_relevance': avg_relevance
+    }
+
+
+def print_qa_category_stats(stats: Dict[str, Any]) -> None:
+    """Print formatted QA category statistics."""
+    total = stats.get('total_qa_pairs', 0)
+    if total == 0:
+        print("\n⚠️ No QA pairs to analyze")
+        return
+    
+    print("\n" + "=" * 60)
+    print("QA CATEGORY BREAKDOWN")
+    print("=" * 60)
+    print(f"\nTotal QA Pairs: {total}")
+    print()
+    print("Category Counts:")
+    print(f"  Multihop:                      {stats['multihop_count']:>3} ({stats['multihop_pct']:>5.1f}%)")
+    print(f"  Multimodal:                    {stats['multimodal_count']:>3} ({stats['multimodal_pct']:>5.1f}%)")
+    print(f"  Both (Multihop ∩ Multimodal):  {stats['multihop_multimodal_count']:>3} ({stats['multihop_multimodal_pct']:>5.1f}%)")
+    print()
+    print("Exclusive Breakdown:")
+    print(f"  Multihop only (text):          {stats['multihop_only_count']:>3} ({stats['multihop_only_pct']:>5.1f}%)")
+    print(f"  Multimodal only (single-hop):  {stats['multimodal_only_count']:>3} ({stats['multimodal_only_pct']:>5.1f}%)")
+    print(f"  Both (multihop + multimodal):  {stats['multihop_multimodal_count']:>3} ({stats['multihop_multimodal_pct']:>5.1f}%)")
+    print(f"  Neither (single-hop, text):   {stats['text_only_count']:>3} ({stats['text_only_pct']:>5.1f}%)")
+    print()
+    print("Quality Scores (0-10 scale):")
+    print(f"  Avg Difficulty:                {stats.get('avg_difficulty', 0.0):>5.2f}")
+    print(f"  Avg Relevance:                 {stats.get('avg_relevance', 0.0):>5.2f}")
+    print("=" * 60)
+
+
+def main():
+    """Main entry point."""
+    data_dir = Path(__file__).parent
+    project_root = data_dir.parent
+    results_dir = project_root / "trials" / "results"
+    
+    datasets = find_datasets(data_dir, results_dir)
+    
+    if not datasets:
+        print("No datasets found.")
+        return
+    
+    print(f"Found {len(datasets)} datasets to analyze...")
+    
+    # Show which have chunks.json
+    with_chunks = sum(1 for d in datasets if d["chunks_path"])
+    print(f"  - {with_chunks} with chunks.json (processed)")
+    print(f"  - {len(datasets) - with_chunks} zip-only (not yet processed)")
+    
+    # Load tokenizer once
+    tokenizer = get_tokenizer()
+    if tokenizer:
+        print("Using GPT2 tokenizer for token counting")
+    else:
+        print("Using word-based token estimation")
+    
+    all_stats = []
+    for ds in tqdm(datasets, desc="Analyzing datasets"):
+        stats = analyze_dataset(
+            ds["name"],
+            ds["chunks_path"],
+            ds["zip_path"],
+            tokenizer
+        )
+        all_stats.append(stats)
+    
+    # Print summary
+    print_summary(all_stats)
+    
+    # Save detailed results
+    output_file = data_dir / "dataset_stats.json"
+    with open(output_file, 'w') as f:
+        json.dump(all_stats, f, indent=2)
+    print(f"\nDetailed stats saved to: {output_file}")
+
+
+if __name__ == "__main__":
+    main()