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

"""Build manifest of shards that lacked topics (based on audit CSV)."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Create manifest for missing-topic shards"
    )
    parser.add_argument(
        "--audit",
        type=Path,
        default=Path("artifacts/dolma_eda/topic_missing_audit.csv"),
        help="Audit CSV from topic_missing_audit.py",
    )
    parser.add_argument(
        "--workers-dir",
        type=Path,
        default=Path("runs/dolma_enriched/eda/workers"),
        help="EDA worker state directory with manifest_done.jsonl files",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("runs/dolma_enriched/manifests/missing_topic.txt"),
        help="Output manifest (one line per worker, space-separated shard paths)",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    missing_workers = set()
    with args.audit.open() as f:
        reader = csv.DictReader(f)
        for row in reader:
            if int(row["missing_topic"]) > 0:
                missing_workers.add(int(row["manifest_idx"]))

    lines: list[str] = []
    for idx in sorted(missing_workers):
        manifest_path = args.workers_dir / str(idx) / "manifest_done.jsonl"
        if not manifest_path.exists():
            continue
        shard_paths: list[str] = []
        with manifest_path.open() as f:
            for line in f:
                shard_paths.append(json.loads(line)["path"])
        if shard_paths:
            lines.append(" ".join(shard_paths))

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"wrote {len(lines)} lines to {args.output}")
    return 0


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