social-data-attribution-w2/scripts/validation/validate_pool_sample.py

#!/usr/bin/env python3
"""Validate a pool sample against SOC-51 acceptance criteria.

Checks:
1. Total tokens >= 150B from manifest parquet
2. Spot-check 50 random docs: non-empty text in JSONL.zst shards
3. weborganizer_topic distribution for bias inspection
"""

from __future__ import annotations

import argparse
import io
import json
import logging
import random
import sys
from pathlib import Path

import pyarrow.parquet as pq
import zstandard as zstd

logger = logging.getLogger(__name__)

TARGET_TOKENS = 150_000_000_000
SPOT_CHECK_COUNT = 50


def check_token_total(manifest_path: Path) -> tuple[bool, int]:
    table = pq.read_table(manifest_path, columns=["token_count"])
    total = table.column("token_count").to_pylist()
    token_sum = sum(t for t in total if t is not None)
    passed = token_sum >= TARGET_TOKENS
    return passed, token_sum


def spot_check_docs(
    manifest_path: Path,
    shards_dir: Path,
    count: int = SPOT_CHECK_COUNT,
    seed: int = 42,
) -> tuple[bool, int, int]:
    table = pq.read_table(manifest_path, columns=["doc_id", "shard_path"])
    rows = table.to_pylist()
    rng = random.Random(seed)
    sample = rng.sample(rows, min(count, len(rows)))

    shard_to_docs: dict[str, set[str]] = {}
    for row in sample:
        shard = row["shard_path"]
        doc_id = row["doc_id"]
        shard_to_docs.setdefault(shard, set()).add(doc_id)

    found = 0
    checked = 0

    for shard_name, target_ids in shard_to_docs.items():
        shard_path = _find_shard(shards_dir, shard_name)
        if shard_path is None:
            logger.warning("Shard not found: %s", shard_name)
            continue

        remaining = set(target_ids)
        dctx = zstd.ZstdDecompressor()
        with shard_path.open("rb") as fh:
            with dctx.stream_reader(fh) as reader:
                for line in io.TextIOWrapper(reader, encoding="utf-8"):
                    if not remaining:
                        break
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        record = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    doc_id = record.get("id")
                    if doc_id in remaining:
                        remaining.discard(doc_id)
                        checked += 1
                        text = record.get("text", "")
                        if text and text.strip():
                            found += 1
                        else:
                            logger.warning(
                                "Empty text for doc_id=%s in %s",
                                doc_id,
                                shard_name,
                            )

    return found == checked, found, checked


def _find_shard(shards_dir: Path, shard_name: str) -> Path | None:
    direct = shards_dir / shard_name
    if direct.exists():
        return direct
    for candidate in shards_dir.rglob(shard_name):
        return candidate
    return None


def report_topic_distribution(manifest_path: Path) -> dict[str, int]:
    table = pq.read_table(manifest_path, columns=["weborganizer_topic"])
    topics = table.column("weborganizer_topic").to_pylist()
    dist: dict[str, int] = {}
    for topic in topics:
        key = topic if topic is not None else "<null>"
        dist[key] = dist.get(key, 0) + 1
    return dict(sorted(dist.items(), key=lambda x: -x[1]))


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="Validate pool sample")
    parser.add_argument(
        "--manifest", type=Path, required=True, help="Manifest parquet path"
    )
    parser.add_argument(
        "--shards-dir", type=Path, required=True, help="Directory with output shards"
    )
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--verbose", action="store_true", default=False)
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )

    all_passed = True

    logger.info("Check 1: Token total >= 150B")
    passed, token_sum = check_token_total(args.manifest)
    logger.info(
        "  Total tokens: %s (%.2fB) - %s",
        f"{token_sum:,}",
        token_sum / 1e9,
        "PASS" if passed else "FAIL",
    )
    all_passed = all_passed and passed

    logger.info("Check 2: Spot-check %d docs", SPOT_CHECK_COUNT)
    passed, found, checked = spot_check_docs(
        args.manifest, args.shards_dir, seed=args.seed
    )
    logger.info(
        "  Non-empty text: %d / %d - %s",
        found,
        checked,
        "PASS" if passed else "FAIL",
    )
    all_passed = all_passed and passed

    logger.info("Check 3: Topic distribution")
    dist = report_topic_distribution(args.manifest)
    for topic, count in list(dist.items())[:20]:
        logger.info("  %-40s %s", topic, f"{count:,}")
    if len(dist) > 20:
        logger.info("  ... and %d more topics", len(dist) - 20)

    if all_passed:
        logger.info("All checks PASSED")
        sys.exit(0)
    else:
        logger.error("Some checks FAILED")
        sys.exit(1)


if __name__ == "__main__":
    main()