scripts/prepare_bfd_dataset.py

#!/usr/bin/env python3
"""Build viewer-friendly source index Parquet splits for LiteFold/BFD."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
from math import ceil
from pathlib import Path
from typing import Any

import pandas as pd
from huggingface_hub import HfApi


INDEX_COLUMNS = [
    "index_id",
    "repo_id",
    "source_file",
    "source_sha",
    "source_format",
    "chunk_index",
    "byte_start",
    "byte_end_exclusive",
    "chunk_size_bytes",
    "total_size_bytes",
    "chunk_size_gib",
    "is_first_chunk",
    "is_last_chunk",
    "access_note",
    "split_bucket",
]


def load_token() -> str | None:
    for key in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN"):
        value = os.environ.get(key)
        if value:
            return value
    env_path = Path(".env")
    if env_path.exists():
        for line in env_path.read_text().splitlines():
            stripped = line.strip()
            if not stripped or stripped.startswith("#") or "=" not in stripped:
                continue
            key, value = stripped.split("=", 1)
            if key.strip() in {"HF_TOKEN", "HUGGINGFACE_HUB_TOKEN"}:
                value = value.strip().strip('"').strip("'")
                if value:
                    return value
    return None


def stable_bucket(value: str, buckets: int = 10) -> int:
    digest = hashlib.sha256(value.encode("utf-8")).hexdigest()[:16]
    return int(digest, 16) % buckets


def build_dataset(repo_id: str, out_dir: Path, chunk_size_gib: int) -> dict[str, Any]:
    token = load_token()
    api = HfApi(token=token)
    info = api.dataset_info(repo_id, files_metadata=True)
    source = next(
        sibling for sibling in info.siblings or [] if sibling.rfilename.endswith(".tar.gz")
    )
    source_file = source.rfilename
    total_size = int(getattr(source, "size", 0) or 0)
    chunk_size = int(chunk_size_gib * 1024**3)
    chunk_count = ceil(total_size / chunk_size)

    rows = []
    for chunk_index in range(chunk_count):
        byte_start = chunk_index * chunk_size
        byte_end = min(byte_start + chunk_size, total_size)
        index_id = f"{source_file}:chunk-{chunk_index:06d}"
        rows.append(
            {
                "index_id": index_id,
                "repo_id": repo_id,
                "source_file": source_file,
                "source_sha": info.sha,
                "source_format": "tar.gz",
                "chunk_index": chunk_index,
                "byte_start": byte_start,
                "byte_end_exclusive": byte_end,
                "chunk_size_bytes": byte_end - byte_start,
                "total_size_bytes": total_size,
                "chunk_size_gib": chunk_size_gib,
                "is_first_chunk": chunk_index == 0,
                "is_last_chunk": chunk_index == chunk_count - 1,
                "access_note": "Compressed byte-range index for the BFD source archive; download or stream the original tar.gz for sequence records.",
                "split_bucket": stable_bucket(index_id),
            }
        )

    if out_dir.exists():
        shutil.rmtree(out_dir)
    data_dir = out_dir / "data"
    metadata_dir = out_dir / "metadata"
    data_dir.mkdir(parents=True, exist_ok=True)
    metadata_dir.mkdir(parents=True, exist_ok=True)

    df = pd.DataFrame.from_records(rows, columns=INDEX_COLUMNS)
    train = df[df["split_bucket"].ne(0)].sort_values("chunk_index", kind="mergesort")
    test = df[df["split_bucket"].eq(0)].sort_values("chunk_index", kind="mergesort")
    train.to_parquet(data_dir / "train-00000-of-00001.parquet", index=False, compression="zstd")
    test.to_parquet(data_dir / "test-00000-of-00001.parquet", index=False, compression="zstd")

    source_files = pd.DataFrame.from_records(
        [
            {
                "repo_id": repo_id,
                "filename": sibling.rfilename,
                "size_bytes": int(getattr(sibling, "size", 0) or 0),
                "source_sha": info.sha,
            }
            for sibling in sorted(info.siblings or [], key=lambda item: item.rfilename)
        ]
    )
    source_files.to_parquet(metadata_dir / "source_files.parquet", index=False, compression="zstd")

    summary = {
        "source": repo_id,
        "source_sha": info.sha,
        "viewer_table_scope": "compressed archive byte-range index",
        "source_file": source_file,
        "source_size_bytes": total_size,
        "chunk_size_gib": chunk_size_gib,
        "chunk_rows": int(len(df)),
        "splits": {"train": int(len(train)), "test": int(len(test))},
        "split_strategy": "deterministic sha256(index_id) % 10; bucket 0 is test, buckets 1-9 are train",
        "columns": INDEX_COLUMNS,
    }
    (out_dir / "dataset_summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    return summary


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo-id", default="LiteFold/BFD")
    parser.add_argument("--out-dir", type=Path, default=Path("LiteFold_BFD_processed"))
    parser.add_argument("--chunk-size-gib", type=int, default=1)
    args = parser.parse_args()
    summary = build_dataset(args.repo_id, args.out_dir, args.chunk_size_gib)
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()