social-data-attribution-w2/tests/quality/validation/test_analysis.py

from __future__ import annotations

from dolma.quality.validation import analyze_validation_sources


def test_analyze_validation_sources_aggregates_joins_and_review_rows(
    monkeypatch,
) -> None:
    source_a = "soc127/phase1_pool_shared/data/c4-en/shard_000.jsonl.zst"
    source_b = "soc127/phase2_nonpool_final/data/wiki-en/shard_001.jsonl.zst"

    quality_rows = {
        source_a: [
            {
                "doc_id": "a",
                "quality_label_id": 1,
                "quality_score": 0.95,
                "quality_high_prob": 0.95,
                "quality_low_prob": 0.05,
                "quality_confidence": 0.95,
            },
            {
                "doc_id": "b",
                "quality_label_id": 0,
                "quality_score": 0.05,
                "quality_high_prob": 0.05,
                "quality_low_prob": 0.95,
                "quality_confidence": 0.95,
            },
        ],
        source_b: [
            {
                "doc_id": "c",
                "quality_label_id": 1,
                "quality_score": 0.60,
                "quality_high_prob": 0.60,
                "quality_low_prob": 0.45,
                "quality_confidence": 0.60,
            }
        ],
    }
    raw_docs = {
        source_a: {
            "a": {
                "source_family": "common_crawl",
                "text_snippet": "high score sample",
                "url": "https://example.com/a",
            },
            "b": {
                "source_family": "common_crawl",
                "text_snippet": "low score sample",
                "url": "https://example.com/b",
            },
        },
        source_b: {},
    }
    soc91_docs = {
        source_a: {
            "a": {
                "topic_url_label": "shopping",
                "format_url_label": "product_page",
            },
            "b": {
                "topic_url_label": "science",
                "format_url_label": "academic_writing",
            },
        },
        source_b: None,
    }

    monkeypatch.setattr(
        "dolma.quality.validation.analysis.read_quality_rows",
        lambda *_args, source_key, **_kwargs: quality_rows[source_key],
    )
    monkeypatch.setattr(
        "dolma.quality.validation.analysis.read_raw_doc_map",
        lambda *_args, source_key, **_kwargs: raw_docs[source_key],
    )
    monkeypatch.setattr(
        "dolma.quality.validation.analysis.read_soc91_doc_map",
        lambda *_args, source_key, **_kwargs: soc91_docs[source_key],
    )

    payload = analyze_validation_sources(
        object(),
        bucket="bucket",
        source_keys=[source_a, source_b],
        output_prefix="soc139-quality-validation/test",
        soc91_prefix="soc91-labels",
    )

    assert payload["summary"]["processed_shards"] == 2
    assert payload["summary"]["total_docs"] == 3
    assert payload["summary"]["raw_join_missing"] == 1
    assert payload["summary"]["soc91_join_hits"] == 2
    assert payload["summary"]["soc91_join_missing_rows"] == 1
    assert payload["summary"]["soc91_missing_shards"] == 1
    assert payload["summary"]["label_counts"] == {"high": 2, "low": 1}
    assert payload["summary"]["label_id_counts"] == {0: 1, 1: 2}
    assert payload["summary"]["consistency_checks"] == {
        "probability_sum_failures": 1,
        "quality_score_failures": 0,
        "quality_confidence_failures": 0,
    }

    assert payload["join_coverage"]["raw_docs"] == {
        "hit_count": 2,
        "total_count": 3,
        "rate": 2 / 3,
    }
    assert payload["join_coverage"]["soc91_sidecars"] == {
        "hit_count": 2,
        "total_count": 3,
        "rate": 2 / 3,
    }
    assert payload["join_coverage"]["missing_soc91_shards"] == 1

    histogram_total = sum(row["count"] for row in payload["histogram_rows"])
    assert histogram_total == 3

    source_family_rows = {
        row["source_family"]: row for row in payload["source_family_rows"]
    }
    assert source_family_rows["common_crawl"]["count"] == 2
    assert source_family_rows["wiki"]["count"] == 1

    topic_rows = {row["topic_url_label"]: row for row in payload["topic_rows"]}
    assert set(topic_rows) == {"science", "shopping"}

    format_rows = {row["format_url_label"]: row for row in payload["format_rows"]}
    assert set(format_rows) == {"academic_writing", "product_page"}

    review_buckets = {row["review_bucket"] for row in payload["review_rows"]}
    assert review_buckets == {"anomaly", "high", "low", "mid"}