social-data-attribution-w2/scripts/corrections/verify_flip.py

"""Verify that flipped shards have correct column relationships.

Reads shards marked with .label-fix.done and checks:
1. quality_score == quality_high_prob (invariant)
2. quality_confidence == max(quality_high_prob, quality_low_prob)
3. quality_high_prob + quality_low_prob ≈ 1.0
4. Stats JSON label_histogram and score_histogram are consistent
"""

from __future__ import annotations

import io
import json
import logging
import os
import sys

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger(__name__)

R2_BUCKET = "soc127-dedup"
R2_ENDPOINT_URL = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"
SIDECAR_PREFIX = "soc139-quality-sidecars"
DONE_SUFFIX = ".label-fix.done"


def main() -> None:
    import pyarrow.parquet as pq
    import boto3

    limit = int(sys.argv[1]) if len(sys.argv) > 1 else 5

    client = boto3.client(
        "s3",
        endpoint_url=R2_ENDPOINT_URL,
        aws_access_key_id=os.environ["R2_ACCESS_KEY_ID"],
        aws_secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"],
        region_name="auto",
    )

    paginator = client.get_paginator("list_objects_v2")
    done_keys: list[str] = []
    prefix = SIDECAR_PREFIX.rstrip("/") + "/"
    for page in paginator.paginate(Bucket=R2_BUCKET, Prefix=prefix):
        for item in page.get("Contents", []):
            if item["Key"].endswith(DONE_SUFFIX):
                done_keys.append(item["Key"])
                if len(done_keys) >= limit:
                    break
        if len(done_keys) >= limit:
            break

    log.info("Found %d flipped shards to verify", len(done_keys))
    errors = 0

    for done_key in done_keys:
        parquet_key = done_key.removesuffix(DONE_SUFFIX)
        stats_key = f"{parquet_key}.stats.json"

        raw = client.get_object(Bucket=R2_BUCKET, Key=parquet_key)["Body"].read()
        table = pq.read_table(io.BytesIO(raw))

        high = table.column("quality_high_prob").to_pylist()
        low = table.column("quality_low_prob").to_pylist()
        score = table.column("quality_score").to_pylist()
        conf = table.column("quality_confidence").to_pylist()

        shard_errors = 0
        for i in range(table.num_rows):
            if abs(score[i] - high[i]) > 1e-7:
                shard_errors += 1
            if abs(conf[i] - max(high[i], low[i])) > 1e-7:
                shard_errors += 1
            if abs(high[i] + low[i] - 1.0) > 0.01:
                shard_errors += 1

        try:
            stats_raw = client.get_object(Bucket=R2_BUCKET, Key=stats_key)[
                "Body"
            ].read()
            stats = json.loads(stats_raw)
            hist = stats.get("label_histogram", {})
            docs_classified = stats.get("docs_classified", 0)
            hist_total = hist.get("high", 0) + hist.get("low", 0)
            if docs_classified > 0 and hist_total != docs_classified:
                log.info(
                    "  WARN: %s histogram sum %d != docs_classified %d",
                    parquet_key,
                    hist_total,
                    docs_classified,
                )
                shard_errors += 1

            score_hist = stats.get("score_histogram", [])
            score_hist_total = sum(score_hist)
            if docs_classified > 0 and score_hist_total != docs_classified:
                log.info(
                    "  WARN: %s score histogram sum %d != docs_classified %d",
                    parquet_key,
                    score_hist_total,
                    docs_classified,
                )
                shard_errors += 1
        except Exception:
            log.info("  No stats JSON for %s", parquet_key)

        status = "OK" if shard_errors == 0 else f"ERRORS: {shard_errors}"
        log.info(
            "[%s] %s (%d rows, high_mean=%.4f, low_mean=%.4f)",
            status,
            parquet_key,
            table.num_rows,
            sum(high) / max(len(high), 1),
            sum(low) / max(len(low), 1),
        )
        errors += shard_errors

    log.info(
        "\nVerification complete: %d total errors across %d shards",
        errors,
        len(done_keys),
    )
    sys.exit(1 if errors > 0 else 0)


if __name__ == "__main__":
    main()