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

"""Report per-subcategory completion for the shard manifest vs R2 .done files."""

import argparse
import os
import re
from collections import Counter
from pathlib import Path

import boto3

R2_BUCKET = "soc127-dedup"
R2_ENDPOINT = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"


def get_done_set():
    s3 = 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",
    )
    p = s3.get_paginator("list_objects_v2")
    done = set()
    for pg in p.paginate(Bucket=R2_BUCKET, Prefix="soc91-labels/"):
        for o in pg.get("Contents", []):
            k = o["Key"]
            if k.endswith(".done"):
                base = k.split("/")[-1].replace(".parquet.done", "")
                done.add(base)
    return done


def categorize_shard(shard_path):
    fname = shard_path.split("/")[-1]
    if "phase2" in shard_path:
        return "phase2_nonpool"
    elif "olmocr" in shard_path:
        m = re.match(r"data__olmocr_science_pdfs-(.+?)__", fname)
        return "olmocr_" + (m.group(1) if m else "unknown")
    else:
        m = re.match(r"data__common_crawl-(.+?)-\d{4}__", fname)
        return "cc_" + (m.group(1) if m else "unknown")


def main():
    parser = argparse.ArgumentParser(
        description="Report per-subcategory shard completion"
    )
    parser.add_argument("--manifest", type=Path, required=True)
    args = parser.parse_args()

    print("Fetching R2 .done files...")
    done = get_done_set()
    print(f"Found {len(done)} completed shards\n")

    lines = args.manifest.read_text().splitlines()

    total = Counter()
    done_count = Counter()

    for line in lines:
        for shard in line.split():
            fname = shard.split("/")[-1]
            base = fname.replace(".jsonl.zst", "")
            cat = categorize_shard(shard)
            total[cat] += 1
            if base in done:
                done_count[cat] += 1

    fmt = "{:<45} {:>6} {:>6} {:>9} {:>6}"
    hdr = fmt.format("Category", "Total", "Done", "Remain", "Done%")
    print(hdr)
    print("-" * len(hdr))

    for cat in sorted(total, key=lambda c: done_count.get(c, 0) / max(total[c], 1)):
        t = total[cat]
        d = done_count.get(cat, 0)
        r = t - d
        pct = d / t * 100 if t else 0
        print(fmt.format(cat, t, d, r, f"{pct:.1f}%"))

    print("-" * len(hdr))
    tot_t = sum(total.values())
    tot_d = sum(done_count.values())
    print(
        fmt.format("TOTAL", tot_t, tot_d, tot_t - tot_d, f"{tot_d / tot_t * 100:.1f}%")
    )


if __name__ == "__main__":
    main()