social-data-attribution-w2/src/data_attribution/analysis/enriched_stats.py

"""Statistics for enriched dataset label distributions.

Computes hard and soft label frequencies, plus entropy statistics,
from an enriched JSONL dataset.
"""

from __future__ import annotations

import argparse
import json
import logging
import math
import statistics
from collections import Counter, defaultdict
from pathlib import Path

logger = logging.getLogger(__name__)


def entropy(probs: list[float]) -> float:
    """Shannon entropy, assumes probs sum to ~1."""
    return -sum(p * math.log(p + 1e-12) for p in probs if p > 0)


def compute_enriched_stats(input_jsonl: Path, output_json: Path) -> None:
    """Compute label statistics from enriched JSONL dataset."""
    hard_counts: Counter[str] = Counter()
    soft_sum: dict[str, float] = defaultdict(float)
    entropies: list[float] = []
    num_docs = 0

    with input_jsonl.open("r", encoding="utf-8") as f:
        for line in f:
            obj = json.loads(line)
            meta = obj.get("metadata", {})
            fmt = meta.get("weborganizer_format")
            fmt_max = meta.get("weborganizer_format_max")

            if not fmt or not fmt_max:
                continue

            num_docs += 1
            hard_counts[fmt_max] += 1

            probs = []
            for label, score in fmt.items():
                soft_sum[label] += float(score)
                probs.append(float(score))

            entropies.append(entropy(probs))

    if num_docs == 0:
        logger.warning("No valid documents found in %s", input_jsonl)
        hard_freq = {}
        soft_mean = {}
    else:
        hard_freq = {k: v / num_docs for k, v in hard_counts.items()}
        soft_mean = {k: v / num_docs for k, v in soft_sum.items()}

    stats = {
        "num_documents": num_docs,
        "hard_label_counts": dict(hard_counts),
        "hard_label_frequencies": hard_freq,
        "soft_label_mean": soft_mean,
        "soft_label_total": dict(soft_sum),
        "entropy": {
            "mean": statistics.mean(entropies) if entropies else None,
            "median": statistics.median(entropies) if entropies else None,
            "p90": (
                statistics.quantiles(entropies, n=10)[8]
                if len(entropies) >= 10
                else None
            ),
        },
    }

    output_json.parent.mkdir(parents=True, exist_ok=True)
    with output_json.open("w", encoding="utf-8") as f:
        json.dump(stats, f, indent=2)

    logger.info("Wrote stats to %s", output_json)
    logger.info("Processed %d documents", num_docs)


def main() -> None:
    """CLI entry point for enriched dataset statistics."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    parser = argparse.ArgumentParser(
        description="Compute label statistics from enriched JSONL dataset."
    )
    parser.add_argument(
        "--input_jsonl",
        type=Path,
        required=True,
        help="Path to enriched dataset JSONL",
    )
    parser.add_argument(
        "--output_json",
        type=Path,
        required=True,
        help="Output path for label statistics JSON",
    )

    args = parser.parse_args()
    compute_enriched_stats(args.input_jsonl, args.output_json)


if __name__ == "__main__":
    main()