social-data-attribution-w2/scripts/analysis/check_nonpool_coverage.py

"""Read-only: does the non-pool (phase2) portion already have WebOrganizer
labels? Lists R2 source shards under soc127/phase2_nonpool_final and checks
soc91-labels/ for any sidecars referencing non-pool families. Determines
whether completing the corpus needs a GPU labeling run or just a manifest
rebuild. Run via scripts/bootstrap/with_r2_credentials.sh."""

from __future__ import annotations

import os
from collections import Counter

import boto3

R2_ENDPOINT = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"
R2_BUCKET = "soc127-dedup"
NONPOOL_PREFIX = "soc127/phase2_nonpool_final/"
LABELS_PREFIX = "soc91-labels/"
NONPOOL_FAMILIES = ("stack_edu", "finemath", "proofpile", "arxiv", "wiki", "dolma1_7")


def _client():
    return boto3.client(
        "s3",
        endpoint_url=R2_ENDPOINT,
        aws_access_key_id=os.environ["R2_ACCESS_KEY_ID"],
        aws_secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"],
        region_name="auto",
    )


def _list(client, prefix: str, suffix: str | None = None) -> list[str]:
    keys: list[str] = []
    paginator = client.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=R2_BUCKET, Prefix=prefix):
        for obj in page.get("Contents", []):
            k = obj["Key"]
            if suffix is None or k.endswith(suffix):
                keys.append(k)
    return keys


def main() -> None:
    client = _client()

    src = _list(client, NONPOOL_PREFIX)
    print(f"phase2_nonpool_final objects: {len(src)}")
    for k in src[:5]:
        print("  src:", k)

    labels = _list(client, LABELS_PREFIX, suffix=".parquet")
    print(f"\nsoc91-labels parquet sidecars (total): {len(labels)}")
    fam = Counter()
    nonpool_hits = []
    for k in labels:
        name = k.split("/")[-1].lower()
        for f in NONPOOL_FAMILIES:
            if f in name:
                fam[f] += 1
                if len(nonpool_hits) < 5:
                    nonpool_hits.append(k)
    print(f"soc91-labels sidecars matching non-pool families: {sum(fam.values())}")
    print(f"  by family: {dict(fam)}")
    for k in nonpool_hits:
        print("  nonpool-label sample:", k)


if __name__ == "__main__":
    main()