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

"""Label statistics and cluster analysis for attribution results.

Computes hard and soft label frequencies, weighted by influence scores,
along with entropy statistics for attributed dataset records.
"""

from __future__ import annotations

import argparse
import csv
import json
import logging
import math
import statistics
import sys
import time
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


PROGRESS_EVERY_ROWS = 200000
PROGRESS_EVERY_LINES = 2_000_000


def safe_entropy(prob_dict: dict[str, float]) -> float:
    h = 0.0
    for p in prob_dict.values():
        if p and p > 0:
            h -= p * math.log(p + 1e-12)
    return h


def normalize_dict(d: dict[str, float]) -> dict[str, float]:
    s = float(sum(d.values()))
    if s <= 0:
        return {k: 0.0 for k in d}
    return {k: float(v) / s for k, v in d.items()}


def weight_transform(w: float, mode: str) -> float:
    if mode == "raw":
        return float(w)
    if mode == "abs":
        return float(abs(w))
    if mode == "pos_only":
        return float(w) if w > 0 else 0.0
    raise ValueError(f"Unknown weight_mode={mode}")


def load_needed_lines_jsonl(
    path: Path, needed_line_idxs: set[int]
) -> dict[int, dict[str, Any]]:
    """Stream-read JSONL and return dict: line_index -> parsed_json."""
    out: dict[int, dict[str, Any]] = {}
    if not needed_line_idxs:
        return out

    remaining = set(needed_line_idxs)
    max_needed = max(needed_line_idxs)

    with path.open("r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            if i > max_needed and not remaining:
                break

            if i in remaining:
                s = line.strip()
                if not s:
                    out[i] = {}
                else:
                    try:
                        out[i] = json.loads(s)
                    except json.JSONDecodeError:
                        out[i] = {"_raw": s}
                remaining.remove(i)

            if i and (i % PROGRESS_EVERY_LINES == 0):
                logger.info(
                    "[jsonl] scanned %d lines; remaining needed=%d", i, len(remaining)
                )

            if not remaining:
                break

    return out


def read_csv_rows_and_needed_indices(
    csv_path: Path,
) -> tuple[list[tuple[int, float]], set[int]]:
    """Read CSV and return rows with needed dataset indices."""
    rows: list[tuple[int, float]] = []
    needed_data_idxs: set[int] = set()

    with csv_path.open("r", encoding="utf-8-sig", newline="") as f:
        reader = csv.DictReader(f)
        for n, r in enumerate(reader, start=1):
            idx = int(r["index_example_idx"])
            score = float(r["attribution"])
            rows.append((idx, score))
            needed_data_idxs.add(idx)

            if n % PROGRESS_EVERY_ROWS == 0:
                logger.info(
                    "[csv] read %d rows; unique indices=%d", n, len(needed_data_idxs)
                )

    return rows, needed_data_idxs


def compute_label_stats(
    csv_path: Path,
    data_jsonl: Path,
    output_json: Path,
    weight_mode: str = "raw",
) -> None:
    """Compute label statistics from attribution CSV and dataset JSONL."""
    t0 = time.time()
    logger.info("compute_label_stats")
    logger.info("csv_path:   %s (exists=%s)", csv_path, csv_path.exists())
    logger.info("data_jsonl: %s (exists=%s)", data_jsonl, data_jsonl.exists())
    logger.info("weight_mode: %s", weight_mode)

    if not csv_path.exists():
        logger.error("missing CSV: %s", csv_path)
        sys.exit(1)
    if not data_jsonl.exists():
        logger.error("missing data_jsonl: %s", data_jsonl)
        sys.exit(1)

    rows, needed_data_idxs = read_csv_rows_and_needed_indices(csv_path)
    logger.info(
        "[csv] total_rows=%d unique_data_indices=%d", len(rows), len(needed_data_idxs)
    )

    data_by_idx = load_needed_lines_jsonl(data_jsonl, needed_data_idxs)
    logger.info("[jsonl] loaded %d needed dataset records", len(data_by_idx))

    hard_counts: Counter[str] = Counter()
    soft_sum: dict[str, float] = defaultdict(float)
    hard_weight_sum: dict[str, float] = defaultdict(float)
    soft_weight_sum: dict[str, float] = defaultdict(float)
    entropies_unweighted: list[float] = []
    entropy_weighted_num = 0.0
    entropy_weighted_den = 0.0
    num_rows = 0
    missing_record = 0
    missing_meta = 0

    for idx, score in rows:
        num_rows += 1
        w = weight_transform(score, weight_mode)

        dr = data_by_idx.get(idx)
        if not isinstance(dr, dict):
            missing_record += 1
            continue

        meta = dr.get("metadata", {}) if isinstance(dr, dict) else {}
        fmt = meta.get("weborganizer_format")
        fmt_max = meta.get("weborganizer_format_max")

        if not fmt or not fmt_max:
            missing_meta += 1
            continue

        hard_counts[fmt_max] += 1
        hard_weight_sum[fmt_max] += w

        for label, p in fmt.items():
            try:
                p = float(p)
            except Exception:
                continue
            soft_sum[label] += p
            soft_weight_sum[label] += w * p

        try:
            ent = safe_entropy({k: float(v) for k, v in fmt.items()})
            entropies_unweighted.append(ent)
            entropy_weighted_num += w * ent
            entropy_weighted_den += w
        except Exception:
            pass

        if num_rows % PROGRESS_EVERY_ROWS == 0:
            logger.info(
                "[accum] rows=%d missing_record=%d missing_meta=%d",
                num_rows,
                missing_record,
                missing_meta,
            )

    denom_items = num_rows - missing_record - missing_meta
    if denom_items <= 0:
        denom_items = 1

    hard_freq = {k: v / denom_items for k, v in hard_counts.items()}
    soft_mean = {k: v / denom_items for k, v in soft_sum.items()}
    hard_weight_freq = normalize_dict(hard_weight_sum)
    hard_weight_total = float(sum(hard_weight_sum.values()))
    soft_weight_mean = {
        k: (v / hard_weight_total) if hard_weight_total > 0 else 0.0
        for k, v in soft_weight_sum.items()
    }

    entropy_stats = {
        "unweighted": {
            "mean": (
                statistics.mean(entropies_unweighted) if entropies_unweighted else None
            ),
            "median": (
                statistics.median(entropies_unweighted)
                if entropies_unweighted
                else None
            ),
            "p90": (
                statistics.quantiles(entropies_unweighted, n=10)[8]
                if len(entropies_unweighted) >= 10
                else None
            ),
        },
        "influence_weighted": {
            "mean": (
                (entropy_weighted_num / entropy_weighted_den)
                if entropy_weighted_den > 0
                else None
            )
        },
    }

    stats = {
        "csv_file": str(csv_path),
        "data_jsonl": str(data_jsonl),
        "weight_mode": weight_mode,
        "num_attribution_rows_total": len(rows),
        "num_rows_missing_dataset_record": missing_record,
        "num_rows_missing_label_metadata": missing_meta,
        "num_rows_with_labels": (len(rows) - missing_record - missing_meta),
        "hard_label_counts_unweighted": dict(hard_counts),
        "hard_label_frequencies_unweighted": hard_freq,
        "hard_label_weight_mass": dict(hard_weight_sum),
        "hard_label_weight_frequencies": hard_weight_freq,
        "soft_label_total_unweighted": dict(soft_sum),
        "soft_label_mean_unweighted": soft_mean,
        "soft_label_total_influence_weighted": dict(soft_weight_sum),
        "soft_label_mean_influence_weighted": soft_weight_mean,
        "entropy": entropy_stats,
    }

    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)

    dt = time.time() - t0
    logger.info("Wrote: %s (%d bytes)", output_json, output_json.stat().st_size)
    logger.info("Done in %.2fs", dt)


def main() -> None:
    """CLI entry point for cluster label statistics."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    parser = argparse.ArgumentParser(
        description="Compute label statistics from attribution results."
    )
    parser.add_argument(
        "--csv_path",
        type=Path,
        required=True,
        help="Path to attribution CSV (topk_attributions.csv or attributions.csv)",
    )
    parser.add_argument(
        "--data_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",
    )
    parser.add_argument(
        "--weight_mode",
        type=str,
        default="raw",
        choices=["raw", "abs", "pos_only"],
        help="How to transform attribution weights",
    )

    args = parser.parse_args()
    compute_label_stats(
        args.csv_path, args.data_jsonl, args.output_json, args.weight_mode
    )


if __name__ == "__main__":
    main()