social-data-attribution-w2/scripts/manifests/make_pool_manifest.py

#!/usr/bin/env python3
"""Generate shard manifest for HCAI-Lab/archive-dolma3-pool-150b.

Lists all .jsonl.zst files from the HF repo and writes one path per line.
Line number (0-indexed) maps to SLURM_ARRAY_TASK_ID.

Usage:
    python scripts/manifests/make_pool_manifest.py
    python scripts/manifests/make_pool_manifest.py --repo-id HCAI-Lab/archive-dolma3-pool-150b --output manifest.txt
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

from huggingface_hub import HfApi


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Generate pool shard manifest")
    parser.add_argument(
        "--repo-id",
        default="HCAI-Lab/archive-dolma3-pool-150b",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("scripts/slurm/pool_shard_manifest.txt"),
    )
    parser.add_argument(
        "--pattern",
        default="*.jsonl.zst",
        help="Glob pattern to match shard files",
    )
    args = parser.parse_args(argv)

    token = os.environ.get("HF_TOKEN")
    api = HfApi(token=token)

    print(f"Listing files in {args.repo_id} matching {args.pattern}...")
    all_files = api.list_repo_files(args.repo_id, repo_type="dataset", token=token)
    shards = sorted(f for f in all_files if f.endswith(".jsonl.zst") and "/" in f)

    if not shards:
        print(f"ERROR: No .jsonl.zst files found in {args.repo_id}")
        return 1

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w") as fh:
        for shard in shards:
            fh.write(shard + "\n")

    print(f"Wrote {len(shards)} shard paths to {args.output}")
    print(f"Array range: 0-{len(shards) - 1}")
    return 0


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