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

"""Read-only probe of HF manifest + working-sample repos to settle the
1.0986B (pool) vs 1.26B (mix) question. Uses the cached HF token."""

from __future__ import annotations

import json
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_download

api = HfApi()

MANIFEST_REPOS = [
    "HCAI-Lab/dolma3-6t-corpus-manifest",
    "HCAI-Lab/dolma3-olmo3-corpus-manifest",
]
SAMPLE_REPOS = [
    "HCAI-Lab/dolma3-6t-sample-10000-docs",
    "HCAI-Lab/dolma3_6T_sample_10000_docs",
]


def probe_repo_files(repo: str) -> None:
    try:
        info = api.repo_info(repo, repo_type="dataset", files_metadata=True)
    except Exception as exc:  # noqa: BLE001
        print(f"[{repo}] repo_info FAILED: {exc}")
        return
    parquet = [s for s in info.siblings if s.rfilename.endswith(".parquet")]
    total_bytes = sum((s.size or 0) for s in parquet)
    print(f"[{repo}] parquet_files={len(parquet)} total_bytes={total_bytes/1e9:.1f}GB")
    # sample a few filenames
    for s in parquet[:3]:
        print(f"    e.g. {s.rfilename} ({(s.size or 0)/1e6:.1f}MB)")


def probe_sample(repo: str) -> bool:
    try:
        files = api.list_repo_files(repo, repo_type="dataset")
    except Exception as exc:  # noqa: BLE001
        print(f"[{repo}] list FAILED: {exc}")
        return False
    print(f"[{repo}] exists; top files: {[f for f in files if not f.startswith('data/')][:8]}")
    out = Path("/tmp/hf_probe")
    out.mkdir(exist_ok=True)
    for fn in ("sample_contract.json", "bin_summary.csv"):
        if fn in files:
            p = hf_hub_download(repo, fn, repo_type="dataset", local_dir=str(out))
            print(f"    downloaded {fn} -> {p}")
    return True


if __name__ == "__main__":
    print("=== MANIFEST REPOS ===")
    for r in MANIFEST_REPOS:
        probe_repo_files(r)
    print("=== SAMPLE REPOS ===")
    for r in SAMPLE_REPOS:
        if probe_sample(r):
            break