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

"""Validation analysis for quality sidecars."""

from __future__ import annotations

from dolma.quality.sidecar import QUALITY_SCORE_BINS
from dolma.quality.validation.aggregate import QualityValidationAggregator, coverage
from dolma.quality.validation.io import (
    read_quality_rows,
    read_raw_doc_map,
    read_soc91_doc_map,
)
from dolma.quality.validation.manifest import source_family_from_key


def analyze_validation_sources(
    client: object,
    *,
    bucket: str,
    source_keys: list[str],
    output_prefix: str,
    soc91_prefix: str,
) -> dict[str, object]:
    aggregator = QualityValidationAggregator()
    for source_key in source_keys:
        raw_docs = read_raw_doc_map(client, bucket=bucket, source_key=source_key)
        quality_rows = read_quality_rows(
            client,
            bucket=bucket,
            source_key=source_key,
            output_prefix=output_prefix,
        )
        soc91_docs = read_soc91_doc_map(
            client,
            bucket=bucket,
            source_key=source_key,
            soc91_prefix=soc91_prefix,
        )
        if soc91_docs is None:
            aggregator.soc91_missing_shards += 1
            soc91_docs = {}
        aggregator.shard_rows.append(
            {
                "source_key": source_key,
                "doc_count": len(quality_rows),
                "source_family": source_family_from_key(source_key),
                "soc91_sidecar_found": bool(soc91_docs),
            }
        )
        for row in quality_rows:
            doc_id = str(row["doc_id"])
            aggregator.update_row(
                source_key=source_key,
                row=row,
                raw_doc=raw_docs.get(doc_id),
                soc91_doc=soc91_docs.get(doc_id),
            )
    return _payload_for_aggregator(aggregator, source_keys)


def _payload_for_aggregator(
    aggregator: QualityValidationAggregator,
    source_keys: list[str],
) -> dict[str, object]:
    total_docs = aggregator.total_docs
    return {
        "summary": {
            "processed_shards": len(source_keys),
            "total_docs": total_docs,
            "raw_join_missing": aggregator.raw_join_missing,
            "soc91_join_hits": aggregator.soc91_join_hits,
            "soc91_join_missing_rows": aggregator.soc91_join_missing_rows,
            "soc91_missing_shards": aggregator.soc91_missing_shards,
            "label_counts": dict(sorted(aggregator.label_counts.items())),
            "label_id_counts": dict(sorted(aggregator.label_id_counts.items())),
            "quality_score": aggregator.scores.summary(),
            "quality_confidence": aggregator.confidences.summary(),
            "consistency_checks": {
                "probability_sum_failures": aggregator.probability_sum_failures,
                "quality_score_failures": aggregator.quality_score_failures,
                "quality_confidence_failures": aggregator.quality_confidence_failures,
            },
        },
        "join_coverage": {
            "raw_docs": coverage(total_docs - aggregator.raw_join_missing, total_docs),
            "soc91_sidecars": coverage(aggregator.soc91_join_hits, total_docs),
            "missing_soc91_shards": aggregator.soc91_missing_shards,
        },
        "shard_rows": sorted(
            aggregator.shard_rows, key=lambda item: str(item["source_key"])
        ),
        "histogram_rows": [
            {
                "bin_start": index / QUALITY_SCORE_BINS,
                "bin_end": (index + 1) / QUALITY_SCORE_BINS,
                "count": count,
            }
            for index, count in enumerate(aggregator.score_histogram)
        ],
        "label_rows": [
            {"quality_label": label, "count": count}
            for label, count in sorted(aggregator.label_counts.items())
        ],
        "source_family_rows": [
            summary.row("source_family", key)
            for key, summary in sorted(aggregator.by_source_family.items())
        ],
        "topic_rows": [
            summary.row("topic_url_label", key)
            for key, summary in sorted(aggregator.by_topic.items())
        ],
        "format_rows": [
            summary.row("format_url_label", key)
            for key, summary in sorted(aggregator.by_format.items())
        ],
        "review_rows": aggregator.review_sampler.rows(),
    }


__all__ = ["analyze_validation_sources"]