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

"""Aggregation helpers for quality validation analysis."""

from __future__ import annotations

from collections import Counter, defaultdict
from dataclasses import dataclass, field
from typing import Any

from dolma.quality.sidecar import QUALITY_SCORE_BINS
from dolma.quality.validation.manifest import source_family_from_key
from dolma.quality.validation.review import ReviewSampler
from dolma.quality.validation.stats import GroupSummary, ScoreAccumulator


def coverage(hit_count: int, total_count: int) -> dict[str, float | int]:
    rate = hit_count / total_count if total_count > 0 else 0.0
    return {"hit_count": hit_count, "total_count": total_count, "rate": rate}


@dataclass
class QualityValidationAggregator:
    total_docs: int = 0
    raw_join_missing: int = 0
    soc91_join_hits: int = 0
    soc91_join_missing_rows: int = 0
    soc91_missing_shards: int = 0
    probability_sum_failures: int = 0
    quality_score_failures: int = 0
    quality_confidence_failures: int = 0
    label_id_counts: Counter[int] = field(default_factory=Counter)
    label_counts: Counter[str] = field(default_factory=Counter)
    score_histogram: list[int] = field(
        default_factory=lambda: [0 for _ in range(QUALITY_SCORE_BINS)]
    )
    scores: ScoreAccumulator = field(default_factory=ScoreAccumulator)
    confidences: ScoreAccumulator = field(default_factory=ScoreAccumulator)
    by_source_family: dict[str, GroupSummary] = field(
        default_factory=lambda: defaultdict(GroupSummary)
    )
    by_topic: dict[str, GroupSummary] = field(
        default_factory=lambda: defaultdict(GroupSummary)
    )
    by_format: dict[str, GroupSummary] = field(
        default_factory=lambda: defaultdict(GroupSummary)
    )
    shard_rows: list[dict[str, object]] = field(default_factory=list)
    review_sampler: ReviewSampler = field(default_factory=ReviewSampler)

    def update_row(
        self,
        *,
        source_key: str,
        row: dict[str, Any],
        raw_doc: dict[str, Any] | None,
        soc91_doc: dict[str, Any] | None,
    ) -> None:
        score = float(row["quality_score"])
        high_prob = float(row["quality_high_prob"])
        low_prob = float(row["quality_low_prob"])
        confidence = float(row["quality_confidence"])
        label = "high" if high_prob >= low_prob else "low"
        self.total_docs += 1
        self.label_counts[label] += 1
        self.label_id_counts[int(row["quality_label_id"])] += 1
        self.scores.update(score)
        self.confidences.update(confidence)
        index = min(int(score * QUALITY_SCORE_BINS), QUALITY_SCORE_BINS - 1)
        self.score_histogram[index] += 1
        if abs((high_prob + low_prob) - 1.0) > 1e-3:
            self.probability_sum_failures += 1
        if abs(score - high_prob) > 1e-6:
            self.quality_score_failures += 1
        if abs(confidence - max(high_prob, low_prob)) > 1e-6:
            self.quality_confidence_failures += 1
        source_family = source_family_from_key(source_key)
        if raw_doc is None:
            self.raw_join_missing += 1
        else:
            source_family = str(raw_doc.get("source_family") or source_family)
        self.by_source_family[source_family].update(score, label)
        topic_label = (
            soc91_doc.get("topic_url_label") if soc91_doc is not None else None
        )
        format_label = (
            soc91_doc.get("format_url_label") if soc91_doc is not None else None
        )
        if soc91_doc is None:
            self.soc91_join_missing_rows += 1
        else:
            self.soc91_join_hits += 1
            if isinstance(topic_label, str):
                self.by_topic[topic_label].update(score, label)
            if isinstance(format_label, str):
                self.by_format[format_label].update(score, label)
        self.review_sampler.update(
            {
                "doc_id": str(row["doc_id"]),
                "source_key": source_key,
                "source_family": source_family,
                "quality_score": score,
                "quality_confidence": confidence,
                "quality_label": label,
                "quality_label_id": int(row["quality_label_id"]),
                "topic_url_label": topic_label,
                "format_url_label": format_label,
                "text_snippet": raw_doc.get("text_snippet") if raw_doc else "",
                "url": raw_doc.get("url") if raw_doc else None,
            }
        )


__all__ = ["QualityValidationAggregator", "coverage"]