OneScience-Group/AlphaFold3

1#!/usr/bin/env python32from __future__ import annotations3 4import argparse5import hashlib6import json7from pathlib import Path8 9try:10    import yaml11except ImportError as exc:12    raise SystemExit("PyYAML is required for preflight: pip install pyyaml") from exc13 14 15MODEL_FILES = {16    "af3.bin": (1146811260, "a43670ea1fae790cd30ccc8c8bf836c4098014c45250fb310a5338407e321de7"),17    "libflash_atten_c.so": (75882104, "2ea956e8d13317dd795ffce846079f20e502a39df66802f485d5f25900ef493b"),18    "mmseqs/bin/mmseqs": (14417192, "d166c67ca79089ea5f89ec1133f67ab1c216e1734a0408bbc7205e588fe865ed"),19    "mmseqs/lib/libmarv.so": (58277568, "fdc787467b6d05d36db1f034a2cb8944b5b072bbe6d46172013810a066583928"),20}21 22DATASET_REQUIRED = [23    "infer_input_data/readme.md",24    "infer_input_data/all_data/7r6r_data.json",25    "infer_input_data/all_data/t1119_data.json",26    "public_databases/pdb_seqres_2022_09_28.fasta",27    "mmseqsDB/small_bfd_db",28    "mmseqsDB/small_bfd_db.dbtype",29    "jackhmmer_split/bfd-first_non_consensus_sequences.fasta-00000-of-00064",30]31 32 33def sha256(path: Path) -> str:34    h = hashlib.sha256()35    with path.open("rb") as f:36        for chunk in iter(lambda: f.read(1024 * 1024), b""):37            h.update(chunk)38    return h.hexdigest()39 40 41def check_json(path: Path) -> None:42    data = json.loads(path.read_text(encoding="utf-8"))43    if data.get("dialect") != "alphafold3":44        raise AssertionError(f"{path} dialect is not alphafold3")45    if not data.get("name"):46        raise AssertionError(f"{path} missing name")47    if not isinstance(data.get("sequences"), list) or not data["sequences"]:48        raise AssertionError(f"{path} missing sequences")49 50 51def main() -> None:52    parser = argparse.ArgumentParser()53    parser.add_argument("--root", default=".")54    parser.add_argument("--dataset-root", default="data/alphafold3_dataset")55    parser.add_argument("--model-dir", default="checkpoints/AlphaFold3")56    parser.add_argument("--skip-hash", action="store_true")57    args = parser.parse_args()58 59    root = Path(args.root).resolve()60    model_dir = (root / args.model_dir).resolve()61    dataset_root = (root / args.dataset_root).resolve()62 63    config = yaml.safe_load((root / "conf" / "alphafold3_paths.yaml").read_text(encoding="utf-8"))64    if config["model"]["repo_id"] != "OneScience/AlphaFold3/":65        raise AssertionError("model repo_id mismatch in conf/alphafold3_paths.yaml")66    if config["dataset"]["repo_id"] != "OneScience/AlphaFold3_dataset":67        raise AssertionError("dataset repo_id mismatch in conf/alphafold3_paths.yaml")68 69    for rel, (size, digest) in MODEL_FILES.items():70        path = model_dir / rel71        if not path.is_file():72            raise AssertionError(f"missing model file: {path}")73        if path.stat().st_size != size:74            raise AssertionError(f"model size mismatch: {path}")75        if not args.skip_hash and sha256(path) != digest:76            raise AssertionError(f"model sha256 mismatch: {path}")77 78    for rel in DATASET_REQUIRED:79        path = dataset_root / rel80        if not path.exists():81            raise AssertionError(f"missing dataset dependency: {path}")82 83    check_json(root / "inputs" / "7r6r_data.json")84    check_json(root / "inputs" / "t1119_search.json")85    print("model_preflight_ok: true")86 87 88if __name__ == "__main__":89    main()90