social-data-attribution-w2/src/dolma/quality/validation/stats.py

"""Streaming score statistics for quality validation."""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Any

SCORE_HISTOGRAM_BINS = 1000


@dataclass
class ScoreAccumulator:
    count: int = 0
    total: float = 0.0
    total_sq: float = 0.0
    minimum: float | None = None
    maximum: float | None = None
    histogram: list[int] = field(
        default_factory=lambda: [0 for _ in range(SCORE_HISTOGRAM_BINS)]
    )

    def update(self, value: float | None) -> None:
        if value is None:
            return
        clipped = min(max(float(value), 0.0), 1.0)
        self.count += 1
        self.total += clipped
        self.total_sq += clipped * clipped
        self.minimum = clipped if self.minimum is None else min(self.minimum, clipped)
        self.maximum = clipped if self.maximum is None else max(self.maximum, clipped)
        index = min(int(clipped * SCORE_HISTOGRAM_BINS), SCORE_HISTOGRAM_BINS - 1)
        self.histogram[index] += 1

    def merge(self, other: "ScoreAccumulator") -> None:
        self.count += other.count
        self.total += other.total
        self.total_sq += other.total_sq
        if other.minimum is not None:
            self.minimum = (
                other.minimum
                if self.minimum is None
                else min(self.minimum, other.minimum)
            )
        if other.maximum is not None:
            self.maximum = (
                other.maximum
                if self.maximum is None
                else max(self.maximum, other.maximum)
            )
        self.histogram = [
            left + right
            for left, right in zip(self.histogram, other.histogram, strict=True)
        ]

    def quantile(self, value: float) -> float | None:
        if self.count == 0:
            return None
        target = value * self.count
        running = 0
        for index, count in enumerate(self.histogram):
            if count == 0:
                continue
            next_running = running + count
            if next_running >= target:
                left = index / SCORE_HISTOGRAM_BINS
                step = 1 / SCORE_HISTOGRAM_BINS
                offset = (target - running) / count if count else 0.0
                return left + max(offset, 0.0) * step
            running = next_running
        return 1.0

    def summary(self) -> dict[str, float | int | None]:
        if self.count == 0:
            return {
                "count": 0,
                "mean": None,
                "std": None,
                "min": None,
                "p25": None,
                "p50": None,
                "p75": None,
                "p95": None,
                "max": None,
            }
        mean = self.total / self.count
        variance = max((self.total_sq / self.count) - (mean * mean), 0.0)
        return {
            "count": self.count,
            "mean": mean,
            "std": math.sqrt(variance),
            "min": self.minimum,
            "p25": self.quantile(0.25),
            "p50": self.quantile(0.5),
            "p75": self.quantile(0.75),
            "p95": self.quantile(0.95),
            "max": self.maximum,
        }

    def to_dict(self) -> dict[str, Any]:
        return {
            "count": self.count,
            "total": self.total,
            "total_sq": self.total_sq,
            "minimum": self.minimum,
            "maximum": self.maximum,
            "histogram": self.histogram,
        }


@dataclass
class GroupSummary:
    scores: ScoreAccumulator = field(default_factory=ScoreAccumulator)
    high_count: int = 0
    low_count: int = 0

    def update(self, score: float, label: str) -> None:
        self.scores.update(score)
        if label == "high":
            self.high_count += 1
            return
        self.low_count += 1

    def row(self, name: str, value: str) -> dict[str, object]:
        return {
            name: value,
            **self.scores.summary(),
            "high_count": self.high_count,
            "low_count": self.low_count,
        }


__all__ = ["GroupSummary", "SCORE_HISTOGRAM_BINS", "ScoreAccumulator"]