social-data-attribution-w2/src/dolma/eda/stats.py

"""Common stats helpers for Dolma EDA."""

from __future__ import annotations

import bisect
from dataclasses import dataclass, field
from typing import Any, Iterable

DEFAULT_WORD_BINS = [
    0,
    50,
    100,
    200,
    500,
    1000,
    2000,
    5000,
    10000,
    20000,
    50000,
    100000,
    200000,
    500000,
    1000000,
]


@dataclass
class Histogram:
    edges: list[int]
    counts: list[int]

    @classmethod
    def with_edges(cls, edges: Iterable[int]) -> "Histogram":
        edge_list = list(edges)
        return cls(edge_list, [0 for _ in range(len(edge_list) - 1)])

    def update(self, value: float) -> None:
        if not self.counts:
            return
        index = bisect.bisect_right(self.edges, value) - 1
        index = max(0, min(index, len(self.counts) - 1))
        self.counts[index] += 1

    def total(self) -> int:
        return sum(self.counts)

    def quantile(self, quantile: float) -> float | None:
        total = self.total()
        if total == 0:
            return None
        target = quantile * total
        running = 0
        for idx, count in enumerate(self.counts):
            running += count
            if running >= target:
                left = self.edges[idx]
                right = self.edges[idx + 1]
                return float(left + (right - left) * (running - target) / count)
        return float(self.edges[-1])

    def to_dict(self) -> dict[str, Any]:
        return {"edges": self.edges, "counts": self.counts}

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "Histogram":
        return cls(list(payload["edges"]), list(payload["counts"]))


@dataclass
class LengthStats:
    histogram: Histogram = field(
        default_factory=lambda: Histogram.with_edges(DEFAULT_WORD_BINS)
    )
    count: int = 0
    total: float = 0.0
    minimum: int | None = None
    maximum: int | None = None
    missing: int = 0

    def update_value(self, value: float | int | None) -> None:
        if value is None:
            self.missing += 1
            return
        self.count += 1
        self.total += float(value)
        int_value = int(value)
        self.minimum = (
            int_value if self.minimum is None else min(self.minimum, int_value)
        )
        self.maximum = (
            int_value if self.maximum is None else max(self.maximum, int_value)
        )
        self.histogram.update(float(value))

    def summary(self) -> dict[str, Any]:
        mean = self.total / self.count if self.count else None
        return {
            "count": self.count,
            "missing": self.missing,
            "min": self.minimum,
            "max": self.maximum,
            "mean": mean,
            "p50": self.histogram.quantile(0.5),
            "p90": self.histogram.quantile(0.9),
            "p99": self.histogram.quantile(0.99),
            "histogram": self.histogram.to_dict(),
        }

    def merge(self, other: "LengthStats") -> None:
        if self.histogram.edges != other.histogram.edges:
            raise ValueError("Histogram edges do not match")
        self.count += other.count
        self.total += other.total
        self.missing += other.missing
        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.counts = [
            left + right
            for left, right in zip(self.histogram.counts, other.histogram.counts)
        ]

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

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "LengthStats":
        histogram = Histogram.from_dict(payload["histogram"])
        return cls(
            histogram=histogram,
            count=payload["count"],
            total=payload["total"],
            minimum=payload["minimum"],
            maximum=payload["maximum"],
            missing=payload["missing"],
        )


__all__ = ["DEFAULT_WORD_BINS", "Histogram", "LengthStats"]