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

#!/usr/bin/env python3
"""Check enrichment label coverage for A-02 acceptance criteria.

Reads enriched JSONL (or JSONL.zst) output and reports coverage of
``weborganizer_topic_max`` and ``weborganizer_format_max`` fields.

Acceptance: both must be ≥ 99.5%.

Usage:
    python scripts/validation/check_enrichment_coverage.py --input-dir runs/dolma_pool_enriched
    python scripts/validation/check_enrichment_coverage.py --input-files out.enriched.jsonl.zst
"""

from __future__ import annotations

import argparse
import io
import json
import sys
from pathlib import Path
from typing import Iterable


def iter_jsonl(path: Path) -> Iterable[dict]:
    with path.open("r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if line:
                yield json.loads(line)


def iter_jsonlzst(path: Path) -> Iterable[dict]:
    import zstandard as zstd

    dctx = zstd.ZstdDecompressor()
    with path.open("rb") as fh:
        with dctx.stream_reader(fh) as reader:
            for line in io.TextIOWrapper(reader, encoding="utf-8"):
                line = line.strip()
                if line:
                    yield json.loads(line)


def iter_records(path: Path) -> Iterable[dict]:
    if path.suffix == ".zst":
        return iter_jsonlzst(path)
    return iter_jsonl(path)


def find_enriched_files(input_dir: Path) -> list[Path]:
    """Find all enriched JSONL files in a directory."""
    patterns = ["*.enriched.jsonl.zst", "*.enriched.jsonl", "*.jsonl.zst", "*.jsonl"]
    files: list[Path] = []
    for pattern in patterns:
        files.extend(sorted(input_dir.glob(pattern)))
    # Deduplicate while preserving order
    seen: set[Path] = set()
    unique: list[Path] = []
    for f in files:
        if f not in seen:
            seen.add(f)
            unique.append(f)
    return unique


def check_coverage(
    files: list[Path], *, verbose: bool = False, max_failures: int = 100
) -> dict:
    total = 0
    has_topic = 0
    has_format = 0
    failures: list[dict] = []

    for filepath in files:
        print(f"Reading: {filepath}")
        for record in iter_records(filepath):
            total += 1
            doc_id = record.get("id", f"<unknown-{total}>")
            metadata = record.get("metadata", {})

            topic_ok = bool(metadata.get("weborganizer_topic_max"))
            format_ok = bool(metadata.get("weborganizer_format_max"))

            if topic_ok:
                has_topic += 1
            if format_ok:
                has_format += 1

            if not topic_ok or not format_ok:
                reason_parts = []
                text = record.get("text")
                if not isinstance(text, str) or not text.strip():
                    reason_parts.append("empty/missing text")
                if not topic_ok:
                    reason_parts.append("missing topic")
                if not format_ok:
                    reason_parts.append("missing format")

                if len(failures) < max_failures:
                    failures.append(
                        {
                            "id": doc_id,
                            "reasons": reason_parts,
                            "file": str(filepath),
                        }
                    )

            if total % 1_000_000 == 0:
                print(f"  ... processed {total:,} records")

    topic_pct = (has_topic / total * 100) if total > 0 else 0.0
    format_pct = (has_format / total * 100) if total > 0 else 0.0

    return {
        "total_documents": total,
        "topic_labeled": has_topic,
        "format_labeled": has_format,
        "topic_coverage_pct": topic_pct,
        "format_coverage_pct": format_pct,
        "topic_missing": total - has_topic,
        "format_missing": total - has_format,
        "failures_sample": failures,
        "pass": topic_pct >= 99.5 and format_pct >= 99.5,
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Check enrichment label coverage")
    parser.add_argument("--input-dir", type=Path, help="Directory with enriched files")
    parser.add_argument("--input-files", nargs="+", type=Path, help="Specific files")
    parser.add_argument(
        "--threshold",
        type=float,
        default=99.5,
        help="Minimum coverage %% (default: 99.5)",
    )
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--output-json", type=Path, help="Write results to JSON file")
    args = parser.parse_args(argv)

    files: list[Path] = []
    if args.input_files:
        files = args.input_files
    elif args.input_dir:
        files = find_enriched_files(args.input_dir)
    else:
        parser.error("Provide --input-dir or --input-files")

    if not files:
        print("ERROR: No enriched files found")
        return 1

    print(f"Checking {len(files)} file(s)...")
    result = check_coverage(files, verbose=args.verbose)

    print()
    print("=" * 60)
    print("  A-02 Enrichment Coverage Report")
    print("=" * 60)
    print(f"  Total documents:    {result['total_documents']:>12,}")
    print(
        f"  Topic labeled:      {result['topic_labeled']:>12,}  "
        f"({result['topic_coverage_pct']:.4f}%)"
    )
    print(
        f"  Format labeled:     {result['format_labeled']:>12,}  "
        f"({result['format_coverage_pct']:.4f}%)"
    )
    print(f"  Topic missing:      {result['topic_missing']:>12,}")
    print(f"  Format missing:     {result['format_missing']:>12,}")
    print(f"  Threshold:          {args.threshold:.1f}%")
    print(f"  PASS:               {'✓ YES' if result['pass'] else '✗ NO'}")
    print("=" * 60)

    if result["failures_sample"]:
        print(
            f"\nSample of unlabeled documents ({len(result['failures_sample'])} shown):"
        )
        for f in result["failures_sample"][:20]:
            print(f"  {f['id']}: {', '.join(f['reasons'])}  [{f['file']}]")

    if args.output_json:
        args.output_json.parent.mkdir(parents=True, exist_ok=True)
        with args.output_json.open("w") as fh:
            json.dump(result, fh, indent=2)
        print(f"\nResults written to: {args.output_json}")

    return 0 if result["pass"] else 1


if __name__ == "__main__":
    sys.exit(main())