social-data-attribution-w2/src/dolma/bin_analysis/core.py

"""Bin population analysis orchestrator for the 24x24 WebOrganizer grid."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path

import polars as pl

from dolma.constants import (
    BIN_EMPTY,
    BIN_FULL,
    BIN_PARTIAL,
    BIN_SPARSE,
    FORMATS,
    TARGET_DOCS_PER_BIN,
    TOPICS,
)

from dolma.bin_analysis.stats import (
    DEFAULT_DOC_THRESHOLD,
    DEFAULT_TOKEN_CAP,
    DEFAULT_TOKEN_FLOOR,
    compute_bin_stats,
    detect_length_mismatches,
    load_manifest,
    validate_labels,
)

logger = logging.getLogger(__name__)


@dataclass
class BinAnalysisResult:
    bin_stats: pl.DataFrame
    mismatches: pl.DataFrame
    recommendation: dict[str, object]
    output_dir: Path


def _raise_on_unknown_labels(
    unknown_topics: list[str],
    unknown_formats: list[str],
) -> None:
    problems: list[str] = []
    if unknown_topics:
        problems.append(f"topics={unknown_topics}")
    if unknown_formats:
        problems.append(f"formats={unknown_formats}")
    if problems:
        raise ValueError(
            "Manifest contains unrecognized WebOrganizer labels: " + ", ".join(problems)
        )


def build_recommendation(
    bin_stats: pl.DataFrame,
    mismatches: pl.DataFrame,
) -> dict[str, object]:
    classification_counts = (
        bin_stats.group_by("classification")
        .agg(pl.len().alias("count"))
        .sort("classification")
    )
    counts_dict = {
        row["classification"]: row["count"]
        for row in classification_counts.iter_rows(named=True)
    }

    total_docs = int(bin_stats["doc_count"].sum())
    total_tokens = int(bin_stats["token_count"].sum())
    non_empty = bin_stats.filter(pl.col("doc_count") > 0)
    n_non_empty = len(non_empty)

    recommendation: dict[str, object] = {
        "classification_counts": counts_dict,
        "total_docs": total_docs,
        "total_tokens": total_tokens,
        "non_empty_bins": n_non_empty,
        "total_bins": len(bin_stats),
        "flagged_bins": len(mismatches),
        "median_docs_per_non_empty_bin": (
            int(non_empty["doc_count"].median()) if n_non_empty > 0 else 0
        ),
        "median_tokens_per_non_empty_bin": (
            int(non_empty["token_count"].median()) if n_non_empty > 0 else 0
        ),
    }

    cls_lines = [
        f"  {c:>8}: {counts_dict.get(c, 0):>4}"
        for c in [BIN_FULL, BIN_PARTIAL, BIN_SPARSE, BIN_EMPTY]
    ]
    lines = [
        f"Total bins: {len(bin_stats)}, non-empty: {n_non_empty}",
        f"Docs: {total_docs:,}, tokens: {total_tokens:,}",
        "Classification:",
        *cls_lines,
    ]
    if len(mismatches) > 0:
        lines.append(f"Flagged bins: {len(mismatches)}")
    recommendation["summary_text"] = "\n".join(lines)
    return recommendation


def run_bin_analysis(
    manifest_path: Path,
    output_dir: Path,
    target_docs_per_bin: int = TARGET_DOCS_PER_BIN,
    doc_threshold: int = DEFAULT_DOC_THRESHOLD,
    token_floor: int = DEFAULT_TOKEN_FLOOR,
    token_cap: int = DEFAULT_TOKEN_CAP,
) -> BinAnalysisResult:
    output_dir.mkdir(parents=True, exist_ok=True)
    lf = load_manifest(manifest_path)

    unknown_topics, unknown_formats = validate_labels(lf, TOPICS, FORMATS)
    _raise_on_unknown_labels(unknown_topics, unknown_formats)

    bin_stats = compute_bin_stats(lf, TOPICS, FORMATS, target_docs_per_bin)
    mismatches = detect_length_mismatches(
        bin_stats, doc_threshold, token_floor, token_cap
    )
    recommendation = build_recommendation(bin_stats, mismatches)
    logger.info("\n%s", recommendation["summary_text"])

    bin_stats.write_parquet(output_dir / "bin_stats.parquet")
    bin_stats.write_csv(output_dir / "bin_stats.csv")
    mismatches.write_csv(output_dir / "length_mismatches.csv")
    (output_dir / "recommendation.json").write_text(
        json.dumps(recommendation, indent=2, default=str), encoding="utf-8"
    )

    return BinAnalysisResult(bin_stats, mismatches, recommendation, output_dir)


__all__ = ["BinAnalysisResult", "run_bin_analysis"]