| |
| """Run AF3-NA+ from Hugging Face files and emit one K=4 artifact.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from dataclasses import asdict, dataclass |
| import json |
| import os |
| import pathlib |
| import platform |
| import tempfile |
| import time |
| from typing import Any, Mapping, Sequence |
|
|
| import numpy as np |
|
|
| from hub_loader import LoadedRelease, load_release |
| from evotemplate_na.inference import ( |
| load_artifact, |
| materialize_torch_tile_artifact_arrays, |
| write_artifact, |
| ) |
| from evotemplate_na.model import NucleicAcidChain, NucleicAcidTokenizer |
|
|
|
|
| class HubPredictionError(ValueError): |
| """Raised before publishing an invalid or ambiguous artifact.""" |
|
|
|
|
| @dataclass(frozen=True) |
| class PredictionSummary: |
| artifact_directory: str |
| arrays_file: str |
| manifest_file: str |
| artifact_id: str |
| length: int |
| num_chains: int |
| device: str |
| dtype: str |
| wall_seconds: float |
| gpu_seconds: float | None |
| peak_memory_mb: float | None |
|
|
|
|
| def _load_json(path: pathlib.Path) -> Any: |
| try: |
| return json.loads(path.read_text(encoding="utf-8")) |
| except OSError as exc: |
| raise HubPredictionError(f"cannot read input JSON: {exc}") from exc |
| except (UnicodeDecodeError, json.JSONDecodeError) as exc: |
| raise HubPredictionError(f"invalid input JSON: {exc}") from exc |
|
|
|
|
| def _canonical_chain(value: Mapping[str, Any], index: int) -> dict[str, str]: |
| if not isinstance(value, Mapping): |
| raise HubPredictionError(f"chains[{index}] must be an object") |
| if set(value) != {"chain_id", "polymer_type", "sequence"}: |
| raise HubPredictionError( |
| f"chains[{index}] requires exactly chain_id/polymer_type/sequence" |
| ) |
| chain_id = value["chain_id"] |
| polymer = value["polymer_type"] |
| sequence = value["sequence"] |
| if not isinstance(chain_id, str) or not chain_id.strip(): |
| raise HubPredictionError("chain_id must be a non-empty string") |
| if not isinstance(polymer, str) or polymer.strip().upper() not in ("RNA", "DNA"): |
| raise HubPredictionError("polymer_type must be RNA or DNA") |
| if not isinstance(sequence, str): |
| raise HubPredictionError("sequence must be a string") |
| canonical = "".join(sequence.split()).upper() |
| normalized_polymer = polymer.strip().upper() |
| alphabet = set("ACGU" if normalized_polymer == "RNA" else "ACGT") |
| if not canonical or not set(canonical) <= alphabet: |
| raise HubPredictionError( |
| f"chain {chain_id!r} is not canonical {normalized_polymer}" |
| ) |
| return { |
| "chain_id": chain_id.strip(), |
| "polymer_type": normalized_polymer, |
| "sequence": canonical, |
| } |
|
|
|
|
| def _af3_chains(payload: Mapping[str, Any]) -> list[dict[str, str]]: |
| if payload.get("dialect") != "alphafold3" or not isinstance( |
| payload.get("sequences"), list |
| ): |
| raise HubPredictionError("input is not an AlphaFold 3 JSON object") |
| raw: list[dict[str, Any]] = [] |
| for entry in payload["sequences"]: |
| if not isinstance(entry, Mapping) or len(entry) != 1: |
| raise HubPredictionError("malformed AlphaFold 3 sequence entry") |
| kind = next(iter(entry)) |
| if kind not in ("rna", "dna"): |
| continue |
| molecule = entry[kind] |
| if not isinstance(molecule, Mapping): |
| raise HubPredictionError("malformed AlphaFold 3 NA molecule") |
| raw_ids = molecule.get("id") |
| ids = raw_ids if isinstance(raw_ids, list) else [raw_ids] |
| for chain_id in ids: |
| raw.append( |
| { |
| "chain_id": chain_id, |
| "polymer_type": kind.upper(), |
| "sequence": molecule.get("sequence"), |
| } |
| ) |
| if not raw: |
| raise HubPredictionError("AlphaFold 3 input contains no RNA/DNA chains") |
| return [_canonical_chain(chain, index) for index, chain in enumerate(raw)] |
|
|
|
|
| def normalize_input( |
| payload: Any, |
| *, |
| fallback_artifact_id: str, |
| artifact_id_override: str | None = None, |
| ) -> tuple[str, tuple[dict[str, str], ...], tuple[tuple[str, str], ...], int]: |
| """Normalize either AF3 JSON or ``{name, chains}`` release input.""" |
|
|
| if not isinstance(payload, Mapping): |
| raise HubPredictionError("input JSON root must be an object") |
| if payload.get("dialect") == "alphafold3": |
| chains = _af3_chains(payload) |
| explicit_pairs = None |
| declared_max_length = None |
| else: |
| raw_chains = payload.get("chains") |
| if not isinstance(raw_chains, list) or not raw_chains: |
| raise HubPredictionError("simple input requires a non-empty chains list") |
| chains = [ |
| _canonical_chain(chain, index) for index, chain in enumerate(raw_chains) |
| ] |
| explicit_pairs = payload.get("authorized_chain_pairs") |
| declared_max_length = payload.get("max_length") |
| ids = [chain["chain_id"] for chain in chains] |
| if len(set(ids)) != len(ids): |
| raise HubPredictionError("chain_id values must be unique") |
| if explicit_pairs is None: |
| pairs = [ |
| (left, right) |
| for index, left in enumerate(ids) |
| for right in ids[index:] |
| ] |
| else: |
| if not isinstance(explicit_pairs, list) or not explicit_pairs: |
| raise HubPredictionError("authorized_chain_pairs must be a non-empty list") |
| pairs = [] |
| seen: set[tuple[str, str]] = set() |
| for pair in explicit_pairs: |
| if ( |
| not isinstance(pair, list) |
| or len(pair) != 2 |
| or pair[0] not in ids |
| or pair[1] not in ids |
| ): |
| raise HubPredictionError("authorized_chain_pairs contains an invalid pair") |
| canonical = tuple(sorted((str(pair[0]), str(pair[1])))) |
| if canonical in seen: |
| raise HubPredictionError("authorized_chain_pairs contains duplicates") |
| seen.add(canonical) |
| pairs.append((str(pair[0]), str(pair[1]))) |
| length = sum(len(chain["sequence"]) for chain in chains) |
| max_length = length if declared_max_length is None else declared_max_length |
| if ( |
| isinstance(max_length, bool) |
| or not isinstance(max_length, int) |
| or max_length < length |
| ): |
| raise HubPredictionError("max_length must be an integer covering all residues") |
| raw_name = artifact_id_override or payload.get("name") or fallback_artifact_id |
| if not isinstance(raw_name, str) or not raw_name.strip(): |
| raise HubPredictionError("artifact_id/name must be a non-empty string") |
| return raw_name.strip(), tuple(chains), tuple(pairs), max_length |
|
|
|
|
| def authorized_pair_mask( |
| chains: Sequence[Mapping[str, str]], pairs: Sequence[Sequence[str]] |
| ) -> np.ndarray: |
| offsets: dict[str, slice] = {} |
| cursor = 0 |
| for chain in chains: |
| stop = cursor + len(chain["sequence"]) |
| offsets[chain["chain_id"]] = slice(cursor, stop) |
| cursor = stop |
| result = np.zeros((cursor, cursor), dtype=np.bool_) |
| for left, right in pairs: |
| left_slice, right_slice = offsets[left], offsets[right] |
| result[left_slice, right_slice] = True |
| result[right_slice, left_slice] = True |
| return result |
|
|
|
|
| def predict( |
| *, |
| model_dir: str | pathlib.Path, |
| input_path: str | pathlib.Path, |
| output_directory: str | pathlib.Path, |
| device: str = "cpu", |
| dtype: str = "auto", |
| block_size: int | None = None, |
| artifact_id: str | None = None, |
| overwrite: bool = False, |
| ) -> PredictionSummary: |
| """Load one release and publish one AF3-compatible K=4 artifact.""" |
|
|
| try: |
| import torch |
| except ImportError as exc: |
| raise RuntimeError("AF3-NA+ requires torch") from exc |
| input_source = pathlib.Path(input_path).expanduser().resolve(strict=True) |
| payload = _load_json(input_source) |
| name, chains, pairs, max_length = normalize_input( |
| payload, |
| fallback_artifact_id=input_source.stem, |
| artifact_id_override=artifact_id, |
| ) |
| if block_size is not None and ( |
| isinstance(block_size, bool) or not isinstance(block_size, int) or block_size <= 0 |
| ): |
| raise HubPredictionError("block_size must be a positive integer") |
| if device == "auto": |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| loaded: LoadedRelease = load_release(model_dir, device=device, dtype=dtype) |
| if len(chains) > loaded.spec.model_config.student.max_chains: |
| raise HubPredictionError("input chain count exceeds the release max_chains") |
| destination = pathlib.Path(output_directory).expanduser().resolve() |
| if not overwrite and any( |
| (destination / name).exists() for name in ("manifest.json", "arrays.npz") |
| ): |
| raise HubPredictionError("artifact exists; pass --overwrite to replace it") |
| destination.mkdir(parents=True, exist_ok=True) |
| tokenizer = NucleicAcidTokenizer( |
| max_chains=loaded.spec.model_config.student.max_chains |
| ) |
| model_chains = [ |
| NucleicAcidChain( |
| chain_id=chain["chain_id"], |
| polymer_type=chain["polymer_type"], |
| sequence=chain["sequence"], |
| ) |
| for chain in chains |
| ] |
| batch = tokenizer.batch_encode([model_chains]).to(loaded.device) |
| length = sum(len(chain["sequence"]) for chain in chains) |
| residue_mask = np.ones(length, dtype=np.bool_) |
| pair_mask = authorized_pair_mask(chains, pairs) |
| tile_size = block_size or loaded.spec.model_config.pair_decoder.block_size |
| with tempfile.TemporaryDirectory(prefix=".af3-na-plus-", dir=destination) as scratch: |
| start = time.perf_counter() |
| gpu_seconds: float | None = None |
| peak_memory_mb: float | None = None |
| start_event = end_event = None |
| if loaded.device.type == "cuda": |
| if loaded.device.index is not None: |
| torch.cuda.set_device(loaded.device.index) |
| torch.cuda.reset_peak_memory_stats(loaded.device) |
| start_event = torch.cuda.Event(enable_timing=True) |
| end_event = torch.cuda.Event(enable_timing=True) |
| start_event.record() |
| with torch.inference_mode(): |
| student = loaded.model.encode_full_chain(batch) |
| context = loaded.model.prepare_context(batch, student_output=student) |
| arrays = materialize_torch_tile_artifact_arrays( |
| tile_forward=loaded.model.forward_tile, |
| context=context, |
| slot_logits=context.decoder_context.slot_logits, |
| length=length, |
| active_pair_mask=pair_mask, |
| residue_mask=residue_mask, |
| scratch_directory=pathlib.Path(scratch), |
| block_size=tile_size, |
| ) |
| if start_event is not None and end_event is not None: |
| end_event.record() |
| torch.cuda.synchronize(loaded.device) |
| gpu_seconds = float(start_event.elapsed_time(end_event)) / 1000.0 |
| peak_memory_mb = float(torch.cuda.max_memory_allocated(loaded.device)) / (1024.0**2) |
| wall_seconds = time.perf_counter() - start |
| precision = "bf16" if loaded.dtype == torch.bfloat16 else "fp32" |
| write_artifact( |
| destination, |
| arrays=arrays, |
| artifact_id=name, |
| chains=chains, |
| authorized_chain_pairs=pairs, |
| max_length=max_length, |
| view_mode="forward+reverse_complement", |
| runtime={ |
| "python": platform.python_version(), |
| "torch": torch.__version__, |
| "pid": os.getpid(), |
| "wall_seconds": wall_seconds, |
| "gpu_seconds": gpu_seconds, |
| "peak_memory_mb": peak_memory_mb, |
| }, |
| device=str(loaded.device), |
| precision=precision, |
| quantization=None, |
| slot_weight_is_normalized=True, |
| ) |
| |
| |
| artifact = load_artifact(destination) |
| return PredictionSummary( |
| artifact_directory=str(destination), |
| arrays_file=str(destination / "arrays.npz"), |
| manifest_file=str(destination / "manifest.json"), |
| artifact_id=name, |
| length=artifact.length, |
| num_chains=len(chains), |
| device=str(loaded.device), |
| dtype=str(loaded.dtype).removeprefix("torch."), |
| wall_seconds=wall_seconds, |
| gpu_seconds=gpu_seconds, |
| peak_memory_mb=peak_memory_mb, |
| ) |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--model-dir", type=pathlib.Path, default=pathlib.Path(".")) |
| parser.add_argument("--input", type=pathlib.Path, required=True) |
| parser.add_argument("--output", type=pathlib.Path, required=True) |
| parser.add_argument("--device", default="auto") |
| parser.add_argument( |
| "--dtype", choices=("auto", "bfloat16", "float32"), default="auto" |
| ) |
| parser.add_argument("--block-size", type=int) |
| parser.add_argument("--artifact-id") |
| parser.add_argument("--overwrite", action="store_true") |
| args = parser.parse_args(argv) |
| summary = predict( |
| model_dir=args.model_dir, |
| input_path=args.input, |
| output_directory=args.output, |
| device=args.device, |
| dtype=args.dtype, |
| block_size=args.block_size, |
| artifact_id=args.artifact_id, |
| overwrite=args.overwrite, |
| ) |
| print(json.dumps(asdict(summary), sort_keys=True)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
|
|
| __all__ = [ |
| "HubPredictionError", |
| "PredictionSummary", |
| "authorized_pair_mask", |
| "main", |
| "normalize_input", |
| "predict", |
| ] |
|
|
