social-data-attribution-w2/scripts/diagnostics/truncated_score_distribution.py

"""Phase 3: Quality score distribution analysis by format.

Samples shards from R2, joins quality and format sidecars, and computes
per-format quality score histograms and summary statistics.

Outputs CSV with per-format statistics and cross-tabulation of
format_confidence vs quality_score for truncated docs.

Requires R2 credentials via with_r2_credentials.sh.
"""

from __future__ import annotations

import argparse
import csv
import json
import logging
import os
import sys
from collections import defaultdict
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger(__name__)

R2_BUCKET = "soc127-dedup"
R2_ENDPOINT_URL = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"
R2_INPUT_PREFIXES = [
    "soc127/phase1_pool_shared",
    "soc127/phase2_nonpool_final",
]
SOC91_PREFIX = "soc91-labels"
SOC139_PREFIX = "soc139-quality-sidecars"

TARGET_FORMATS = ["truncated", "spam_ads", "academic_writing", "news_article"]
HISTOGRAM_BINS = 100
CONFIDENCE_BUCKETS = [0.0, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
QUALITY_BUCKETS = [0.0, 0.05, 0.1, 0.2, 0.3, 0.5, 1.0]


def bucket_index(value: float, edges: list[float]) -> str:
    for i in range(len(edges) - 1):
        if value < edges[i + 1]:
            return f"[{edges[i]:.2f},{edges[i + 1]:.2f})"
    return f"[{edges[-2]:.2f},{edges[-1]:.2f}]"


def find_shards_with_both_sidecars(
    client, *, bucket: str, shard_count: int
) -> list[str]:
    from dolma.quality.r2 import list_keys

    soc91_parquets = set(
        list_keys(client, bucket=bucket, prefix=SOC91_PREFIX, suffix=".parquet")
    )
    soc139_parquets = set(
        list_keys(client, bucket=bucket, prefix=SOC139_PREFIX, suffix=".parquet")
    )

    soc91_basenames = {Path(k).stem for k in soc91_parquets}
    soc139_basenames = {Path(k).stem for k in soc139_parquets}
    common = soc91_basenames & soc139_basenames

    source_keys: list[str] = []
    for prefix in R2_INPUT_PREFIXES:
        source_keys.extend(
            list_keys(client, bucket=bucket, prefix=prefix, suffix=".jsonl.zst")
        )

    matched: list[str] = []
    for key in sorted(source_keys):
        basename = Path(key).name.removesuffix(".jsonl.zst")
        if basename in common:
            matched.append(key)
            if len(matched) >= shard_count:
                break

    log.info(
        "Found %d shards with both sidecars (requested %d)", len(matched), shard_count
    )
    return matched


def process_shard(client, *, bucket: str, source_key: str) -> list[dict]:
    from dolma.quality.validation.io import read_quality_rows, read_soc91_doc_map

    quality_rows = read_quality_rows(
        client,
        bucket=bucket,
        source_key=source_key,
        output_prefix=SOC139_PREFIX,
    )
    soc91_map = read_soc91_doc_map(
        client,
        bucket=bucket,
        source_key=source_key,
        soc91_prefix=SOC91_PREFIX,
    )
    if soc91_map is None:
        return []

    joined: list[dict] = []
    for row in quality_rows:
        doc_id = str(row["doc_id"])
        soc91_doc = soc91_map.get(doc_id)
        if soc91_doc is None:
            continue
        format_label = soc91_doc.get("format_url_label")
        if not isinstance(format_label, str):
            continue
        joined.append(
            {
                "format": format_label,
                "quality_score": float(row["quality_score"]),
                "quality_label": "high"
                if float(row["quality_high_prob"]) >= float(row["quality_low_prob"])
                else "low",
            }
        )
    return joined


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Phase 3: Quality score distributions by format"
    )
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--shard-count", type=int, default=50)
    args = parser.parse_args()

    from dolma.quality.r2 import R2Config, create_r2_client
    from dolma.quality.validation.stats import GroupSummary

    config = R2Config(
        endpoint_url=R2_ENDPOINT_URL,
        bucket=R2_BUCKET,
        access_key_id=os.environ["R2_ACCESS_KEY_ID"],
        secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"],
        input_prefixes=tuple(R2_INPUT_PREFIXES),
        output_prefix=SOC139_PREFIX,
    )
    client = create_r2_client(config)

    shards = find_shards_with_both_sidecars(
        client,
        bucket=R2_BUCKET,
        shard_count=args.shard_count,
    )
    if not shards:
        log.error("No shards found with both sidecars.")
        sys.exit(1)

    format_groups: dict[str, GroupSummary] = defaultdict(GroupSummary)
    format_histograms: dict[str, list[int]] = defaultdict(lambda: [0] * HISTOGRAM_BINS)

    for i, source_key in enumerate(shards, 1):
        log.info("[%d/%d] %s", i, len(shards), source_key)
        joined = process_shard(client, bucket=R2_BUCKET, source_key=source_key)

        for doc in joined:
            fmt = doc["format"]
            score = doc["quality_score"]
            label = doc["quality_label"]
            format_groups[fmt].update(score, label)

            bin_idx = min(int(score * HISTOGRAM_BINS), HISTOGRAM_BINS - 1)
            format_histograms[fmt][bin_idx] += 1

        log.info("  %d joined docs", len(joined))

    args.output_dir.mkdir(parents=True, exist_ok=True)

    summary_rows: list[dict] = []
    for fmt in sorted(format_groups.keys()):
        group = format_groups[fmt]
        summary_rows.append(group.row("format", fmt))

    summary_path = args.output_dir / "format_quality_summary.csv"
    if summary_rows:
        with open(summary_path, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=list(summary_rows[0].keys()))
            writer.writeheader()
            writer.writerows(summary_rows)
        log.info("Wrote format summary to %s", summary_path)

    hist_path = args.output_dir / "format_quality_histograms.json"
    with open(hist_path, "w") as f:
        json.dump(dict(format_histograms), f, indent=2)
    log.info("Wrote histograms to %s", hist_path)

    log.info("\n=== Format quality summary (target formats) ===\n")
    log.info(
        "%-20s %8s %8s %8s %8s %8s %8s %8s",
        "format",
        "count",
        "mean",
        "p5",
        "p25",
        "p50",
        "p75",
        "p95",
    )
    for row in summary_rows:
        fmt = row["format"]
        if fmt not in TARGET_FORMATS and len(summary_rows) > 10:
            continue
        stats = row
        log.info(
            "%-20s %8d %8.4f %8.4f %8.4f %8.4f %8.4f %8.4f",
            fmt,
            stats.get("count", 0),
            stats.get("mean", 0) or 0,
            stats.get("p5", 0) or 0,
            stats.get("p25", 0) or 0,
            stats.get("p50", 0) or 0,
            stats.get("p75", 0) or 0,
            stats.get("p95", 0) or 0,
        )

    trunc_group = format_groups.get("truncated")
    if trunc_group and trunc_group.scores.count > 0:
        trunc_summary = trunc_group.scores.summary()
        log.info("\n=== Truncated format distribution shape ===")
        log.info("Count: %d", trunc_summary["count"])
        log.info("Mean:  %.4f", trunc_summary["mean"] or 0)
        log.info("Std:   %.4f", trunc_summary["std"] or 0)
        log.info("P5:    %.4f", trunc_summary.get("p5", 0) or 0)
        log.info("P25:   %.4f", trunc_summary.get("p25", 0) or 0)
        log.info("P50:   %.4f", trunc_summary.get("p50", 0) or 0)
        log.info("P75:   %.4f", trunc_summary.get("p75", 0) or 0)
        log.info("P95:   %.4f", trunc_summary.get("p95", 0) or 0)

        hist = format_histograms.get("truncated", [])
        total = sum(hist)
        if total > 0:
            low_mass = sum(hist[:5]) / total
            mid_mass = sum(hist[5:50]) / total
            high_mass = sum(hist[50:]) / total
            log.info(
                "\nMass distribution: [0,0.05)=%.1f%%  [0.05,0.50)=%.1f%%  [0.50,1.0]=%.1f%%",
                low_mass * 100,
                mid_mass * 100,
                high_mass * 100,
            )
            if mid_mass > 0.3:
                log.info(
                    "FINDING: Significant mass in mid-range - smooth rightward shift (supports B)"
                )
            elif high_mass > 0.1 and low_mass > 0.5:
                log.info(
                    "FINDING: Bimodal distribution - mixed causes, investigate subpopulation"
                )


if __name__ == "__main__":
    main()