Buckets:
| #!/usr/bin/env python3 | |
| """Build viewer-friendly Parquet splits for LiteFold/PDB.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import re | |
| from datetime import date, datetime | |
| from pathlib import Path | |
| import pandas as pd | |
| import pyarrow as pa | |
| import pyarrow.parquet as pq | |
| from huggingface_hub import HfApi | |
| ENTRY_COLUMNS = [ | |
| "pdb_id", | |
| "classification", | |
| "accession_date", | |
| "title", | |
| "source_organism", | |
| "authors", | |
| "raw_resolution", | |
| "experimental_method", | |
| ] | |
| def load_hf_token(env_path: Path) -> str | None: | |
| if not env_path.exists(): | |
| return None | |
| for line in env_path.read_text(encoding="utf-8").splitlines(): | |
| line = line.strip() | |
| if not line or line.startswith("#") or "=" not in line: | |
| continue | |
| key, value = line.split("=", 1) | |
| if key.strip() in {"HF_TOKEN", "HUGGINGFACE_HUB_TOKEN"}: | |
| return value.strip().strip('"').strip("'") | |
| return None | |
| def parse_accession_date(value: str, current_year: int) -> str | None: | |
| value = (value or "").strip() | |
| if not value: | |
| return None | |
| try: | |
| month, day, year = [int(part) for part in value.split("/")] | |
| except ValueError: | |
| return None | |
| current_two_digit_year = current_year % 100 | |
| full_year = 2000 + year if year <= current_two_digit_year else 1900 + year | |
| try: | |
| return date(full_year, month, day).isoformat() | |
| except ValueError: | |
| return None | |
| def parse_entries_idx(path: Path) -> pd.DataFrame: | |
| rows = [] | |
| current_year = datetime.utcnow().year | |
| with path.open("r", encoding="utf-8", errors="replace") as handle: | |
| for line_number, line in enumerate(handle, start=1): | |
| line = line.rstrip("\n") | |
| if line_number <= 2 or not line: | |
| continue | |
| parts = line.split("\t") | |
| if len(parts) < len(ENTRY_COLUMNS): | |
| parts = parts + [""] * (len(ENTRY_COLUMNS) - len(parts)) | |
| elif len(parts) > len(ENTRY_COLUMNS): | |
| parts = parts[: len(ENTRY_COLUMNS) - 1] + [" ".join(parts[len(ENTRY_COLUMNS) - 1 :])] | |
| rows.append(dict(zip(ENTRY_COLUMNS, parts))) | |
| df = pd.DataFrame.from_records(rows) | |
| df["pdb_id"] = df["pdb_id"].str.lower() | |
| df["accession_date_iso"] = df["accession_date"].map( | |
| lambda value: parse_accession_date(value, current_year) | |
| ) | |
| df["resolution_angstrom"] = pd.to_numeric(df["raw_resolution"], errors="coerce") | |
| df["resolution_is_unknown"] = df["resolution_angstrom"].isna() | |
| for column in ["classification", "title", "source_organism", "authors", "experimental_method"]: | |
| df[column] = df[column].fillna("").str.strip() | |
| return df | |
| def mmcif_rows_from_hub(repo_id: str, token: str | None) -> pd.DataFrame: | |
| api = HfApi(token=token) | |
| info = api.dataset_info(repo_id, files_metadata=True) | |
| records = [] | |
| for sibling in info.siblings or []: | |
| path = sibling.rfilename | |
| if not path.startswith("mmcif/") or not path.endswith(".cif.gz"): | |
| continue | |
| filename = Path(path).name | |
| pdb_id = filename.removesuffix(".cif.gz").lower() | |
| records.append( | |
| { | |
| "pdb_id": pdb_id, | |
| "mmcif_path": path, | |
| "mmcif_file_size_bytes": int(sibling.size) if sibling.size is not None else None, | |
| "mmcif_blob_id": sibling.blob_id, | |
| } | |
| ) | |
| return pd.DataFrame.from_records(records) | |
| def stable_bucket(value: str, buckets: int = 10) -> int: | |
| digest = hashlib.sha256(value.encode("utf-8")).hexdigest()[:16] | |
| return int(digest, 16) % buckets | |
| def write_parquet(df: pd.DataFrame, path: Path) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| table = pa.Table.from_pandas(df, preserve_index=False) | |
| pq.write_table(table, path, compression="zstd") | |
| def build_dataset(raw_dir: Path, out_dir: Path, repo_id: str, token: str | None) -> dict: | |
| entries = parse_entries_idx(raw_dir / "entries.idx") | |
| mmcif = mmcif_rows_from_hub(repo_id, token) | |
| if mmcif.empty: | |
| raise RuntimeError(f"No mmCIF files found in {repo_id}") | |
| df = mmcif.merge(entries, on="pdb_id", how="left", validate="one_to_one") | |
| df["has_entries_idx_metadata"] = df["title"].notna() | |
| for column in ["classification", "accession_date", "accession_date_iso", "title", "source_organism", "authors", "raw_resolution", "experimental_method"]: | |
| df[column] = df[column].fillna("") | |
| df["resolution_angstrom"] = df["resolution_angstrom"].astype("Float64") | |
| df["resolution_is_unknown"] = df["resolution_angstrom"].isna() | |
| df["pdb_url"] = "https://www.rcsb.org/structure/" + df["pdb_id"].str.upper() | |
| df["rcsb_download_url"] = "https://files.rcsb.org/download/" + df["pdb_id"] + ".cif.gz" | |
| df["split_bucket"] = df["pdb_id"].map(stable_bucket).astype("int64") | |
| df["split"] = df["split_bucket"].map(lambda bucket: "test" if bucket == 0 else "train") | |
| ordered_columns = [ | |
| "pdb_id", | |
| "mmcif_path", | |
| "mmcif_file_size_bytes", | |
| "mmcif_blob_id", | |
| "pdb_url", | |
| "rcsb_download_url", | |
| "classification", | |
| "accession_date", | |
| "accession_date_iso", | |
| "title", | |
| "source_organism", | |
| "authors", | |
| "raw_resolution", | |
| "resolution_angstrom", | |
| "resolution_is_unknown", | |
| "experimental_method", | |
| "has_entries_idx_metadata", | |
| "split_bucket", | |
| ] | |
| df = df[ordered_columns + ["split"]].sort_values(["split", "pdb_id"], kind="mergesort") | |
| data_dir = out_dir / "data" | |
| for split in ["train", "test"]: | |
| split_df = df[df["split"].eq(split)].drop(columns=["split"]) | |
| write_parquet(split_df, data_dir / f"{split}-00000-of-00001.parquet") | |
| metadata_dir = out_dir / "metadata" | |
| write_parquet(entries.sort_values("pdb_id", kind="mergesort"), metadata_dir / "entries_idx.parquet") | |
| method_counts = ( | |
| df["experimental_method"].replace("", "UNKNOWN").value_counts().head(20).to_dict() | |
| ) | |
| class_counts = df["classification"].replace("", "UNKNOWN").value_counts().head(20).to_dict() | |
| summary = { | |
| "source": repo_id, | |
| "full_entries_idx_rows": int(len(entries)), | |
| "mmcif_rows_in_repo": int(len(df)), | |
| "metadata_joined_rows": int(df["has_entries_idx_metadata"].sum()), | |
| "splits": { | |
| "train": int(df["split"].eq("train").sum()), | |
| "test": int(df["split"].eq("test").sum()), | |
| }, | |
| "split_strategy": "deterministic sha256(pdb_id) % 10; bucket 0 is test, buckets 1-9 are train", | |
| "total_mmcif_size_bytes": int(df["mmcif_file_size_bytes"].fillna(0).sum()), | |
| "resolution": { | |
| "known_rows": int(df["resolution_angstrom"].notna().sum()), | |
| "unknown_rows": int(df["resolution_angstrom"].isna().sum()), | |
| "min": float(df["resolution_angstrom"].min(skipna=True)), | |
| "median": float(df["resolution_angstrom"].median(skipna=True)), | |
| "mean": float(df["resolution_angstrom"].mean(skipna=True)), | |
| "max": float(df["resolution_angstrom"].max(skipna=True)), | |
| }, | |
| "top_experimental_methods": method_counts, | |
| "top_classifications": class_counts, | |
| "columns": ordered_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("--raw-dir", type=Path, default=Path("LiteFold_PDB_raw")) | |
| parser.add_argument("--out-dir", type=Path, default=Path("LiteFold_PDB_processed")) | |
| parser.add_argument("--repo-id", default="LiteFold/PDB") | |
| parser.add_argument("--env-file", type=Path, default=Path(".env")) | |
| args = parser.parse_args() | |
| summary = build_dataset(args.raw_dir, args.out_dir, args.repo_id, load_hf_token(args.env_file)) | |
| print(json.dumps(summary, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 8 kB
- Xet hash:
- cb511de2b767a0d4f69b8c51c8fa7f53fd060cd2fc688c0c02e9d3c9e90525b4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.