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

"""Build a per-shard audit of topic label coverage from EDA worker states."""

from __future__ import annotations

import argparse
import csv
import json
import re
from pathlib import Path
from typing import Iterable


def _dataset_from_path(path: str) -> str:
    match = re.search(r"data_([^/]+)_shard", path)
    if match:
        return match.group(1)
    return "unknown"


def _iter_states(worker_root: Path) -> Iterable[tuple[int, Path, dict]]:
    for entry in sorted(worker_root.iterdir(), key=lambda p: int(p.name)):
        state_path = entry / "state.json"
        if not state_path.exists():
            continue
        with state_path.open() as f:
            state = json.load(f)
        yield int(entry.name), state_path, state


def build_rows(worker_root: Path) -> list[dict[str, str | int]]:
    rows: list[dict[str, str | int]] = []
    for manifest_idx, state_path, state in _iter_states(worker_root):
        total = int(state.get("total_records", 0))
        labeled = int(sum(state.get("topic_counts", {}).values()))
        missing = total - labeled
        # Look at first shard from manifest_done to capture dataset name
        manifest_path = state_path.parent / "manifest_done.jsonl"
        first_path = ""
        if manifest_path.exists():
            with manifest_path.open() as f:
                for line in f:
                    first_path = json.loads(line)["path"]
                    break
        dataset = _dataset_from_path(first_path)
        rows.append(
            {
                "manifest_idx": manifest_idx,
                "dataset": dataset,
                "example_shard": first_path,
                "total_records": total,
                "topic_labeled": labeled,
                "missing_topic": missing,
            }
        )
    return rows


def write_csv(rows: list[dict[str, str | int]], output: Path) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "manifest_idx",
        "dataset",
        "example_shard",
        "total_records",
        "topic_labeled",
        "missing_topic",
    ]
    with output.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Audit topic label coverage")
    parser.add_argument(
        "--workers-dir",
        type=Path,
        default=Path("runs/dolma_enriched/eda/workers"),
        help="Path to EDA worker state directory",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("artifacts/dolma_eda/topic_missing_audit.csv"),
        help="Output CSV path",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    rows = build_rows(args.workers_dir)
    write_csv(rows, args.output)
    total = sum(r["total_records"] for r in rows)
    labeled = sum(r["topic_labeled"] for r in rows)
    missing = sum(r["missing_topic"] for r in rows)
    print(f"wrote {len(rows)} rows to {args.output}")
    print(f"total={total} topic_labeled={labeled} missing={missing}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())